"""LLM gateway: stage cascade (small → medium → large), OpenAI-compatible engine, strict JSON output validated by pydantic
schemas, full cost accounting in `llm_jobs`. Deterministic processing always runs before any call here."""
from __future__ import annotations

import asyncio
import json
import logging
import re
import socket
import time
from dataclasses import dataclass
from typing import Any

import httpx
from pydantic import BaseModel, ValidationError

from aiatlas.config import settings
from aiatlas.db import execute, jsonb, transaction
from aiatlas.ids import new_id

log = logging.getLogger(__name__)

STAGES = ("small", "medium", "large")


class LLMUnavailable(Exception):
    pass


@dataclass
class LLMResult:
    ok: bool
    data: dict[str, Any] | None
    raw: str
    model: str
    stage: str
    input_tokens: int | None
    output_tokens: int | None
    duration_ms: int
    error: str | None = None
    llm_job_id: str | None = None


class Engine:
    name = "null"

    async def complete(self, *, model: str, system: str, user: str, max_tokens: int, temperature: float, json_mode: bool) -> tuple[str, dict[str, Any]]:
        raise LLMUnavailable("no LLM engine configured")

    async def embed(self, *, model: str, texts: list[str]) -> list[list[float]]:
        raise LLMUnavailable("no embedding engine configured")

    async def health(self) -> bool:
        return False


class OpenAICompatEngine(Engine):
    """Works with MacLustr llm-api.io, vLLM, llama.cpp server, Ollama (/v1), LM Studio…"""

    name = "openai_compat"

    def __init__(self, base_url: str, api_key: str, timeout_s: float):
        self.base_url = base_url.rstrip("/")
        self.api_key = api_key
        self.timeout_s = timeout_s

    def _client(self) -> httpx.AsyncClient:
        return httpx.AsyncClient(base_url=self.base_url, headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
                                 timeout=httpx.Timeout(self.timeout_s, connect=30))

    async def _post(self, client: httpx.AsyncClient, path: str, body: dict[str, Any]) -> httpx.Response:
        """Local inference servers answer 503/429 while they load or swap a model: wait and retry instead of failing the job."""
        delay = 15.0
        for attempt in range(6):
            r = await client.post(path, json=body)
            if r.status_code in (503, 429, 502, 504) and attempt < 5:
                retry_after = r.headers.get("retry-after")
                wait = min(120.0, float(retry_after)) if retry_after and retry_after.replace(".", "", 1).isdigit() else delay
                log.info("llm server busy, retrying", extra={"status": r.status_code, "wait": wait, "path": path})
                await asyncio.sleep(wait)
                delay = min(120.0, delay * 1.7)
                continue
            return r
        return r

    async def complete(self, *, model: str, system: str, user: str, max_tokens: int, temperature: float, json_mode: bool) -> tuple[str, dict[str, Any]]:
        body: dict[str, Any] = {"model": model, "messages": [{"role": "system", "content": system}, {"role": "user", "content": user}],
                                "max_tokens": max_tokens, "temperature": temperature, "stream": False}
        if json_mode:
            body["response_format"] = {"type": "json_object"}
        async with self._client() as client:
            r = await self._post(client, "/chat/completions", body)
            if r.status_code == 400 and json_mode:
                body.pop("response_format")
                r = await self._post(client, "/chat/completions", body)
            r.raise_for_status()
            data = r.json()
        choice = (data.get("choices") or [{}])[0]
        content = (choice.get("message") or {}).get("content") or ""
        return content, data.get("usage") or {}

    async def embed(self, *, model: str, texts: list[str]) -> list[list[float]]:
        async with self._client() as client:
            r = await self._post(client, "/embeddings", {"model": model, "input": texts})
            r.raise_for_status()
            data = r.json()
        items = sorted(data.get("data") or [], key=lambda d: d.get("index", 0))
        return [d["embedding"] for d in items]

    async def health(self) -> bool:
        try:
            async with self._client() as client:
                r = await client.get("/models", timeout=20)
                return r.status_code == 200
        except Exception:  # noqa: BLE001
            return False


SYSTEM_PROMPT = """You are the extraction engine of AI Atlas, a structured database of the AI ecosystem.
Extract ONLY facts explicitly stated in the document. Never guess, never fill gaps from memory, never invent numbers or dates.
If a field is not stated, use null. Quote units exactly as written. Prefer official statements over marketing phrasing.
Return a single JSON object matching the requested schema and nothing else."""


