"""Client for OpenAI-compatible chat endpoints.

Deliberately plain httpx rather than the official SDK. The target is not just
api.openai.com but LM Studio, vLLM, llama.cpp, Ollama's compatibility layer,
OpenRouter and anything else exposing /v1 -- and they differ in small ways. A
thin client passes request parameters through untouched and is tolerant about
what comes back, which is exactly what talking to all of them requires.
"""

from __future__ import annotations

import json
import logging
from collections.abc import AsyncIterator
from dataclasses import dataclass
from typing import Any

import httpx

from lembas.config import settings
from lembas.db.models import Connection
from lembas.services.crypto import decrypt

log = logging.getLogger(__name__)


class LLMError(Exception):
    """An upstream failure with a message fit to show a user.

    Every failure path in this module raises this rather than letting an httpx
    or JSON exception escape, so callers have exactly one thing to catch and
    the chat UI always has something intelligible to display.
    """

    def __init__(self, message: str, *, status_code: int | None = None) -> None:
        super().__init__(message)
        self.message = message
        self.status_code = status_code


@dataclass(frozen=True)
class Endpoint:
    """Everything needed to call a connection, with the key already decrypted.

    A frozen snapshot rather than the ORM object because streaming outlives the
    request that started it, and a detached SQLAlchemy instance is a trap.
    """

    base_url: str
    api_key: str
    extra_headers: dict[str, str]
    name: str = ""

    @classmethod
    def from_connection(cls, connection: Connection) -> Endpoint:
        return cls(
            base_url=connection.base_url.rstrip("/"),
            api_key=decrypt(connection.api_key_encrypted),
            extra_headers=dict(connection.extra_headers_json or {}),
            name=connection.name,
        )

    def url(self, path: str) -> str:
        # Accept both "http://host:1234" and "http://host:1234/v1" so users do
        # not have to guess which form this expects.
        base = self.base_url
        if not base.endswith("/v1") and "/v1/" not in base:
            base = f"{base}/v1"
        return f"{base}/{path.lstrip('/')}"

    def headers(self) -> dict[str, str]:
        headers = {"Content-Type": "application/json", **self.extra_headers}
        # Local endpoints frequently need no key at all; sending an empty
        # bearer token makes some of them reject the request outright.
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        return headers


def describe_http_error(exc: httpx.HTTPStatusError) -> str:
    """Turn an upstream error response into something worth reading.

    Public because the audio and search clients talk to the same class of
    server and want the same translation; LLMError stays the one thing a
    caller has to catch.

    Providers put the useful part in wildly different places, so try the common
    shapes before falling back to the raw body.
    """
    status = exc.response.status_code
    detail = ""
    try:
        payload = exc.response.json()
        if isinstance(payload, dict):
            error = payload.get("error")
            if isinstance(error, dict):
                detail = error.get("message", "")
            elif isinstance(error, str):
                detail = error
            detail = detail or payload.get("message", "") or payload.get("detail", "")
    except (ValueError, json.JSONDecodeError):
        detail = exc.response.text[:300]

    friendly = {
        401: "The API key was rejected.",
        403: "The API key is not permitted to use this model.",
        404: "The endpoint or model was not found.",
        429: "Rate limited by the provider.",
    }.get(status)

    if friendly and detail:
        return f"{friendly} {detail}"
    return friendly or detail or f"The endpoint returned HTTP {status}."


def wrap_transport_error(exc: httpx.RequestError, endpoint: Endpoint) -> LLMError:
    if isinstance(exc, httpx.ConnectError):
        return LLMError(
            f"Could not reach {endpoint.base_url}. Is the endpoint running and "
            f"the URL correct?"
        )
    if isinstance(exc, httpx.TimeoutException):
        return LLMError(
            f"{endpoint.base_url} did not respond within "
            f"{settings.request_timeout:.0f}s."
        )
    return LLMError(f"Could not reach {endpoint.base_url}: {exc}")


async def list_models(endpoint: Endpoint) -> list[dict[str, Any]]:
    """Fetch the models a connection advertises via GET /v1/models."""
    try:
        async with httpx.AsyncClient(timeout=30.0) as client:
            response = await client.get(endpoint.url("models"), headers=endpoint.headers())
            response.raise_for_status()
            payload = response.json()
    except httpx.HTTPStatusError as exc:
        raise LLMError(describe_http_error(exc), status_code=exc.response.status_code) from exc
    except httpx.RequestError as exc:
        raise wrap_transport_error(exc, endpoint) from exc
    except ValueError as exc:
        raise LLMError("The endpoint returned a response that was not JSON.") from exc

    # The spec says {"data": [...]}, but some servers return a bare list.
    entries = payload.get("data", payload) if isinstance(payload, dict) else payload
    if not isinstance(entries, list):
        raise LLMError("The endpoint's model list was not in the expected format.")

    models = []
    for entry in entries:
        if isinstance(entry, dict) and entry.get("id"):
            models.append(entry)
        elif isinstance(entry, str):
            models.append({"id": entry})
    return models


