"""Async OpenAI-compatible client for the self-hosted aipi endpoint (PLAN A5).

Provides the embeddings surface (importer, retrieval) and chat streaming
(PLAN A15) for the RAG pipeline.

Fail-loud rule (PLAN A6): the ``chunks.embedding`` column is fixed at 768
dimensions when the table is created, so a model that returns a different
dimension must abort the import with an actionable error — never store
vectors that pgvector rejects.
"""
from __future__ import annotations

import logging
from collections.abc import AsyncIterator
from typing import cast

from openai import AsyncOpenAI
from openai.types.chat import ChatCompletionMessageParam

from app.config import Settings, get_settings

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


class EmbeddingError(RuntimeError):
    """The embeddings endpoint failed (network, HTTP, or malformed reply)."""


class EmbeddingDimensionError(EmbeddingError):
    """Embedding dimension != BOR_EMBEDDING_DIM — import must fail loudly."""


class LLMError(RuntimeError):
    """The chat-completions endpoint failed (network, HTTP, or mid-stream)."""


# aipi's local embedding model rejects requests over ~1024 input tokens
# ("input is too large to process"). Batch by estimated tokens, with a
# safety margin under that cap — code-dense text can tokenize at ~3
# chars/token, so stay well below the ceiling. A retry that halves an
# over-large batch (:meth:`LLMClient._embed_batch`) is the backstop.
_TOKENS_PER_CHAR = 0.25
_REQUEST_TOKEN_BUDGET = 700


class _TooLarge(RuntimeError):
    """Internal: the endpoint rejected the request's input size."""