class LLMGateway:
    def __init__(self, engine: Engine | None = None):
        self.engine = engine or self._default_engine()
        self.models = {"small": settings.llm_small_model, "medium": settings.llm_medium_model, "large": settings.llm_large_model}

    @staticmethod
    def _default_engine() -> Engine:
        if settings.llm_available:
            return OpenAICompatEngine(settings.llm_base_url, settings.llm_api_key, settings.llm_timeout_s)
        return Engine()

    @property
    def available(self) -> bool:
        return self.engine.name != "null"

    async def health(self) -> dict[str, Any]:
        return {"engine": self.engine.name, "available": self.available, "reachable": await self.engine.health() if self.available else False,
                "models": self.models, "embedding_model": settings.embedding_model}

    async def extract(self, *, task_type: str, document: str, schema: type[BaseModel], stage: str = "medium", instructions: str = "",
                      snapshot_id: str | None = None, entity_id: str | None = None, job_id: str | None = None, max_tokens: int = 2000,
                      escalate_on_failure: bool = True, max_chars: int = 24000) -> LLMResult:
        if not self.available:
            raise LLMUnavailable("LLM engine not configured (AIA_LLM_BASE_URL / AIA_LLM_API_KEY)")
        if stage not in STAGES:
            raise ValueError(f"stage must be one of {STAGES}")
        doc = document if len(document) <= max_chars else document[:max_chars] + "\n…[truncated]"
        schema_json = json.dumps(schema.model_json_schema(), ensure_ascii=False)
        user = (f"TASK: {task_type}\n{instructions}\n\nJSON SCHEMA:\n{schema_json}\n\nDOCUMENT:\n<<<\n{doc}\n>>>\n\n"
                f"Return only the JSON object.")
        stages = STAGES[STAGES.index(stage):] if escalate_on_failure else (stage,)
        last: LLMResult | None = None
        for st in stages:
            model = self.models[st]
            t0 = time.perf_counter()
            error = None
            data: dict[str, Any] | None = None
            raw = ""
            usage: dict[str, Any] = {}
            status = "ok"
            try:
                raw, usage = await self.engine.complete(model=model, system=SYSTEM_PROMPT, user=user, max_tokens=max_tokens, temperature=0.0, json_mode=True)
                parsed = _parse_json(raw)
                if parsed is None:
                    status, error = "invalid_json", "model did not return JSON"
                else:
                    try:
                        data = schema.model_validate(parsed).model_dump(mode="json")
                    except ValidationError as ve:
                        status, error = "schema_error", str(ve)[:2000]
                        data = parsed  # keep for debugging
            except Exception as exc:  # noqa: BLE001
                status, error = "failed", f"{exc.__class__.__name__}: {exc}"[:2000]
            duration = int((time.perf_counter() - t0) * 1000)
            llm_job_id = await self._account(task_type=task_type, stage=st, model=model, schema_name=schema.__name__, snapshot_id=snapshot_id,
                                             entity_id=entity_id, job_id=job_id, usage=usage, duration=duration, status=status,
                                             output=data if status in ("ok", "schema_error") else None, error=error)
            last = LLMResult(ok=status == "ok", data=data if status == "ok" else None, raw=raw, model=model, stage=st, input_tokens=usage.get("prompt_tokens"),
                             output_tokens=usage.get("completion_tokens"), duration_ms=duration, error=error, llm_job_id=llm_job_id)
            if last.ok or status == "failed":
                break  # transport failures: do not escalate blindly (probably the server), schema errors: try the next stage
        assert last is not None
        return last

    async def classify(self, *, text: str, labels: list[str], task_type: str = "classify", snapshot_id: str | None = None) -> str | None:
        from pydantic import Field, create_model

        Model = create_model("Classification", label=(str, Field(description=f"one of: {', '.join(labels)}")), confidence=(float, Field(ge=0, le=1)))
        res = await self.extract(task_type=task_type, document=text[:6000], schema=Model, stage="small", snapshot_id=snapshot_id,
                                 instructions=f"Choose exactly one label among: {', '.join(labels)}.", max_tokens=100, escalate_on_failure=False)
        if res.ok and res.data and res.data.get("label") in labels:
            return str(res.data["label"])
        return None

    async def embed(self, texts: list[str]) -> list[list[float]]:
        if not self.available:
            raise LLMUnavailable("embedding engine not configured")
        t0 = time.perf_counter()
        vectors = await self.engine.embed(model=settings.embedding_model, texts=texts)
        await self._account(task_type="embed", stage="small", model=settings.embedding_model, schema_name=None, snapshot_id=None, entity_id=None,
                            job_id=None, usage={"prompt_tokens": sum(len(t) // 4 for t in texts)}, duration=int((time.perf_counter() - t0) * 1000),
                            status="ok", output=None, error=None)
        return vectors

    async def _account(self, *, task_type: str, stage: str, model: str, schema_name: str | None, snapshot_id: str | None, entity_id: str | None,
                       job_id: str | None, usage: dict[str, Any], duration: int, status: str, output: dict[str, Any] | None, error: str | None) -> str:
        lid = new_id("llm_job")
        try:
            async with transaction() as conn:
                await execute(conn, """insert into llm_jobs (id, job_id, task_type, stage, engine, model, node, schema_name, snapshot_id, entity_id, input_tokens, output_tokens,
                                       duration_ms, status, output, error) values (:id, :j, :t, :st, :eng, :m, :node, :schema, :snap, :e, :it, :ot, :d, :status, cast(:o as jsonb), :err)""",
                              id=lid, j=job_id, t=task_type, st=stage, eng=self.engine.name, m=model, node=socket.gethostname(), schema=schema_name, snap=snapshot_id,
                              e=entity_id, it=usage.get("prompt_tokens"), ot=usage.get("completion_tokens"), d=duration, status=status,
                              o=jsonb(output) if output is not None else None, err=error)
        except Exception as exc:  # noqa: BLE001
            log.warning("llm accounting failed", extra={"error": str(exc)})
        return lid


def _parse_json(raw: str) -> dict[str, Any] | None:
    s = raw.strip()
    s = re.sub(r"<think>.*?</think>", "", s, flags=re.DOTALL).strip()
    if s.startswith("```"):
        s = re.sub(r"^```(?:json)?\s*", "", s)
        s = re.sub(r"\s*```$", "", s)
    try:
        v = json.loads(s)
        return v if isinstance(v, dict) else None
    except json.JSONDecodeError:
        pass
    m = re.search(r"\{.*\}", s, flags=re.DOTALL)
    if m:
        try:
            v = json.loads(m.group(0))
            return v if isinstance(v, dict) else None
        except json.JSONDecodeError:
            return None
    return None


gateway = LLMGateway()

__all__ = ["Engine", "LLMGateway", "LLMResult", "LLMUnavailable", "OpenAICompatEngine", "gateway"]