from __future__ import annotations

from typing import Any

import httpx

from agent_platform.config import Settings


def completions_url(base_url: str) -> str:
    base = base_url.rstrip("/")
    if base.endswith("/v1"):
        return f"{base}/chat/completions"
    return f"{base}/v1/chat/completions"


class ModelProvider:
    def __init__(self, settings: Settings, client: httpx.AsyncClient | None = None) -> None:
        self.settings = settings
        self.client = client or httpx.AsyncClient(
            timeout=httpx.Timeout(settings.model_timeout_seconds),
            limits=httpx.Limits(max_connections=50, max_keepalive_connections=20),
        )
        self._owns_client = client is None

    @property
    def headers(self) -> dict[str, str]:
        return {
            "Authorization": f"Bearer {self.settings.model_api_key}",
            "Content-Type": "application/json",
        }

    async def close(self) -> None:
        if self._owns_client:
            await self.client.aclose()

    async def complete(
        self,
        *,
        model: str,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None = None,
        temperature: float | None = None,
    ) -> dict[str, Any]:
        payload: dict[str, Any] = {
            "model": model,
            "messages": messages,
            "stream": False,
        }
        if tools:
            payload["tools"] = tools
            payload["tool_choice"] = "auto"
        if temperature is not None:
            payload["temperature"] = temperature
        response = await self.client.post(
            completions_url(self.settings.model_api_base_url),
            headers=self.headers,
            json=payload,
        )
        response.raise_for_status()
        data = response.json()
        if not data.get("choices"):
            raise RuntimeError("Model provider returned no choices")
        return data

    async def forward(self, payload: dict[str, Any]) -> httpx.Response:
        request = self.client.build_request(
            "POST",
            completions_url(self.settings.model_api_base_url),
            headers=self.headers,
            json=payload,
        )
        response = await self.client.send(request, stream=bool(payload.get("stream")))
        if response.status_code >= 400:
            await response.aread()
            response.raise_for_status()
        return response