from abc import ABC, abstractmethod
from collections.abc import Callable
from contextlib import asynccontextmanager, contextmanager
from typing import TypeVar, override

import httpx
from httpx import USE_CLIENT_DEFAULT, Response
from pydantic import BaseModel

from core.domain.exceptions import JSONSchemaValidationError
from core.domain.message import Message
from core.providers._base.abstract_provider import (
    AbstractProvider,
    ProviderConfigInterface,
)
from core.providers._base.llm_usage import LLMUsage
from core.providers._base.models import RawCompletion
from core.providers._base.provider_error import (
    ContentModerationError,
    FailedGenerationError,
    InvalidGenerationError,
    InvalidProviderConfigError,
    ProviderError,
    ProviderInternalError,
    ProviderRateLimitError,
    ProviderTimeoutError,
    ProviderUnavailableError,
    ReadTimeOutError,
    ServerOverloadedError,
    UnknownProviderError,
)
from core.providers._base.provider_options import ProviderOptions
from core.runners.output_factory import OutputFactory
from core.runners.runner_output import RunnerOutput
from core.utils.background import add_background_task
from core.utils.streams import JSONStreamError

ResponseModel = TypeVar("ResponseModel", bound=BaseModel)


def _timeout_object(value: float):
    return httpx.Timeout(read=value, connect=10.0, pool=10.0, write=10.0)