# Where the runners that bother to advertise a context length put it. There is
# no standard field, so this is a list of what the common ones actually emit.
_CONTEXT_KEYS = ("context_length", "max_model_len", "context_window", "max_context_length")

# Below the first, the number is not a context length; above the second it is a
# typo or a different unit. Either way, better to record nothing than a wrong
# figure a percentage would then be computed from.
MIN_CONTEXT = 256
MAX_CONTEXT = 10_000_000


def context_from(entry: dict[str, Any]) -> int:
    """A model's context length as advertised by /v1/models, or 0 if it is not.

    Strings are accepted because some servers quote the number, but only when
    they are digits alone -- "8192 tokens" is a label, not a measurement.
    """
    candidates = [entry.get(key) for key in _CONTEXT_KEYS]
    meta = entry.get("meta")
    if isinstance(meta, dict):
        candidates += [meta.get("n_ctx"), *(meta.get(key) for key in _CONTEXT_KEYS)]

    for value in candidates:
        if isinstance(value, bool):
            continue
        if isinstance(value, str):
            value = value.strip()
            if not value.isdigit():
                continue
            value = int(value)
        if isinstance(value, int) and MIN_CONTEXT <= value <= MAX_CONTEXT:
            return value
    return 0


# Endpoints that rejected `stream_options`, so it is asked for once per base URL
# per process and then never again. Not persisted: it is a property of whatever
# is running there now, and a restart is the right time to find out afresh.
_NO_STREAM_OPTIONS: set[str] = set()


async def stream_chat(
    endpoint: Endpoint,
    payload: dict[str, Any],
) -> AsyncIterator[dict[str, Any]]:
    """Stream a chat completion, yielding each parsed SSE data object.

    Yields the raw upstream chunks; interpreting them is the caller's job. The
    terminating "[DONE]" sentinel is consumed here and not yielded.

    `stream_options` asks for the final usage chunk, which is the only way to
    learn what a streamed reply actually cost. Not every server implements it,
    and an unknown key is a 400 from some of them -- the same hazard as sending
    a `tools` array to an endpoint without support. So it is asked for once,
    and an endpoint that refuses is remembered and never asked again. Retrying
    is safe because the status is checked before a single line is read: nothing
    has been yielded, so there is nothing to duplicate.
    """
    wants_usage = endpoint.base_url not in _NO_STREAM_OPTIONS

    try:
        async for chunk in _stream_once(endpoint, payload, usage=wants_usage):
            yield chunk
    except LLMError as exc:
        if not wants_usage or exc.status_code not in (400, 422):
            raise
        _NO_STREAM_OPTIONS.add(endpoint.base_url)
        log.info(
            "%s rejected stream_options; token counts will be estimated there",
            endpoint.base_url,
        )
        async for chunk in _stream_once(endpoint, payload, usage=False):
            yield chunk


async def _stream_once(
    endpoint: Endpoint,
    payload: dict[str, Any],
    *,
    usage: bool,
) -> AsyncIterator[dict[str, Any]]:
    body = {**payload, "stream": True}
    if usage:
        body["stream_options"] = {"include_usage": True}

    try:
        async with (
            httpx.AsyncClient(timeout=settings.request_timeout) as client,
            client.stream(
                "POST",
                endpoint.url("chat/completions"),
                headers=endpoint.headers(),
                json=body,
            ) as response,
        ):
            if response.status_code >= 400:
                # The body has not been read yet on a streaming response, and
                # the error detail is in it.
                await response.aread()
                response.raise_for_status()

            async for line in response.aiter_lines():
                line = line.strip()
                if not line or line.startswith(":"):
                    continue  # keep-alive comment
                if not line.startswith("data:"):
                    continue
                data = line[5:].strip()
                if data == "[DONE]":
                    return
                try:
                    yield json.loads(data)
                except json.JSONDecodeError:
                    # A malformed frame is not worth killing a reply over.
                    log.warning("skipping unparseable SSE frame: %.120s", data)
                    continue

    except httpx.HTTPStatusError as exc:
        raise LLMError(describe_http_error(exc), status_code=exc.response.status_code) from exc
    except httpx.RequestError as exc:
        raise wrap_transport_error(exc, endpoint) from exc


