# retoor <retoor@molodetz.nl>

from __future__ import annotations

import logging
from typing import Any, Callable

import httpx

from devplacepy import stealth
from .config import Settings
from .cost import record_cost, record_usage
from .errors import LLMError

logger = logging.getLogger("devii.llm")


class LLMClient:
    def __init__(
        self,
        settings: Settings,
        key_resolver: Callable[[], str] | None = None,
    ) -> None:
        self._settings = settings
        self._key_resolver = key_resolver
        self._client = stealth.stealth_async_client(
            timeout=settings.timeout_seconds,
            headers={
                "Authorization": f"Bearer {settings.ai_key}",
                "Content-Type": "application/json",
            },
        )

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

    async def _post(self, payload: dict[str, Any]) -> httpx.Response:
        response = await self._client.post(self._settings.ai_url, json=payload)
        if response.status_code == 401 and self._refresh_key():
            response = await self._client.post(self._settings.ai_url, json=payload)
        return response

    def _refresh_key(self) -> bool:
        if self._key_resolver is None:
            return False
        try:
            key = self._key_resolver()
        except Exception:  # noqa: BLE001 - refresh is best-effort; the original 401 still surfaces
            logger.exception("Credential refresh failed after 401")
            return False
        if not key or self._client.headers.get("Authorization") == f"Bearer {key}":
            return False
        self._client.headers["Authorization"] = f"Bearer {key}"
        logger.info("Refreshed model endpoint credentials after 401")
        return True

    async def complete(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]],
    ) -> dict[str, Any]:
        payload = {
            "model": self._settings.ai_model,
            "messages": messages,
            "tools": tools,
            "tool_choice": "auto",
        }
        try:
            logger.debug("LLM request with %d messages", len(messages))
            response = await self._post(payload)
        except httpx.HTTPError as exc:
            raise LLMError(f"Could not reach the model endpoint: {exc}") from exc

        if response.status_code >= 400:
            raise LLMError(
                f"Model endpoint returned {response.status_code}: {self._reason(response)}",
                status=response.status_code,
                body=response.text[:500],
            )

        try:
            data = response.json()
        except ValueError as exc:
            raise LLMError("Model endpoint returned invalid JSON.") from exc

        choices = data.get("choices")
        if not choices:
            raise LLMError("Model response contained no choices.", body=str(data)[:500])

        message = choices[0].get("message")
        if message is None:
            raise LLMError("Model response contained no message.", body=str(data)[:500])

        record_usage(data.get("usage"))
        self._record_native_cost(response)
        logger.debug(
            "LLM response received (tool_calls=%s)", bool(message.get("tool_calls"))
        )
        return message

    async def complete_text(
        self, messages: list[dict[str, Any]], temperature: float = 0.0
    ) -> str:
        payload = {
            "model": self._settings.ai_model,
            "messages": messages,
            "temperature": temperature,
        }
        try:
            response = await self._post(payload)
        except httpx.HTTPError as exc:
            raise LLMError(f"Could not reach the model endpoint: {exc}") from exc
        if response.status_code >= 400:
            raise LLMError(
                f"Model endpoint returned {response.status_code}: {self._reason(response)}"
            )
        try:
            data = response.json()
            content = data["choices"][0]["message"]["content"] or ""
        except (ValueError, KeyError, IndexError) as exc:
            raise LLMError("Model endpoint returned an unexpected response.") from exc
        record_usage(data.get("usage"))
        self._record_native_cost(response)
        return content

    async def summarize(self, text: str) -> str:
        payload = {
            "model": self._settings.ai_model,
            "messages": [
                {
                    "role": "system",
                    "content": "You are a precise technical summarizer.",
                },
                {"role": "user", "content": text},
            ],
            "temperature": 0.0,
        }
        try:
            response = await self._post(payload)
        except httpx.HTTPError as exc:
            raise LLMError(f"Could not reach the model endpoint: {exc}") from exc
        if response.status_code >= 400:
            raise LLMError(
                f"Model endpoint returned {response.status_code}: {self._reason(response)}"
            )
        try:
            data = response.json()
            content = data["choices"][0]["message"]["content"] or ""
        except (ValueError, KeyError, IndexError) as exc:
            raise LLMError(
                "Model summarization returned an unexpected response."
            ) from exc
        record_usage(data.get("usage"))
        self._record_native_cost(response)
        return content

    @staticmethod
    def _record_native_cost(response: httpx.Response) -> None:
        raw = response.headers.get("X-Gateway-Cost-USD")
        if raw is None:
            return
        try:
            record_cost(float(raw))
        except (TypeError, ValueError):
            return

    @staticmethod
    def _reason(response: httpx.Response) -> str:
        try:
            error = response.json().get("error")
        except ValueError:
            return response.text[:200] or response.reason_phrase
        if isinstance(error, dict):
            return str(error.get("message", error))
        if error:
            return str(error)
        return response.reason_phrase