from typing import Any, Optional

from httpx import URL
from httpx import AsyncClient as HttpxAsyncClient
from httpx import Client as HttpxClient
from httpx import Request, Response
from httpx import codes as HttpxCodes
from httpx._types import AuthTypes

from firebolt.client.auth import Auth
from firebolt.client.auth.base import AuthRequest
from firebolt.client.constants import DEFAULT_API_URL
from firebolt.utils.exception import AccountNotFoundError
from firebolt.utils.urls import ACCOUNT_BY_NAME_URL, ACCOUNT_URL
from firebolt.utils.util import (
    cached_property,
    fix_url_schema,
    merge_urls,
    mixin_for,
)

FireboltClientMixinBase = mixin_for(HttpxClient)  # type: Any


class FireboltClientMixin(FireboltClientMixinBase):
    """HttpxAsyncClient mixin with Firebolt authentication functionality."""

    def __init__(
        self,
        *args: Any,
        account_name: Optional[str] = None,
        api_endpoint: str = DEFAULT_API_URL,
        auth: Optional[Auth] = None,
        **kwargs: Any,
    ):
        self.account_name = account_name
        self._api_endpoint = URL(fix_url_schema(api_endpoint))
        super().__init__(*args, auth=auth, **kwargs)

    def _build_auth(self, auth: Optional[AuthTypes]) -> Optional[Auth]:
        """Create Auth object based on auth provided.

        Overrides ``httpx.Client._build_auth``

        Args:
            auth (AuthTypes): Provided auth

        Returns:
            Optional[Auth]: Auth object

        Raises:
            TypeError: Auth argument has unsupported type
        """
        if auth is None or isinstance(auth, Auth):
            return auth
        raise TypeError(f'Invalid "auth" argument: {auth!r}')

    def _merge_auth_request(self, request: Request) -> Request:
        if isinstance(request, AuthRequest):
            request.url = merge_urls(self._api_endpoint, request.url)
            request._prepare(dict(request.headers))
        return request


class Client(FireboltClientMixin, HttpxClient):
    """An HTTP client, based on httpx.Client.

    Handles the authentication for Firebolt database.
    Authentication can be passed through auth keyword as a tuple or as a
    FireboltAuth instance
    """

    @cached_property
    def account_id(self) -> str:
        """User account ID.

        If account_name was provided during Client construction, returns its ID;
        gets default account otherwise.

        Returns:
            str: Account ID

        Raises:
            AccountNotFoundError: No account found with provided name
        """
        if self.account_name:
            response = self.get(
                url=ACCOUNT_BY_NAME_URL, params={"account_name": self.account_name}
            )
            if response.status_code == HttpxCodes.NOT_FOUND:
                raise AccountNotFoundError(self.account_name)
            # process all other status codes
            response.raise_for_status()
            return response.json()["account_id"]

        # account_name isn't set, use the default account.
        return self.get(url=ACCOUNT_URL).json()["account"]["id"]

    def _send_handling_redirects(
        self, request: Request, *args: Any, **kwargs: Any
    ) -> Response:
        return super()._send_handling_redirects(
            self._merge_auth_request(request), *args, **kwargs
        )


class AsyncClient(FireboltClientMixin, HttpxAsyncClient):
    """An HTTP client, based on httpx.AsyncClient.

    Asynchronously handles authentication for Firebolt database.
    Authentication can be passed through auth keyword as a tuple, or as a
    FireboltAuth instance.
    """

    @cached_property
    async def account_id(self) -> str:
        """User account id.

        If account_name was provided during AsyncClient construction, returns its ID;
        gets default account otherwise.

        Returns:
            str: Account ID

        Raises:
            AccountNotFoundError: No account found with provided name
        """
        if self.account_name:
            response = await self.get(
                url=ACCOUNT_BY_NAME_URL, params={"account_name": self.account_name}
            )
            if response.status_code == HttpxCodes.NOT_FOUND:
                raise AccountNotFoundError(self.account_name)
            # process all other status codes
            response.raise_for_status()
            return response.json()["account_id"]

        # account_name isn't set; use the default account.
        return (await self.get(url=ACCOUNT_URL)).json()["account"]["id"]

    async def _send_handling_redirects(
        self, request: Request, *args: Any, **kwargs: Any
    ) -> Response:
        return await super()._send_handling_redirects(
            self._merge_auth_request(request), *args, **kwargs
        )