class LLMClient:
    """Thin async wrapper over the aipi OpenAI-compatible API."""

    def __init__(self, settings: Settings | None = None) -> None:
        self.settings = settings or get_settings()
        #: Number of embedding HTTP requests made so far (importer logging).
        self.embed_batches: int = 0
        self._client = AsyncOpenAI(
            base_url=self.settings.llm_base_url,
            api_key=self.settings.effective_api_key,
            timeout=120.0,
        )

    async def _post_embeddings(self, texts: list[str]) -> list[list[float]]:
        """One POST /embeddings with a minimal OpenAI-compatible payload.

        The ``openai`` SDK (1.x and 2.x) injects ``encoding_format`` into
        every embeddings request (defaulting to ``"base64"``), and aipi's
        litellm ``openai_like`` model group rejects that parameter in any
        form — so we reuse the openai client's own httpx transport (same
        base URL, TLS, and connection pooling) and send a clean payload.
        The endpoint's default is floats, which is what pgvector needs.
        """
        http = self._client._client  # pyright: ignore[reportAttributeAccessIssue]
        resp = await http.post(
            "embeddings",
            json={"model": self.settings.llm_embed_model, "input": texts},
            headers={"Authorization": f"Bearer {self.settings.effective_api_key}"},
        )
        if resp.status_code >= 400:
            if "too large" in resp.text:
                raise _TooLarge(resp.text[:300])
            raise EmbeddingError(
                f"embeddings endpoint returned HTTP {resp.status_code}: {resp.text[:300]}"
            )
        payload = resp.json()
        rows = payload.get("data") if isinstance(payload, dict) else None
        if not isinstance(rows, list) or len(rows) != len(texts):
            n = len(rows) if isinstance(rows, list) else "?"
            raise EmbeddingError(
                f"embeddings endpoint returned {n} vectors for {len(texts)} inputs — "
                "refusing to guess which is which"
            )
        ordered = sorted(rows, key=lambda r: r["index"])
        return [[float(x) for x in row["embedding"]] for row in ordered]

    async def _embed_batch(self, chunk: list[str]) -> list[list[float]]:
        """Embed *chunk*, halving the request if the endpoint says the input
        is too large (tokenizer estimates can be wrong for dense content).
        A single text that still fails is a hard, actionable error."""
        try:
            return await self._post_embeddings(chunk)
        except _TooLarge:
            if len(chunk) == 1:
                raise EmbeddingError(
                    f"a single {len(chunk[0])}-char chunk exceeded the endpoint's "
                    "per-request input token cap — lower BOR_CHUNK_TARGET_CHARS "
                    "and re-import"
                ) from None
            mid = len(chunk) // 2
            left = await self._embed_batch(chunk[:mid])
            right = await self._embed_batch(chunk[mid:])
            return [*left, *right]

    def _check_dims(self, vecs: list[list[float]]) -> None:
        expected = self.settings.embedding_dim
        dims = sorted({len(v) for v in vecs})
        if dims != [expected]:
            raise EmbeddingDimensionError(
                f"embedding dimension mismatch: model '{self.settings.llm_embed_model}' "
                f"returned dims {dims} but BOR_EMBEDDING_DIM={expected}. The chunks "
                "table stores a fixed dimension — run `uv run python -m scripts.llm_probe`,"
                " update BOR_EMBEDDING_DIM, and recreate the schema (README → "
                "Troubleshooting: 'Embedding dimension mismatch')."
            )

    def _batch_texts(self, texts: list[str]) -> list[list[str]]:
        """Group *texts* into requests under the endpoint's token cap.

        ``BOR_EMBED_BATCH_SIZE`` stays a hard cap on *texts* per request;
        the token budget usually binds first for 2000-char chunks.
        """
        budget_chars = int(_REQUEST_TOKEN_BUDGET / _TOKENS_PER_CHAR)
        max_texts = max(1, self.settings.embed_batch_size)
        batches: list[list[str]] = []
        cur: list[str] = []
        cur_chars = 0
        for text in texts:
            if cur and (len(cur) >= max_texts or cur_chars + len(text) > budget_chars):
                batches.append(cur)
                cur, cur_chars = [], 0
            cur.append(text)
            cur_chars += len(text)
        if cur:
            batches.append(cur)
        return batches

    async def embed(self, texts: list[str]) -> list[list[float]]:
        """Embed *texts* in token-budgeted batches (order always kept)."""
        if not texts:
            return []
        out: list[list[float]] = []
        for chunk in self._batch_texts(texts):
            try:
                vecs = await self._embed_batch(chunk)
            except EmbeddingError:
                raise
            except Exception as e:  # noqa: BLE001 — wrap transport-level failures
                raise EmbeddingError(
                    f"embeddings request to {self.settings.llm_base_url} failed: {e}"
                ) from e
            self.embed_batches += 1
            self._check_dims(vecs)
            out.extend(vecs)
        return out

    async def embed_one(self, text: str) -> list[float]:
        """Convenience: embed a single text (retrieval path)."""
        (vec,) = await self.embed([text])
        return vec

    async def chat_stream(self, messages: list[dict[str, str]]) -> AsyncIterator[str]:
        """Stream assistant text deltas from the chat model (PLAN A5/A15).

        ``stream=True`` against the OpenAI-compatible endpoint; yields only
        non-empty ``delta.content`` pieces. Any failure (network, HTTP,
        malformed stream) surfaces as :class:`LLMError` so the API layer can
        turn it into an SSE ``error`` event instead of a hung request.

        Answers are allowed up to ``BOR_MAX_OUTPUT_TOKENS`` (default 32 768)
        output tokens — the old hard 700-token cap cut long answers off
        mid-sentence (owner report 2026-08-22).
        """
        try:
            # ``{role, content}`` dicts are exactly what the message params
            # accept; the cast keeps pyright honest about the SDK's union.
            stream = await self._client.chat.completions.create(
                model=self.settings.llm_chat_model,
                messages=cast("list[ChatCompletionMessageParam]", messages),
                temperature=0.4,
                max_tokens=self.settings.max_output_tokens,
                stream=True,
            )
            async for chunk in stream:
                if not chunk.choices:
                    continue
                piece = chunk.choices[0].delta.content
                if piece:
                    yield piece
        except LLMError:
            raise
        except Exception as e:  # noqa: BLE001 — wrap transport-level failures
            raise LLMError(f"chat stream from {self.settings.llm_base_url} failed: {e}") from e