from __future__ import annotations

import datetime as dt
import functools
import json
from abc import ABC, abstractmethod
from structlog import BoundLogger
from typing import Any, Callable, Generator, Mapping, Type

import backoff
import requests

from src.api_client.errors import BaseApiError
from src.config import API_RETRY_COUNT


__all__ = ["BaseApiClient", "AuthorizedApiClient"]


class BaseApiClient:
    default_error_cls: Type[BaseApiError] = BaseApiError

    def __init__(self, *, logger: BoundLogger):
        self._logger = logger

    def _handle_retry(self, retry_data: Mapping[str, Any]):
        pass

    @staticmethod
    def _is_auth_error(error: BaseApiError) -> bool:
        return error.response.status_code == 401

    @staticmethod
    def _is_retry_error(error: BaseApiError) -> bool:
        return error.response.status_code == 429

    def _check_response(self, response: requests.Response) -> None:
        try:
            response.raise_for_status()
        except requests.HTTPError as e:
            self._logger.debug(
                json.dumps(
                    {
                        "response": {
                            "code": e.response.status_code,
                            "text": e.response.text,
                            "headers": dict(e.response.headers),
                        }
                        if e.response is not None
                        else None,
                        "request": {
                            "method": e.request.method,
                            "url": e.request.url,
                            "body": e.request.body.decode() if isinstance(e.request.body, bytes) else e.request.body,
                            "headers": dict(e.request.headers),
                        }
                        if e.request is not None
                        else None,
                    }
                )
            )
            raise self.__class__.default_error_cls(response=e.response) from e

    def send_request(self, request: requests.Request, **kwargs) -> requests.Response:
        """
        Send request
        :param request: Request object to send
        :return: Response object
        """
        with requests.Session() as session:
            response = session.send(session.prepare_request(request))
        return response

    def send_request_and_check(
        self, request: requests.Request, *, check_response: Callable[[requests.Response], None] | None = None, **kwargs
    ) -> requests.Response:
        """
        Send request and check the response
        :param request: Request object to send
        :return: Response object
        """
        response = self.send_request(request, **kwargs)
        check = self._check_response if check_response is None else check_response
        check(response)
        return response

    def _giveup(self, error: Exception):
        """Do not retry on 4xx except auth."""
        return (
            isinstance(error, requests.exceptions.HTTPError)
            and error.response is not None
            and error.response.status_code < 500
            and not self._is_auth_error(error)
            and not self._is_retry_error(error)
        )

    def _calc_wait_time(self, count: int, exception: Exception):
        """Calculate wait time."""

        return count

    def _get_wait_generator(self) -> Generator[int, None, None]:
        """Generator for retry wait interval."""

        exception = yield  # type: ignore[misc]
        count = 1
        while True:
            exception = yield self._calc_wait_time(count, exception)
            count += 1

    def send_request_with_retry(
        self, request: requests.Request, *, max_tries: int = API_RETRY_COUNT, **kwargs
    ) -> requests.Response:
        """
        Send request, check the response and retry if it's failed
        :param request: Request object to send
        :param max_tries: Max number of retries before giving up
        :return: Response object
        """

        @backoff.on_exception(
            self._get_wait_generator,
            self.__class__.default_error_cls,
            max_tries=max_tries,
            logger=self._logger,
            on_backoff=self._handle_retry,
            giveup=self._giveup,
            jitter=None,
        )
        @functools.wraps(self.send_request_and_check)
        def wrap(request_: requests.Request, **kwargs_):
            return self.send_request_and_check(request_, **kwargs_)

        return wrap(request, **kwargs)


class AuthorizedApiClient(BaseApiClient, ABC):
    auth_token_type: str = "Bearer"

    def __init__(
        self,
        *,
        token_update_callback: Callable[[str, int, dt.datetime], Any] | None = None,
        refresh_token: str | None = None,
        access_token: str | None = None,
        expires_in: int | None = None,
        updated_at: dt.datetime | None = None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self._refresh_token = refresh_token
        self._access_token = access_token
        self._expires_in = expires_in or 0
        self._updated_at = updated_at or dt.datetime.min
        self._token_update_callback = token_update_callback

    def _handle_retry(self, retry_data: Mapping[str, Any]):
        if self._is_auth_error(retry_data["exception"]) and retry_data["kwargs"].get("authorize", True):
            self.get_access_token(force_update=True)

    def send_request(self, request: requests.Request, *, authorize: bool = True, **kwargs) -> requests.Response:
        if authorize:
            self._authorize_request(request)

        return super().send_request(request)

    def get_access_token(self, force_update: bool = False) -> str:
        if self._access_token is None or force_update or self._is_token_expired():
            self._access_token, self._expires_in = self._refresh_access_token()
            self._updated_at = dt.datetime.utcnow()

            if self._token_update_callback is not None:
                self._token_update_callback(self._access_token, self._expires_in, self._updated_at)

        return self._access_token

    def _is_token_expired(self) -> bool:
        return dt.datetime.utcnow() >= self._updated_at + dt.timedelta(seconds=self._expires_in)

    def _authorize_request(self, request: requests.Request):
        request.headers["Authorization"] = f"{self.auth_token_type} {self.get_access_token()}"

    @abstractmethod
    def _refresh_access_token(self) -> tuple[str, int]:
        """
        Get a new token, and it's expiration
        :return: tuple of token and expires_in
        """