import httpx
import typing

from .const import ATMEEX_API_BASE_URL, COMMON_HEADERS

_NOT_AUTHENTICATED_MSG = "Not authenticated. Call client.signin_with_email() or client.request_sms_code() + signin_with_phone() first"


class AtmeexAuth(httpx.Auth):
    requires_response_body = True

    def __init__(self, access_token: str = "", refresh_token: str = ""):
        self._access_token = access_token
        self._refresh_token = refresh_token

    @staticmethod
    def request_phone_code(phone: str) -> httpx.Request:
        """Create request to send SMS code to phone number."""
        payload = {
            "grant_type": "phone_code",
            "phone": phone,
        }
        return httpx.Request(
            "POST", ATMEEX_API_BASE_URL + "/auth/signup", json=payload, headers=COMMON_HEADERS
        )

    @staticmethod
    async def request_phone_code_async(phone: str) -> None:
        """Send SMS code to phone number."""
        async with httpx.AsyncClient() as client:
            request = AtmeexAuth.request_phone_code(phone)
            response = await client.send(request)
            response.raise_for_status()

    @staticmethod
    def signin_with_phone(phone: str, phone_code: str) -> httpx.Request:
        """Create request to sign in with phone and SMS code."""
        payload = {
            "grant_type": "phone_code",
            "phone": phone,
            "phone_code": phone_code,
        }
        return httpx.Request(
            "POST", ATMEEX_API_BASE_URL + "/auth/signin", json=payload, headers=COMMON_HEADERS
        )

    @staticmethod
    async def signin_with_phone_async(phone: str, phone_code: str) -> tuple[str, str]:
        """Sign in with phone and SMS code. Returns (access_token, refresh_token)."""
        async with httpx.AsyncClient() as client:
            request = AtmeexAuth.signin_with_phone(phone, phone_code)
            response = await client.send(request)
            response.raise_for_status()
            data = response.json()
            return data["access_token"], data["refresh_token"]

    @staticmethod
    def signin_with_email(email: str, password: str) -> httpx.Request:
        """Create request to sign in with email and password."""
        payload = {
            "email": email,
            "password": password,
            "grant_type": "basic",
        }
        return httpx.Request(
            "POST", ATMEEX_API_BASE_URL + "/auth/signin", json=payload, headers=COMMON_HEADERS
        )

    @staticmethod
    async def signin_with_email_async(email: str, password: str) -> tuple[str, str]:
        """Sign in with email and password. Returns (access_token, refresh_token)."""
        async with httpx.AsyncClient() as client:
            request = AtmeexAuth.signin_with_email(email, password)
            response = await client.send(request)
            response.raise_for_status()
            data = response.json()
            return data["access_token"], data["refresh_token"]

    def auth_flow(self, request: httpx.Request) -> typing.Generator[httpx.Request, httpx.Response, None]:
        if self._access_token == "":
            raise RuntimeError(_NOT_AUTHENTICATED_MSG)

        request.headers["authorization"] = f"Bearer {self._access_token}"
        response = yield request

        if response.status_code == 401:
            yield from self._refresh_token()
            request.headers["authorization"] = f"Bearer {self._access_token}"
            yield request

    def _refresh_token(self) -> typing.Generator[httpx.Request, httpx.Response, None]:
        if self._refresh_token == "":
            raise RuntimeError(_NOT_AUTHENTICATED_MSG)

        payload = {
            "grant_type": "refresh_token",
            "refresh_token": self._refresh_token,
        }
        response = yield httpx.Request("POST", ATMEEX_API_BASE_URL + "/auth/signin", json=payload, headers=COMMON_HEADERS)
        if response.status_code == 401:
            raise RuntimeError("Session expired. Please sign in again.")
        else:
            self._handle_auth_response(response)

    def _handle_auth_response(self, response: httpx.Response):
        response.raise_for_status()
        data = response.json()
        self._refresh_token = data["refresh_token"]
        self._access_token = data["access_token"]
