"""LLM provider abstraction (spec §25): OpenAI-compatible chat completions with JSON mode, strict pydantic validation, one repair
retry, backoff on 429/503 (the MacLustr llm-api.io server loads models on demand and answers 503/429 while swapping), token
accounting and health. Keep ONE model per job stream (sticky) to avoid evictions — see `enrich.py`.

    provider = get_provider()
    result = await provider.complete_json("small", system, user_json, ChangeClassification, max_tokens=600)
    result.data  → validated pydantic model · result.request_tokens / response_tokens / latency_ms / model / attempts
"""
from __future__ import annotations

import asyncio
import json
import logging
import re
import time
from abc import ABC, abstractmethod
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from typing import Any

import httpx
from pydantic import BaseModel, ValidationError

from companyatlas.config import settings

log = logging.getLogger(__name__)

Tier = str                                              # "small" | "medium" | "large"
RETRY_STATUS = {408, 425, 429, 500, 502, 503, 504}
_JSON_BLOCK = re.compile(r"```(?:json)?\s*(\{.*?\})\s*```", re.DOTALL)
_THINK = re.compile(r"<think>.*?</think>", re.DOTALL)


class LLMError(Exception):
    def __init__(self, message: str, *, status: int | None = None, retryable: bool = False):
        super().__init__(message)
        self.status = status
        self.retryable = retryable


class LLMNotConfigured(LLMError):
    pass


class LLMValidationError(LLMError):
    """Model output could not be parsed/validated even after the repair retry."""


@dataclass(slots=True)
class LLMResult[T: BaseModel]:
    data: T
    model: str
    request_tokens: int = 0
    response_tokens: int = 0
    latency_ms: int = 0
    attempts: int = 1
    repaired: bool = False
    raw: str = ""


@dataclass(slots=True)
class Health:
    ok: bool
    base_url: str
    models: list[str] = field(default_factory=list)
    latency_ms: int = 0
    error: str | None = None


class LLMProvider(ABC):
    """Vendor-neutral interface. Implementations must be safe to share across tasks in one process."""

    @abstractmethod
    def model_for(self, tier: Tier) -> str: ...

    @abstractmethod
    async def complete_json[T: BaseModel](self, model_tier: Tier, system: str, user: str, schema: type[T], *, max_tokens: int = 800,
                            temperature: float = 0.1) -> LLMResult[T]: ...

    @abstractmethod
    async def complete_text(self, model_tier: Tier, system: str, user: str, *, max_tokens: int = 400, temperature: float = 0.2) -> LLMResult[Any]: ...

    @abstractmethod
    async def embed(self, texts: list[str]) -> list[list[float]]: ...

    @abstractmethod
    async def health(self) -> Health: ...

    async def close(self) -> None:
        return None


