"""Any OpenAI-style `/chat/completions` endpoint: OpenRouter, OpenAI, YandexGPT compatibility API, a local server."""

import asyncio
import logging
from collections.abc import Awaitable, Callable
from typing import Any

import httpx

from app.config import LLMSettings
from app.llm.parsing import parse_json_object
from app.llm.provider import ChatMessage, JsonSchema, LLMError, LLMResult

log = logging.getLogger(__name__)

Sleep = Callable[[float], Awaitable[None]]

RETRY_STATUSES = {408, 409, 429, 500, 502, 503, 504}


class OpenAICompatProvider:
    def __init__(
        self,
        settings: LLMSettings,
        transport: httpx.AsyncBaseTransport | None = None,
        sleep: Sleep = asyncio.sleep,
    ) -> None:
        if not settings.base_url or not settings.model:
            raise LLMError("LLM_BASE_URL and LLM_MODEL are required for the openai_compat provider")
        self._settings = settings
        self._sleep = sleep
        self._http = httpx.AsyncClient(
            base_url=settings.base_url.rstrip("/"),
            headers={"Authorization": f"Bearer {settings.api_key.get_secret_value()}"},
            timeout=settings.timeout,
            transport=transport,
            proxy=settings.proxy_url.get_secret_value() if settings.proxy_url else None,
        )

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

    async def complete(self, messages: list[ChatMessage], schema: JsonSchema | None = None) -> LLMResult:
        body: dict[str, Any] = {
            "model": self._settings.model,
            "messages": [{"role": m.role, "content": m.content} for m in messages],
            "temperature": self._settings.temperature,
        }
        if schema is not None:
            body["response_format"] = {
                "type": "json_schema",
                "json_schema": {"name": schema.name, "schema": schema.schema, "strict": True},
            }
        data = await self._post(body)
        try:
            text = data["choices"][0]["message"]["content"] or ""
        except (KeyError, IndexError, TypeError) as e:
            raise LLMError(f"unexpected completion shape: {list(data)}") from e
        usage = data.get("usage") or {}
        return LLMResult(
            text=text,
            data=parse_json_object(text) if schema is not None else None,
            prompt_tokens=int(usage.get("prompt_tokens", 0)),
            completion_tokens=int(usage.get("completion_tokens", 0)),
            raw=data,
        )

    async def _post(self, body: dict[str, Any]) -> dict[str, Any]:
        attempts = self._settings.max_attempts
        for attempt in range(1, attempts + 1):
            try:
                response = await self._http.post("/chat/completions", json=body)
            except httpx.TransportError as e:
                if attempt == attempts:
                    raise LLMError(f"LLM unreachable: {type(e).__name__}") from e
            else:
                if response.status_code < 400:
                    result: dict[str, Any] = response.json()
                    return result
                if response.status_code not in RETRY_STATUSES or attempt == attempts:
                    # never echo the request: it may contain user text; the response body is the provider's error
                    raise LLMError(f"LLM HTTP {response.status_code}: {response.text[:200]}")
            log.warning("llm request failed, retry %d/%d", attempt, attempts - 1)
            await self._sleep(min(1.0 * 2 ** (attempt - 1), 8.0))
        raise AssertionError("unreachable")