async def complete(endpoint: Endpoint, payload: dict[str, Any]) -> str:
    """Non-streaming completion. Used for short internal calls like auto-titling."""
    body = {**payload, "stream": False}
    try:
        async with httpx.AsyncClient(timeout=60.0) as client:
            response = await client.post(
                endpoint.url("chat/completions"), headers=endpoint.headers(), json=body
            )
            response.raise_for_status()
            data = response.json()
    except httpx.HTTPStatusError as exc:
        raise LLMError(describe_http_error(exc), status_code=exc.response.status_code) from exc
    except httpx.RequestError as exc:
        raise wrap_transport_error(exc, endpoint) from exc
    except ValueError as exc:
        raise LLMError("The endpoint returned a response that was not JSON.") from exc

    try:
        return data["choices"][0]["message"]["content"] or ""
    except (KeyError, IndexError, TypeError) as exc:
        raise LLMError("The endpoint returned no completion.") from exc


def delta_reasoning(chunk: dict[str, Any]) -> str:
    """Pull a reasoning delta out of one streamed chunk.

    Providers disagree on the field name -- llama.cpp, llama-swap and vLLM use
    ``reasoning_content``, some others just ``reasoning`` -- so both are read.
    Models that emit ``<think>`` tags inline in ``content`` instead are handled
    by lembas.services.reasoning.
    """
    try:
        choices = chunk.get("choices") or []
        if not choices:
            return ""
        delta = choices[0].get("delta") or {}
        for field in ("reasoning_content", "reasoning"):
            value = delta.get(field)
            if isinstance(value, str) and value:
                return value
        return ""
    except (AttributeError, TypeError):
        return ""


def delta_tool_calls(chunk: dict[str, Any]) -> list[dict[str, Any]]:
    """Pull tool-call fragments out of one streamed chunk.

    Each entry carries an ``index`` and, across chunks, a name that arrives
    once and an ``arguments`` string that arrives in pieces. Reassembling them
    is lembas.services.tools.ToolCallAccumulator's job; this only extracts.
    """
    try:
        choices = chunk.get("choices") or []
        if not choices:
            return []
        calls = (choices[0].get("delta") or {}).get("tool_calls")
        return calls if isinstance(calls, list) else []
    except (AttributeError, TypeError):
        return []


def finish_reason(chunk: dict[str, Any]) -> str:
    """Why the model stopped, when the chunk says so.

    ``tool_calls`` here is the signal that the reply is not an answer but a
    request to run something and come back. Some servers send ``stop`` even
    when they emitted tool calls, so the accumulator's contents are the real
    authority and this is only a hint.
    """
    try:
        choices = chunk.get("choices") or []
        if not choices:
            return ""
        return choices[0].get("finish_reason") or ""
    except (AttributeError, TypeError):
        return ""


def chunk_usage(chunk: dict[str, Any]) -> dict[str, int] | None:
    """Token counts from a usage chunk, or None if this is not one.

    A usage chunk carries `choices: []`, which is exactly the shape delta_text,
    delta_reasoning, delta_tool_calls and finish_reason all return early on --
    they have always tolerated it, so nothing else needs to change to let one
    through.

    Fields are read defensively because "the endpoint returned something odd"
    must never be the reason a reply fails; a bad shape simply means no counts.
    """
    try:
        raw = chunk.get("usage")
    except AttributeError:
        return None
    if not isinstance(raw, dict):
        return None

    counts: dict[str, int] = {}
    for key in ("prompt_tokens", "completion_tokens", "total_tokens"):
        value = raw.get(key)
        if isinstance(value, bool) or not isinstance(value, (int, float)):
            continue
        if value >= 0:
            counts[key] = int(value)

    # Some servers send a usage object of zeros on every chunk and the real
    # numbers only at the end. All-zero is indistinguishable from that, and
    # treating it as an answer would freeze the count at nothing.
    if not counts or not any(counts.values()):
        return None
    counts.setdefault(
        "total_tokens", counts.get("prompt_tokens", 0) + counts.get("completion_tokens", 0)
    )
    return counts


def delta_text(chunk: dict[str, Any]) -> str:
    """Pull the text out of one streamed chunk, tolerating provider variation."""
    try:
        choices = chunk.get("choices") or []
        if not choices:
            return ""
        delta = choices[0].get("delta") or {}
        content = delta.get("content")
        if isinstance(content, str):
            return content
        # Some providers send content as a list of typed parts even in deltas.
        if isinstance(content, list):
            return "".join(
                part.get("text", "")
                for part in content
                if isinstance(part, dict) and part.get("type") == "text"
            )
        return ""
    except (AttributeError, TypeError):
        return ""