# TODO: The fact that the HTTPXProvider class uses a plain dict as a request is blocking for OpenAIImageProvider
# Ultimately HTTPXProvider should also use a templated request type
class HTTPXProviderBase[ProviderConfigVar: ProviderConfigInterface, ProviderRequestVar](
    AbstractProvider[ProviderConfigVar, ProviderRequestVar],
    ABC,
):
    _shared_client: httpx.AsyncClient = httpx.AsyncClient(
        # 5 minutes timeout by default
        timeout=_timeout_object(300.0),
        # max_connections are per origin
        limits=httpx.Limits(max_connections=500, max_keepalive_connections=100),
    )

    @classmethod
    async def close(cls):
        await cls._shared_client.aclose()

    @classmethod
    def _invalid_json_error(
        cls,
        response: Response | None,
        exception: Exception | None,
        raw_completion: str,
        error_msg: str,
        retry: bool = False,
    ) -> Exception:
        moderation_patterns = ["inappropriate", "offensive"]
        if "apologize" in raw_completion.lower() and any(
            pattern in raw_completion.lower() for pattern in moderation_patterns
        ):
            return ContentModerationError(retry=retry, provider_error=raw_completion)

        # Ok to have a wide catch since this is only called on failed generation
        if "sorry" in raw_completion.lower():
            return FailedGenerationError(
                msg=f"Model refused to generate a response: {raw_completion}",
                response=response,
            )
        return FailedGenerationError(msg=error_msg, raw_completion=raw_completion, retry=retry)

    def _provider_rate_limit_error(self, response: Response):
        return ProviderRateLimitError(retry_after=10, response=response)

    def _provider_timeout_error(self, response: Response):
        return ProviderTimeoutError(retry_after=10, response=response)

    def _provider_internal_error(self, response: Response):
        return ProviderInternalError(retry_after=10, response=response)

    def _server_overloaded_error(self, response: Response):
        return ServerOverloadedError(retry_after=10, response=response)

    def _provider_unavailable_error(self, response: Response):
        return ProviderInternalError(retry_after=10, response=response)

    def _unknown_error_message(self, response: Response):
        """Method called to extract the error message from the response when"""
        return f"Unknown error status {response.status_code}"

    def _unknown_error(self, response: Response) -> ProviderError:
        return UnknownProviderError(msg=self._unknown_error_message(response), response=response)

    def _handle_error_status_code(self, response: Response):
        match response.status_code:
            case 401 | 403:
                err = InvalidProviderConfigError(f"Config {self._config_id} seems invalid", response=response)
                if not self._config_id:
                    # if no config id is provided, then it's the local config that is invalid
                    # so it should still log to sentry

                    self.logger.exception("Provider request timed out", response=response.text, exc_info=err)
                raise err
            case 402:
                # Payment required, let's raise an invalid config so the provider pipeline knows to go to the next provider
                raise InvalidProviderConfigError(
                    f"Payment required for provider {self._config_id}",
                    response=response,
                    capture=True,
                )
            case 408:
                raise self._provider_timeout_error(response)
            case 429:
                raise self._provider_rate_limit_error(response)
            case 500 | 520 | 530:
                raise self._provider_internal_error(response)
            case 502 | 503 | 522 | 524:
                raise self._provider_unavailable_error(response)
            case 529:
                raise self._server_overloaded_error(response)
            case _:
                # if no exception is raised, then it's an unknown error
                # which will be handled by the caller
                pass

    def _assign_raw_completion_response_on_error(self, raw_completion: RawCompletion, error: httpx.HTTPStatusError):
        """Set raw completion to None on status error. Usually non 200 status codes from providers means that the
        call did not succeed and we should not have a cost"""
        raw_completion.response = None

    @contextmanager
    def _wrap_errors(
        self,
        options: ProviderOptions,
        raw_completion: RawCompletion,
        finally_block: Callable[..., None] | None = None,
    ):
        """Remap errors for proper handling."""
        # TODO: this method is called in this file for completion but in httpx_provider for streaming
        try:
            yield
        except (httpx.ConnectError, httpx.ReadError) as e:
            raise ProviderUnavailableError(
                msg=f"Failed to reach provider: {e}",
                retry=True,
                capture=True,
                max_attempt_count=3,
            ) from e
        except httpx.HTTPStatusError as e:
            self._assign_raw_completion_response_on_error(raw_completion, e)
            self._handle_error_status_code(response=e.response)
            # if no exception is raised, then it's an unknown error
            raise self._unknown_error(e.response) from e
        except ProviderError as e:
            # Just forward provider errors
            e.provider_options = options
            e.provider = self.name()
            raise e
        except (JSONSchemaValidationError, JSONStreamError) as e:
            raise InvalidGenerationError(
                msg=f"Model failed to generate a valid json: {e}",
                provider_status_code=200,
            ) from e
        except httpx.TimeoutException as e:
            self._get_logger().warning("Provider request timed out", request_url=e.request.url, exc_info=e)
            raise ReadTimeOutError(retry=True, retry_after=10) from e
        except httpx.RemoteProtocolError as e:
            raise ProviderInternalError(
                msg="Provider has disconnected without sending a response.",
                retry_after=10,
            ) from e
        finally:
            if finally_block:
                finally_block()

    @asynccontextmanager
    async def _open_client(self, url: str):
        # We don't open or close the client here
        # Since we re-use them from a shared pool
        yield self._shared_client

    @classmethod
    def timeout_or_default(cls, value: float | None):
        if not value:
            return USE_CLIENT_DEFAULT
        return _timeout_object(value)

    @classmethod
    def _initial_usage(cls, messages: list[Message]) -> LLMUsage:
        image_count = 0
        has_audio = False
        for m in messages:
            for content in m.content:
                if content.file:
                    if content.file.is_image or content.file.is_pdf:
                        image_count += 1
                    if content.file.is_audio:
                        has_audio = True
        usage = LLMUsage(prompt_image_count=image_count)
        if not has_audio:
            usage.prompt_audio_duration_seconds = 0
            usage.prompt_audio_token_count = 0
        return usage

    async def _extract_and_log_rate_limits(self, response: Response, options: ProviderOptions):
        """Use _log_rate_limit from the base class to track rate limits"""

    @abstractmethod
    async def _execute_request(self, request: ProviderRequestVar, options: ProviderOptions) -> Response:
        pass

    @abstractmethod
    def _parse_response(
        self,
        response: Response,
        output_factory: OutputFactory,
        raw_completion: RawCompletion,
        request: ProviderRequestVar,
    ) -> RunnerOutput:
        pass

    @override
    async def _single_complete(
        self,
        request: ProviderRequestVar,
        output_factory: OutputFactory,
        raw_completion: RawCompletion,
        options: ProviderOptions,
    ) -> RunnerOutput:
        with self._wrap_errors(options=options, raw_completion=raw_completion):
            response = await self._execute_request(request, options)
            response = response.raise_for_status()
            add_background_task(self._extract_and_log_rate_limits(response, options=options))
            parsed_response = self._parse_response(
                response,
                output_factory=output_factory,
                raw_completion=raw_completion,
                request=request,
            )
            return parsed_response

    # TODO: we should likely handle streaming here as well, maybe when we merge this class
    # with httpx_provider.py

    @override
    async def check_valid(self) -> bool:
        options = ProviderOptions(model=self.default_model(), max_tokens=100, temperature=0)

        try:
            _ = await self.complete(
                messages=[Message.with_text(text="Respond with an empty json", role="user")],
                options=options,
                output_factory=lambda x: x,
            )
            return True
        except InvalidProviderConfigError:
            return False
