from __future__ import annotations

import asyncio
from typing import Any

import httpx

from agent_platform.config import Settings

TRANSIENT_COMPLETE_STATUS_CODES = frozenset({408, 425, 429, 500, 502, 503, 504, 524})
COMPLETE_MAX_ATTEMPTS = 2
COMPLETE_MAX_TOKENS = 4096
MAX_RETRY_DELAY_SECONDS = 5.0


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"


def retry_delay_seconds(response: httpx.Response | None, attempt: int) -> float:
    if response is not None:
        retry_after = response.headers.get("Retry-After")
        if retry_after:
            try:
                return min(MAX_RETRY_DELAY_SECONDS, max(0.0, float(retry_after)))
            except ValueError:
                pass
    return min(MAX_RETRY_DELAY_SECONDS, float(2 ** (attempt - 1)))


async def decode_complete_response(response: httpx.Response) -> dict[str, Any]:
    content_type = response.headers.get("Content-Type", "").lower()
    if "text/event-stream" not in content_type:
        await response.aread()
        data = response.json()
        if not data.get("choices"):
            raise RuntimeError("Model provider returned no choices")
        return data

    envelope: dict[str, Any] = {"object": "chat.completion"}
    message: dict[str, Any] = {"role": "assistant", "content": None}
    content_parts: list[str] = []
    reasoning_parts: list[str] = []
    tool_calls: dict[int, dict[str, Any]] = {}
    finish_reason: str | None = None
    usage: dict[str, Any] | None = None
    saw_choice = False

    async for line in response.aiter_lines():
        if not line.startswith("data:"):
            continue
        raw = line[5:].strip()
        if not raw or raw == "[DONE]":
            continue
        chunk = httpx.Response(200, content=raw).json()
        if chunk.get("error"):
            raise RuntimeError(f"Model provider stream error: {chunk['error']}")
        for key in ("id", "created", "model", "system_fingerprint"):
            if chunk.get(key) is not None:
                envelope[key] = chunk[key]
        if chunk.get("usage"):
            usage = chunk["usage"]
        choices = chunk.get("choices") or []
        if not choices:
            continue
        saw_choice = True
        choice = choices[0]
        delta = choice.get("delta") or {}
        if delta.get("role"):
            message["role"] = delta["role"]
        if isinstance(delta.get("content"), str):
            content_parts.append(delta["content"])
        reasoning = delta.get("reasoning_content")
        if not isinstance(reasoning, str):
            reasoning = delta.get("reasoning")
        if isinstance(reasoning, str):
            reasoning_parts.append(reasoning)
        for position, fragment in enumerate(delta.get("tool_calls") or []):
            index = int(fragment.get("index", position))
            merged = tool_calls.setdefault(
                index,
                {
                    "id": "",
                    "type": "function",
                    "function": {"name": "", "arguments": ""},
                },
            )
            if fragment.get("id"):
                merged["id"] = fragment["id"]
            if fragment.get("type"):
                merged["type"] = fragment["type"]
            function = fragment.get("function") or {}
            if function.get("name"):
                merged["function"]["name"] += function["name"]
            if function.get("arguments"):
                merged["function"]["arguments"] += function["arguments"]
        if choice.get("finish_reason") is not None:
            finish_reason = choice["finish_reason"]

    if not saw_choice:
        raise RuntimeError("Model provider stream returned no choices")
    if content_parts:
        message["content"] = "".join(content_parts)
    if reasoning_parts:
        message["reasoning_content"] = "".join(reasoning_parts)
    if tool_calls:
        message["tool_calls"] = [tool_calls[index] for index in sorted(tool_calls)]

    envelope["choices"] = [
        {
            "index": 0,
            "message": message,
            "finish_reason": finish_reason,
        }
    ]
    if usage is not None:
        envelope["usage"] = usage
    return envelope


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": True,
            "max_tokens": COMPLETE_MAX_TOKENS,
        }
        if tools:
            payload["tools"] = tools
            payload["tool_choice"] = "auto"
        if temperature is not None:
            payload["temperature"] = temperature
        for attempt in range(1, COMPLETE_MAX_ATTEMPTS + 1):
            response: httpx.Response | None = None
            try:
                async with self.client.stream(
                    "POST",
                    completions_url(self.settings.model_api_base_url),
                    headers=self.headers,
                    json=payload,
                ) as response:
                    if response.status_code in TRANSIENT_COMPLETE_STATUS_CODES and attempt < COMPLETE_MAX_ATTEMPTS:
                        await response.aread()
                        await asyncio.sleep(retry_delay_seconds(response, attempt))
                        continue

                    response.raise_for_status()
                    return await decode_complete_response(response)
            except httpx.TransportError:
                if attempt >= COMPLETE_MAX_ATTEMPTS:
                    raise
                await asyncio.sleep(retry_delay_seconds(response, attempt))
                continue

        raise RuntimeError("Model provider retry loop exited unexpectedly")

    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