"""LLMClient — single entry point for every model call (OpenRouter, OpenAI-compatible).

Handles: streaming SSE parsing (keep-alive comments, fragmented tool_calls, final usage
chunk), native model fallback via `models`, retries on 429/5xx, attribution headers.
"""

from __future__ import annotations

import asyncio
import json
from collections.abc import AsyncIterator
from typing import Any

import httpx

from app.core.config import Settings, get_settings
from app.core.logging import get_logger
from app.llm.schemas import StreamEvent

log = get_logger("llm.openrouter")

RETRY_STATUS = {429, 500, 502, 503, 504}


class LLMError(Exception):
    def __init__(self, message: str, status: int | None = None) -> None:
        super().__init__(message)
        self.status = status


class LLMClient:
    def __init__(self, settings: Settings | None = None) -> None:
        self.settings = settings or get_settings()
        self._client: httpx.AsyncClient | None = None

    def _headers(self) -> dict[str, str]:
        return {
            "Authorization": f"Bearer {self.settings.OPENROUTER_API_KEY.get_secret_value()}",
            "HTTP-Referer": self.settings.OPENROUTER_APP_URL,
            "X-Title": self.settings.OPENROUTER_APP_NAME,
            "Content-Type": "application/json",
        }

    def client(self) -> httpx.AsyncClient:
        if self._client is None or self._client.is_closed:
            self._client = httpx.AsyncClient(
                base_url=self.settings.OPENROUTER_BASE_URL,
                timeout=httpx.Timeout(self.settings.LLM_TIMEOUT_SECONDS, connect=15),
                http2=True,
            )
        return self._client

    async def aclose(self) -> None:
        if self._client and not self._client.is_closed:
            await self._client.aclose()

    # ------------------------------------------------------------------ body
    def _body(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None,
        models: list[str],
        temperature: float,
        max_tokens: int,
        reasoning: dict[str, Any] | None,
        provider: dict[str, Any] | None,
        user_id_hash: str | None,
        stream: bool,
        response_format: dict[str, Any] | None,
    ) -> dict[str, Any]:
        body: dict[str, Any] = {
            "messages": messages,
            "stream": stream,
            "temperature": temperature,
            "max_tokens": max_tokens,
            "usage": {"include": True},
            "provider": provider or {"data_collection": "deny", "allow_fallbacks": True},
        }
        if len(models) == 1:
            body["model"] = models[0]
        else:
            body["model"] = models[0]
            body["models"] = models
        if tools:
            body["tools"] = tools
            body["tool_choice"] = "auto"
        if reasoning:
            body["reasoning"] = reasoning
        if user_id_hash:
            body["user"] = user_id_hash
        if response_format:
            body["response_format"] = response_format
        return body

    # --------------------------------------------------------------- stream
    async def stream_chat(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None,
        models: list[str],
        temperature: float = 0.3,
        max_tokens: int | None = None,
        reasoning: dict[str, Any] | None = None,
        provider: dict[str, Any] | None = None,
        user_id_hash: str | None = None,
        response_format: dict[str, Any] | None = None,
    ) -> AsyncIterator[StreamEvent]:
        body = self._body(
            messages, tools, models, temperature,
            max_tokens or self.settings.LLM_MAX_OUTPUT_TOKENS,
            reasoning, provider, user_id_hash, True, response_format,
        )
        attempt = 0
        while True:
            attempt += 1
            try:
                async with self.client().stream(
                    "POST", "/chat/completions", headers=self._headers(), json=body
                ) as resp:
                    if resp.status_code in RETRY_STATUS and attempt < 3:
                        await resp.aread()
                        log.warning("llm_retry", status=resp.status_code, attempt=attempt)
                        await asyncio.sleep(1.5 * attempt)
                        continue
                    if resp.status_code >= 400:
                        text = (await resp.aread()).decode(errors="replace")
                        raise LLMError(f"OpenRouter {resp.status_code}: {text[:400]}",
                                       resp.status_code)
                    async for ev in self._parse_sse(resp):
                        yield ev
                    return
            except (httpx.TransportError, httpx.TimeoutException) as exc:
                if attempt >= 3:
                    yield StreamEvent("error", {"code": "network", "message": str(exc)})
                    return
                log.warning("llm_network_retry", attempt=attempt, error=str(exc))
                await asyncio.sleep(1.5 * attempt)

    async def _parse_sse(self, resp: httpx.Response) -> AsyncIterator[StreamEvent]:
        """Parse OpenRouter SSE. Yields normalised StreamEvents."""
        tool_buf: dict[int, dict[str, Any]] = {}
        finish_reason: str | None = None
        model_used = ""
        gen_id = ""
        usage: dict[str, Any] | None = None
        async for raw_line in resp.aiter_lines():
            line = raw_line.strip()
            if not line or line.startswith(":"):  # keep-alive ": OPENROUTER PROCESSING"
                continue
            if not line.startswith("data:"):
                continue
            data = line[5:].strip()
            if data == "[DONE]":
                break
            try:
                chunk = json.loads(data)
            except json.JSONDecodeError:
                continue
            if "error" in chunk and chunk["error"]:
                err = chunk["error"]
                yield StreamEvent("error", {"code": str(err.get("code", "upstream")),
                                            "message": str(err.get("message", err))})
                return
            model_used = chunk.get("model") or model_used
            gen_id = chunk.get("id") or gen_id
            if chunk.get("usage"):
                usage = chunk["usage"]
            for choice in chunk.get("choices", []) or []:
                delta = choice.get("delta") or {}
                if delta.get("content"):
                    yield StreamEvent("text_delta", {"delta": delta["content"]})
                if delta.get("reasoning"):
                    yield StreamEvent("reasoning_delta", {"delta": delta["reasoning"]})
                for tc in delta.get("tool_calls") or []:
                    if "index" in tc and tc["index"] is not None:
                        idx = int(tc["index"])
                    elif tc.get("id") and any(b["id"] == tc["id"] for b in tool_buf.values()):
                        idx = next(k for k, b in tool_buf.items() if b["id"] == tc["id"])
                    elif tc.get("id") or not tool_buf:
                        idx = len(tool_buf)  # provider omitted index: new call
                    else:
                        idx = max(tool_buf)  # continuation of the last call
                    buf = tool_buf.get(idx)
                    if buf is None:
                        buf = {"id": tc.get("id") or f"call_{idx}", "name": "", "args": ""}
                        tool_buf[idx] = buf
                        yield StreamEvent("tool_call_start", {"index": idx, "id": buf["id"]})
                    if tc.get("id"):
                        buf["id"] = tc["id"]
                    fn = tc.get("function") or {}
                    if fn.get("name") and not buf["name"]:
                        buf["name"] = fn["name"]
                    if fn.get("arguments"):
                        buf["args"] += fn["arguments"]
                        yield StreamEvent("tool_call_delta", {"index": idx, "delta": fn["arguments"]})
                if choice.get("finish_reason"):
                    finish_reason = choice["finish_reason"]
        for idx in sorted(tool_buf):
            b = tool_buf[idx]
            yield StreamEvent("tool_call_end", {"index": idx, "id": b["id"], "name": b["name"],
                                                "arguments": b["args"]})
        if usage:
            yield StreamEvent("usage", {
                "input_tokens": usage.get("prompt_tokens", 0),
                "output_tokens": usage.get("completion_tokens", 0),
                "cost_usd": float(usage.get("cost", 0.0) or 0.0),
                "model": model_used,
                "generation_id": gen_id,
            })
        if finish_reason is None and tool_buf:
            finish_reason = "tool_calls"
        yield StreamEvent("done", {"finish_reason": finish_reason or "stop", "model": model_used,
                                   "generation_id": gen_id})

    # ------------------------------------------------------------- complete
    async def complete(
        self,
        messages: list[dict[str, Any]],
        models: list[str],
        temperature: float = 0.0,
        max_tokens: int = 1024,
        response_format: dict[str, Any] | None = None,
        user_id_hash: str | None = None,
    ) -> tuple[str, dict[str, Any]]:
        """Non-streaming helper for small tasks (titles, classification, quiz JSON)."""
        body = self._body(messages, None, models, temperature, max_tokens, None, None,
                          user_id_hash, False, response_format)
        for attempt in range(1, 4):
            try:
                resp = await self.client().post("/chat/completions", headers=self._headers(),
                                                json=body)
            except (httpx.TransportError, httpx.TimeoutException) as exc:
                if attempt == 3:
                    raise LLMError(f"network: {exc}") from exc
                await asyncio.sleep(1.5 * attempt)
                continue
            if resp.status_code in RETRY_STATUS and attempt < 3:
                await asyncio.sleep(1.5 * attempt)
                continue
            if resp.status_code >= 400:
                raise LLMError(f"OpenRouter {resp.status_code}: {resp.text[:400]}", resp.status_code)
            data = resp.json()
            choice = (data.get("choices") or [{}])[0]
            text = (choice.get("message") or {}).get("content") or ""
            usage = data.get("usage") or {}
            return text, {
                "input_tokens": usage.get("prompt_tokens", 0),
                "output_tokens": usage.get("completion_tokens", 0),
                "cost_usd": float(usage.get("cost", 0.0) or 0.0),
                "model": data.get("model", ""),
                "generation_id": data.get("id", ""),
            }
        raise LLMError("OpenRouter: retries exhausted")

    async def embed(self, texts: list[str]) -> list[list[float]] | None:
        """Optional embeddings via an OpenAI-compatible endpoint (EMBEDDINGS_BASE_URL)."""
        s = self.settings
        if not s.EMBEDDINGS_BASE_URL or not s.MODEL_EMBEDDINGS:
            return None
        headers = {"Content-Type": "application/json"}
        key = s.EMBEDDINGS_API_KEY.get_secret_value() or s.OPENROUTER_API_KEY.get_secret_value()
        if key:
            headers["Authorization"] = f"Bearer {key}"
        async with httpx.AsyncClient(timeout=60) as c:
            r = await c.post(f"{s.EMBEDDINGS_BASE_URL.rstrip('/')}/embeddings", headers=headers,
                             json={"model": s.MODEL_EMBEDDINGS, "input": texts})
            r.raise_for_status()
            data = r.json()["data"]
            return [d["embedding"] for d in sorted(data, key=lambda d: d["index"])]


_client: LLMClient | None = None


def get_llm() -> LLMClient:
    global _client
    if _client is None:
        _client = LLMClient()
    return _client