from __future__ import annotations

import json
import time
import logging
from typing import Any, Dict, Type, Union, TypeVar, Optional

import httpx
from httpx import URL

from . import _exceptions
from ._utils import SensitiveHeadersFilter

ResponseT = TypeVar("ResponseT")

MAX_RETRIES = 2
RETRY_STATUS_CODES = {429, 500, 502, 503, 504}
INITIAL_RETRY_DELAY = 0.5
MAX_RETRY_DELAY = 8.0


log: logging.Logger = logging.getLogger(__name__)
log.addFilter(SensitiveHeadersFilter())


class BaseClient(httpx.Client):
    _max_retries: int

    def __init__(
        self,
        *,
        base_url: URL | str,
        headers: Optional[Dict[str, str]] = None,
        timeout: Union[float, httpx.Timeout, None] = None,
        max_retries: int = MAX_RETRIES,
        **kwargs: Any,
    ):
        self._max_retries = max_retries
        super().__init__(base_url=base_url, headers=headers, timeout=timeout, **kwargs)

    @property
    def auth_headers(self) -> dict[str, str]:
        return {}

    @property
    def default_headers(self) -> dict[str, str]:
        return {
            "Accept": "application/json",
            "Content-Type": "application/json",
            **self.auth_headers,
        }

    def _request_cast(
        self,
        method: str,
        url: str,
        *,
        cast_to: Optional[Type[ResponseT]] = None,
        body: Optional[Any] = None,
        params: Optional[Dict[str, Any]] = None,
        headers: Optional[Dict[str, str]] = None,
        **kwargs: Any,
    ) -> Union[ResponseT, httpx.Response]:
        combined_headers = {**self.default_headers, **(headers or {})}
        retries_left = self._max_retries
        delay = INITIAL_RETRY_DELAY

        while True:
            response = super().request(
                method=method,
                url=url,
                json=body,
                params=params,
                headers=combined_headers,
                **kwargs,
            )

            if response.status_code in RETRY_STATUS_CODES and retries_left > 0:
                retry_after = response.headers.get("retry-after")
                sleep_time = float(retry_after) if retry_after else delay
                sleep_time = min(sleep_time, MAX_RETRY_DELAY)
                log.debug(
                    "Retrying request after %.1fs (status %d, %d retries left)",
                    sleep_time,
                    response.status_code,
                    retries_left,
                )
                time.sleep(sleep_time)
                delay = min(delay * 2, MAX_RETRY_DELAY)
                retries_left -= 1
                continue

            try:
                response.raise_for_status()
            except httpx.HTTPStatusError as err:
                log.debug("Encountered httpx.HTTPStatusError", exc_info=True)
                log.debug("Re-raising status error")
                raise self._make_status_error_from_response(err.response) from None

            if cast_to:
                data = response.json()
                return cast_to(**data)
            return response

    def get_cast(
        self,
        url: str,
        *,
        cast_to: Optional[Type[ResponseT]] = None,
        params: Optional[Dict[str, Any]] = None,
        headers: Optional[Dict[str, str]] = None,
        **kwargs: Any,
    ) -> Union[ResponseT, httpx.Response]:
        return self._request_cast("GET", url, cast_to=cast_to, params=params, headers=headers, **kwargs)

    def post_cast(
        self,
        url: str,
        *,
        cast_to: Optional[Type[ResponseT]] = None,
        body: Optional[Any] = None,
        headers: Optional[Dict[str, str]] = None,
        **kwargs: Any,
    ) -> Union[ResponseT, httpx.Response]:
        return self._request_cast("POST", url, cast_to=cast_to, body=body, headers=headers, **kwargs)

    def patch_cast(
        self,
        url: str,
        *,
        cast_to: Optional[Type[ResponseT]] = None,
        body: Optional[Any] = None,
        headers: Optional[Dict[str, str]] = None,
        **kwargs: Any,
    ) -> Union[ResponseT, httpx.Response]:
        return self._request_cast("PATCH", url, cast_to=cast_to, body=body, headers=headers, **kwargs)

    def delete_cast(
        self,
        url: str,
        *,
        cast_to: Optional[Type[ResponseT]] = None,
        params: Optional[Dict[str, Any]] = None,
        headers: Optional[Dict[str, str]] = None,
        **kwargs: Any,
    ) -> Union[ResponseT, httpx.Response]:
        return self._request_cast("DELETE", url, cast_to=cast_to, params=params, headers=headers, **kwargs)

    def _make_status_error_from_response(
        self,
        response: httpx.Response,
    ) -> _exceptions.APIStatusError:
        err_text = response.text.strip()
        body = err_text

        try:
            body = json.loads(err_text)
            err_msg = f"Error code: {response.status_code} - {body}"
        except Exception:
            err_msg = err_text or f"Error code: {response.status_code}"

        return self._make_status_error(err_msg, body=body, response=response)

    def _make_status_error(
        self,
        err_msg: str,
        *,
        body: object,
        response: httpx.Response,
    ) -> _exceptions.APIStatusError:
        raise NotImplementedError()


