import json
from http import HTTPStatus
from typing import Any, Literal, Self, cast, final

from httpx import AsyncClient, Response

from draive.ollama.config import OllamaChatConfig
from draive.ollama.errors import OllamaException
from draive.ollama.models import (
    ChatCompletionResponse,
    ChatMessage,
)
from draive.parameters import DataModel
from draive.scope import ScopeDependency
from draive.utils import getenv_str, not_missing

__all__ = [
    "OllamaClient",
]


@final
class OllamaClient(ScopeDependency):
    @classmethod
    def prepare(cls) -> Self:
        return cls(
            endpoint=getenv_str("OLLAMA_ENDPOINT", "http://localhost:11434"),
            timeout=90,
        )

    def __init__(
        self,
        endpoint: str,
        timeout: float | None = None,
    ) -> None:
        self._client: AsyncClient = AsyncClient(
            base_url=endpoint,
            timeout=timeout,
        )

    async def chat_completion(
        self,
        *,
        config: OllamaChatConfig,
        messages: list[ChatMessage],
    ) -> ChatCompletionResponse:
        return await self._create_chat_completion(
            messages=messages,
            model=config.model,
            temperature=config.temperature,
            top_k=config.top_k if not_missing(config.top_k) else None,
            top_p=config.top_p if not_missing(config.top_p) else None,
            max_tokens=config.max_tokens if not_missing(config.max_tokens) else None,
            response_format=cast(Literal["text", "json"], config.response_format)
            if not_missing(config.response_format)
            else "text",
            seed=config.seed if not_missing(config.seed) else None,
        )

    async def _create_chat_completion(  # noqa: PLR0913
        self,
        model: str,
        temperature: float,
        top_k: float | None,
        top_p: float | None,
        seed: int | None,
        max_tokens: int | None,
        response_format: Literal["text", "json"],
        messages: list[ChatMessage],
    ) -> ChatCompletionResponse:
        request_body: dict[str, Any] = {
            "model": model,
            "messages": [message.as_dict() for message in messages],
            "options": {
                "temperature": temperature,
            },
            "stream": False,
        }

        if max_tokens:
            request_body["options"]["num_predict"] = max_tokens
        if top_k is not None:
            request_body["options"]["top_k"] = top_k
        if top_p is not None:
            request_body["options"]["top_p"] = top_p
        if seed is not None:
            request_body["options"]["seed"] = seed
        if response_format == "json":
            request_body["format"] = "json"

        return await self._request(
            model=ChatCompletionResponse,
            method="POST",
            url="/api/chat",
            body=request_body,
        )

    async def dispose(self) -> None:
        await self._client.aclose()

    async def _request[Requested: DataModel](  # noqa: PLR0913
        self,
        model: type[Requested],
        method: str,
        url: str,
        query: dict[str, str] | None = None,
        headers: dict[str, str] | None = None,
        body: DataModel | dict[str, Any] | None = None,
        follow_redirects: bool | None = None,
        timeout: float | None = None,
    ) -> Requested:
        request_headers: dict[str, str]
        if headers:
            request_headers = headers
        else:
            request_headers = {
                "Accept": "application/json",
            }

        body_content: str | None
        match body:
            case None:
                body_content = None

            case body_model if isinstance(body_model, DataModel):
                body_content = body_model.as_json()

            case values:
                body_content = json.dumps(values)

        if body_content:
            request_headers["Content-Type"] = "application/json"

        response: Response
        try:
            response = await self._client.request(
                method=method,
                url=url,
                headers=request_headers,
                params=query,
                content=body_content,
                follow_redirects=follow_redirects or False,
                timeout=timeout,
            )

        except Exception as exc:
            raise OllamaException("Network request failed") from exc

        status: HTTPStatus = HTTPStatus(value=response.status_code)
        if status.is_success:
            try:
                return model.from_json(await response.aread())

            except Exception as exc:
                raise OllamaException("Failed to decode Ollama response", response) from exc

        elif status.is_client_error:
            error_body: bytes = await response.aread()
            raise OllamaException("Ollama request error: %s", error_body.decode("utf-8"))

        else:
            raise OllamaException("Network request failed", response)