class OpenAICompatibleProvider(LLMProvider):
    def __init__(self, *, base_url: str | None = None, api_key: str | None = None, timeout_s: float | None = None,
                 max_tries: int | None = None, backoff_initial_s: float | None = None, backoff_max_s: float | None = None,
                 sleep: Callable[[float], Awaitable[None]] | None = None, transport: httpx.AsyncBaseTransport | None = None):
        self.base_url = (base_url if base_url is not None else settings.llm_base_url).rstrip("/")
        self.api_key = api_key if api_key is not None else settings.llm_api_key
        self.timeout_s = timeout_s or settings.llm_timeout_s
        self.max_tries = max_tries or settings.llm_max_tries
        self.backoff_initial_s = backoff_initial_s if backoff_initial_s is not None else settings.llm_backoff_initial_s
        self.backoff_max_s = backoff_max_s if backoff_max_s is not None else settings.llm_backoff_max_s
        self._sleep = sleep or asyncio.sleep
        self._transport = transport
        self._client: httpx.AsyncClient | None = None
        self._models = {"small": settings.llm_small_model, "medium": settings.llm_medium_model, "large": settings.llm_large_model,
                        "embedding": settings.llm_embedding_model}

    # ------------------------------------------------------------------ plumbing
    @property
    def configured(self) -> bool:
        return bool(self.base_url)

    def model_for(self, tier: Tier) -> str:
        return self._models.get(tier) or tier                   # a literal model name is accepted too

    def _headers(self) -> dict[str, str]:
        h = {"Content-Type": "application/json", "Accept": "application/json", "User-Agent": "CompanyAtlas-LLM/0.1"}
        if self.api_key:
            h["Authorization"] = f"Bearer {self.api_key}"
        return h

    def client(self) -> httpx.AsyncClient:
        if self._client is None:
            self._client = httpx.AsyncClient(base_url=self.base_url, headers=self._headers(), timeout=httpx.Timeout(self.timeout_s, connect=20),
                                             transport=self._transport)
        return self._client

    async def close(self) -> None:
        if self._client is not None:
            await self._client.aclose()
            self._client = None

    def _backoff(self, attempt: int) -> float:
        return min(self.backoff_max_s, self.backoff_initial_s * (2 ** attempt))

    async def _post(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
        if not self.configured:
            raise LLMNotConfigured("CA_LLM_BASE_URL is not set")
        last: LLMError | None = None
        for attempt in range(self.max_tries):
            try:
                r = await self.client().post(path, json=payload)
            except httpx.TimeoutException as exc:
                last = LLMError(f"timeout: {exc}", retryable=True)
            except httpx.HTTPError as exc:
                last = LLMError(f"transport: {exc}", retryable=True)
            else:
                if r.status_code < 300:
                    try:
                        return r.json()
                    except ValueError as exc:
                        raise LLMError(f"non-JSON response: {exc}", status=r.status_code) from exc
                body = r.text[:300]
                if r.status_code in RETRY_STATUS:
                    last = LLMError(f"HTTP {r.status_code}: {body}", status=r.status_code, retryable=True)
                    retry_after = r.headers.get("Retry-After")
                    if retry_after and retry_after.isdigit():
                        await self._sleep(min(self.backoff_max_s, float(retry_after)))
                        continue
                else:
                    raise LLMError(f"HTTP {r.status_code}: {body}", status=r.status_code)
            if attempt + 1 < self.max_tries:
                delay = self._backoff(attempt)
                log.warning("llm retry", extra={"attempt": attempt + 1, "delay_s": delay, "error": str(last)})
                await self._sleep(delay)
        raise last or LLMError("unknown LLM failure", retryable=True)

    # ------------------------------------------------------------------ completions
    async def _chat(self, model: str, messages: list[dict[str, str]], *, max_tokens: int, temperature: float, json_mode: bool) -> tuple[str, dict[str, int], int]:
        payload: dict[str, Any] = {"model": model, "messages": messages, "max_tokens": max_tokens, "temperature": temperature, "stream": False}
        if json_mode:
            payload["response_format"] = {"type": "json_object"}
        t0 = time.monotonic()
        data = await self._post("/chat/completions", payload)
        latency = int((time.monotonic() - t0) * 1000)
        try:
            content = data["choices"][0]["message"]["content"] or ""
        except (KeyError, IndexError, TypeError) as exc:
            raise LLMError(f"malformed completion: {json.dumps(data)[:200]}") from exc
        usage = data.get("usage") or {}
        tokens = {"request": int(usage.get("prompt_tokens") or 0), "response": int(usage.get("completion_tokens") or 0)}
        return str(content), tokens, latency

    async def complete_text(self, model_tier: Tier, system: str, user: str, *, max_tokens: int = 400, temperature: float = 0.2) -> LLMResult[Any]:
        model = self.model_for(model_tier)
        content, tokens, latency = await self._chat(model, [{"role": "system", "content": system}, {"role": "user", "content": user}],
                                                    max_tokens=max_tokens, temperature=temperature, json_mode=False)
        text = _THINK.sub("", content).strip()
        return LLMResult(data=text, model=model, request_tokens=tokens["request"], response_tokens=tokens["response"], latency_ms=latency, raw=content)

    async def complete_json[T: BaseModel](self, model_tier: Tier, system: str, user: str, schema: type[T], *, max_tokens: int = 800,
                            temperature: float = 0.1) -> LLMResult[T]:
        model = self.model_for(model_tier)
        sys_prompt = system.rstrip() + "\n\nRespond with a single JSON object only. JSON schema:\n" + json.dumps(_compact_schema(schema), ensure_ascii=False)
        messages = [{"role": "system", "content": sys_prompt}, {"role": "user", "content": user}]
        content, tokens, latency = await self._chat(model, messages, max_tokens=max_tokens, temperature=temperature, json_mode=True)
        req, resp = tokens["request"], tokens["response"]
        try:
            return LLMResult(data=parse_json_object(content, schema), model=model, request_tokens=req, response_tokens=resp, latency_ms=latency, raw=content)
        except (ValueError, ValidationError) as exc:
            error = str(exc)[:800]
            log.info("llm output invalid, repairing", extra={"model": model, "schema": schema.__name__, "error": error[:160]})
        # one repair retry: show the model its own output and the validation error
        messages += [{"role": "assistant", "content": content[:4000]},
                     {"role": "user", "content": "Your previous answer was not valid for the schema. Error:\n" + error +
                      "\nReturn ONLY the corrected JSON object with the same meaning. No prose, no markdown."}]
        content2, tokens2, latency2 = await self._chat(model, messages, max_tokens=max_tokens, temperature=0.0, json_mode=True)
        req += tokens2["request"]
        resp += tokens2["response"]
        try:
            data = parse_json_object(content2, schema)
        except (ValueError, ValidationError) as exc:
            raise LLMValidationError(f"invalid after repair: {str(exc)[:300]}") from exc
        return LLMResult(data=data, model=model, request_tokens=req, response_tokens=resp, latency_ms=latency + latency2, attempts=2, repaired=True, raw=content2)

    async def embed(self, texts: list[str]) -> list[list[float]]:
        if not texts:
            return []
        data = await self._post("/embeddings", {"model": self.model_for("embedding"), "input": texts})
        items = sorted(data.get("data") or [], key=lambda d: d.get("index", 0))
        return [list(map(float, d["embedding"])) for d in items]

    async def health(self) -> Health:
        if not self.configured:
            return Health(ok=False, base_url="", error="not configured")
        t0 = time.monotonic()
        try:
            r = await self.client().get("/models")
            latency = int((time.monotonic() - t0) * 1000)
            if r.status_code >= 300:
                return Health(ok=False, base_url=self.base_url, latency_ms=latency, error=f"HTTP {r.status_code}")
            body = r.json()
            models = [m.get("id") for m in (body.get("data") or []) if isinstance(m, dict) and m.get("id")]
            return Health(ok=True, base_url=self.base_url, models=models, latency_ms=latency)
        except (httpx.HTTPError, ValueError) as exc:
            return Health(ok=False, base_url=self.base_url, latency_ms=int((time.monotonic() - t0) * 1000), error=str(exc)[:200])


# ---------------------------------------------------------------------------------------------------------------- parsing


def extract_json_text(content: str) -> str:
    """Tolerate reasoning tags, code fences and prose around the object; return the outermost {...} span."""
    text = _THINK.sub("", content or "").strip()
    m = _JSON_BLOCK.search(text)
    if m:
        return m.group(1)
    start = text.find("{")
    end = text.rfind("}")
    if start == -1 or end == -1 or end < start:
        raise ValueError("no JSON object in model output")
    return text[start:end + 1]


def parse_json_object[T: BaseModel](content: str, schema: type[T]) -> T:
    raw = extract_json_text(content)
    try:
        obj = json.loads(raw)
    except json.JSONDecodeError as exc:
        raise ValueError(f"invalid JSON: {exc.msg} at {exc.pos}") from exc
    if isinstance(obj, dict):
        return schema.model_validate(obj)
    raise ValueError("top-level JSON value must be an object")


def _compact_schema(schema: type[BaseModel]) -> dict[str, Any]:
    full = schema.model_json_schema()
    props = {}
    for name, spec in (full.get("properties") or {}).items():
        t = spec.get("type") or ("/".join(x.get("type", "?") for x in spec.get("anyOf", [])) if spec.get("anyOf") else "any")
        entry: dict[str, Any] = {"type": t}
        if spec.get("description"):
            entry["description"] = spec["description"]
        for k in ("maxLength", "minimum", "maximum", "maxItems"):
            if k in spec:
                entry[k] = spec[k]
        props[name] = entry
    return {"type": "object", "properties": props, "required": full.get("required", [])}


# ---------------------------------------------------------------------------------------------------------------- singleton

_provider: LLMProvider | None = None


def get_provider() -> LLMProvider:
    global _provider
    if _provider is None:
        _provider = OpenAICompatibleProvider()
    return _provider


def set_provider(provider: LLMProvider | None) -> None:
    """Tests / alternative vendors."""
    global _provider
    _provider = provider


__all__ = ["Health", "LLMError", "LLMNotConfigured", "LLMProvider", "LLMResult", "LLMValidationError", "OpenAICompatibleProvider", "extract_json_text",
           "get_provider", "parse_json_object", "set_provider"]