class BaseAsyncClient(httpx.AsyncClient):
    _max_retries: int

    def __init__(
        self,
        *,
        base_url: URL | str,
        headers: Optional[Dict[str, str]] = None,
        timeout: Union[float, httpx.Timeout, None] = None,
        max_retries: int = MAX_RETRIES,
        **kwargs: Any,
    ):
        self._max_retries = max_retries
        super().__init__(base_url=base_url, headers=headers, timeout=timeout, **kwargs)

    @property
    def auth_headers(self) -> dict[str, str]:
        return {}

    @property
    def default_headers(self) -> dict[str, str]:
        return {
            "Accept": "application/json",
            "Content-Type": "application/json",
            **self.auth_headers,
        }

    async def _request_cast(
        self,
        method: str,
        url: str,
        *,
        cast_to: Optional[Type[ResponseT]] = None,
        body: Optional[Any] = None,
        params: Optional[Dict[str, Any]] = None,
        headers: Optional[Dict[str, str]] = None,
        **kwargs: Any,
    ) -> Union[ResponseT, httpx.Response]:
        import asyncio

        combined_headers = {**self.default_headers, **(headers or {})}
        retries_left = self._max_retries
        delay = INITIAL_RETRY_DELAY

        while True:
            response = await super().request(
                method=method,
                url=url,
                json=body,
                params=params,
                headers=combined_headers,
                **kwargs,
            )

            if response.status_code in RETRY_STATUS_CODES and retries_left > 0:
                retry_after = response.headers.get("retry-after")
                sleep_time = float(retry_after) if retry_after else delay
                sleep_time = min(sleep_time, MAX_RETRY_DELAY)
                log.debug(
                    "Retrying request after %.1fs (status %d, %d retries left)",
                    sleep_time,
                    response.status_code,
                    retries_left,
                )
                await asyncio.sleep(sleep_time)
                delay = min(delay * 2, MAX_RETRY_DELAY)
                retries_left -= 1
                continue

            try:
                response.raise_for_status()
            except httpx.HTTPStatusError as err:
                log.debug("Encountered httpx.HTTPStatusError", exc_info=True)
                log.debug("Re-raising status error")
                raise self._make_status_error_from_response(err.response) from None

            if cast_to:
                data = response.json()
                return cast_to(**data)
            return response

    async def get_cast(
        self,
        url: str,
        *,
        cast_to: Optional[Type[ResponseT]] = None,
        params: Optional[Dict[str, Any]] = None,
        headers: Optional[Dict[str, str]] = None,
        **kwargs: Any,
    ) -> Union[ResponseT, httpx.Response]:
        return await self._request_cast("GET", url, cast_to=cast_to, params=params, headers=headers, **kwargs)

    async def post_cast(
        self,
        url: str,
        *,
        cast_to: Optional[Type[ResponseT]] = None,
        body: Optional[Any] = None,
        headers: Optional[Dict[str, str]] = None,
        **kwargs: Any,
    ) -> Union[ResponseT, httpx.Response]:
        return await self._request_cast("POST", url, cast_to=cast_to, body=body, headers=headers, **kwargs)

    async def patch_cast(
        self,
        url: str,
        *,
        cast_to: Optional[Type[ResponseT]] = None,
        body: Optional[Any] = None,
        headers: Optional[Dict[str, str]] = None,
        **kwargs: Any,
    ) -> Union[ResponseT, httpx.Response]:
        return await self._request_cast("PATCH", url, cast_to=cast_to, body=body, headers=headers, **kwargs)

    async def delete_cast(
        self,
        url: str,
        *,
        cast_to: Optional[Type[ResponseT]] = None,
        params: Optional[Dict[str, Any]] = None,
        headers: Optional[Dict[str, str]] = None,
        **kwargs: Any,
    ) -> Union[ResponseT, httpx.Response]:
        return await self._request_cast("DELETE", url, cast_to=cast_to, params=params, headers=headers, **kwargs)

    def _make_status_error_from_response(
        self,
        response: httpx.Response,
    ) -> _exceptions.APIStatusError:
        err_text = response.text.strip()
        body = err_text

        try:
            body = json.loads(err_text)
            err_msg = f"Error code: {response.status_code} - {body}"
        except Exception:
            err_msg = err_text or f"Error code: {response.status_code}"

        return self._make_status_error(err_msg, body=body, response=response)

    def _make_status_error(
        self,
        err_msg: str,
        *,
        body: object,
        response: httpx.Response,
    ) -> _exceptions.APIStatusError:
        raise NotImplementedError()
