"""LLM access (spec A5, A11).

One client interface, present from the first real LLM call. All traffic goes
to the LiteLLM proxy — the VM's single LLM destination; the agent holds one
LiteLLM token and no vendor credentials. Config comes from the environment
(``NISSE_LLM_BASE_URL``, ``NISSE_LLM_TOKEN``), nothing hardcoded.

Model identity is expressed as roles (A11), resolved to concrete model strings
from env config; the resolved mapping is pinned at task start so a mid-task
config change cannot split a task across two planners.
"""

import json
import os
import time
from collections.abc import Mapping
from dataclasses import asdict, dataclass
from typing import Any, ClassVar, Protocol

import httpx

from nisse_dev.canonical import (
    AssistantText,
    CanonicalEvent,
    StoredEvent,
    ToolCall,
    mint_call_id,
)
from nisse_dev.render import OpenAIChatRenderer
from nisse_dev.render_anthropic import AnthropicMessagesRenderer


class LLMError(RuntimeError):
    pass


# A rate limit is a transient condition, not a bad request. Losing a
# ten-turn delegation to one wastes far more tokens than waiting out the
# window, so retry once before surfacing it (A13's caps still bound the
# total).
RATE_LIMIT_STATUS = 429
RATE_LIMIT_RETRY_S = 20.0


class LLMClient(Protocol):
    def complete(
        self,
        events: list[StoredEvent],
        tools: list[dict[str, Any]] | None,
        model: str,
        system: str | None = None,
    ) -> list[CanonicalEvent]: ...


@dataclass(frozen=True)
class RoleMap:
    executor: str
    planner: str
    delegate: str
    consultant: str

    _ENV_VARS: ClassVar[dict[str, str]] = {
        "executor": "NISSE_ROLE_EXECUTOR",
        "planner": "NISSE_ROLE_PLANNER",
        "delegate": "NISSE_ROLE_DELEGATE",
        "consultant": "NISSE_ROLE_CONSULTANT",
    }

    @classmethod
    def from_env(cls, env: Mapping[str, str] = os.environ) -> "RoleMap":
        missing = [var for var in cls._ENV_VARS.values() if not env.get(var)]
        if missing:
            raise LLMError(f"role mapping incomplete; set: {', '.join(missing)}")
        return cls(**{role: env[var] for role, var in cls._ENV_VARS.items()})

    def resolve(self, role: str) -> str:
        if role not in self._ENV_VARS:
            raise LLMError(f"unknown role: {role!r}")
        return str(getattr(self, role))

    def pinned(self) -> dict[str, str]:
        """The full resolved mapping, to record at task start (A11)."""
        return asdict(self)


class LiteLLMClient:
    def __init__(
        self,
        base_url: str,
        token: str,
        *,
        timeout_s: float = 120.0,
        transport: httpx.BaseTransport | None = None,
        sleep: Any = time.sleep,
    ) -> None:
        self._http = httpx.Client(
            base_url=base_url.rstrip("/"),
            headers={"Authorization": f"Bearer {token}"},
            timeout=timeout_s,
            transport=transport,
        )
        self._renderer = OpenAIChatRenderer()
        self._sleep = sleep

    @classmethod
    def from_env(cls, env: Mapping[str, str] = os.environ) -> "LiteLLMClient":
        base_url = env.get("NISSE_LLM_BASE_URL")
        token = env.get("NISSE_LLM_TOKEN")
        if not base_url or not token:
            raise LLMError("set NISSE_LLM_BASE_URL and NISSE_LLM_TOKEN")
        return cls(base_url, token)

    def complete(
        self,
        events: list[StoredEvent],
        tools: list[dict[str, Any]] | None,
        model: str,
        system: str | None = None,
    ) -> list[CanonicalEvent]:
        body = self._renderer.render(events, tools, model, system=system)
        response = _post_with_rate_limit_retry(
            self._http, "/chat/completions", body, self._sleep
        )
        if response.status_code != 200:
            raise LLMError(
                f"LLM request returned {response.status_code}: {response.text[:500]}"
            )
        return self._parse(response.json(), requested_model=model)

    def _parse(
        self, data: dict[str, Any], requested_model: str
    ) -> list[CanonicalEvent]:
        try:
            message = data["choices"][0]["message"]
        except (KeyError, IndexError) as exc:
            raise LLMError(f"malformed completion response: {exc}") from exc
        produced_by = data.get("model") or requested_model
        out: list[CanonicalEvent] = []
        if message.get("content"):
            out.append(AssistantText(text=message["content"], model=produced_by))
        # Provider tool-call IDs are mapped at the boundary (A10): each call
        # gets a locally-minted canonical ID; the provider's is dropped, since
        # every future request is rendered wholly from the canonical log.
        for call in message.get("tool_calls") or []:
            raw = call["function"]["arguments"] or "{}"
            # Weak models emit truncated or non-JSON argument strings; that is
            # a bad model turn, not a client crash (A15).
            try:
                arguments = json.loads(raw)
            except json.JSONDecodeError as exc:
                raise LLMError(
                    f"tool call {call['function']['name']!r} had non-JSON "
                    f"arguments: {raw[:200]!r}"
                ) from exc
            if not isinstance(arguments, dict):
                raise LLMError(
                    f"tool call {call['function']['name']!r} arguments were "
                    f"{type(arguments).__name__}, expected object"
                )
            out.append(
                ToolCall(
                    call_id=mint_call_id(),
                    tool=call["function"]["name"],
                    arguments=arguments,
                )
            )
        return out

    def close(self) -> None:
        self._http.close()


class AnthropicPassthroughClient:
    """LiteLLM's Anthropic passthrough (A5). Same interface as LiteLLMClient,
    so a caller chooses a route rather than a different kind of client.

    Delegations take this path because fidelity matters there (A13);
    ask_expert and fallbacks use the normalized route, where it does not.
    """

    def __init__(
        self,
        base_url: str,
        token: str,
        *,
        timeout_s: float = 300.0,
        max_tokens: int = 4096,
        transport: httpx.BaseTransport | None = None,
        sleep: Any = time.sleep,
    ) -> None:
        self._http = httpx.Client(
            base_url=base_url.rstrip("/"),
            headers={
                "Authorization": f"Bearer {token}",
                "x-api-key": token,
                "anthropic-version": "2023-06-01",
            },
            timeout=timeout_s,
            transport=transport,
        )
        self._renderer = AnthropicMessagesRenderer()
        self._max_tokens = max_tokens
        self._sleep = sleep

    @classmethod
    def from_env(
        cls, env: Mapping[str, str] = os.environ
    ) -> "AnthropicPassthroughClient":
        base_url = env.get("NISSE_LLM_BASE_URL")
        token = env.get("NISSE_LLM_TOKEN")
        if not base_url or not token:
            raise LLMError("set NISSE_LLM_BASE_URL and NISSE_LLM_TOKEN")
        return cls(base_url, token)

    def complete(
        self,
        events: list[StoredEvent],
        tools: list[dict[str, Any]] | None,
        model: str,
        system: str | None = None,
    ) -> list[CanonicalEvent]:
        body = self._renderer.render(
            events, tools, model, system=system, max_tokens=self._max_tokens
        )
        response = _post_with_rate_limit_retry(
            self._http, "/v1/messages", body, self._sleep
        )
        if response.status_code != 200:
            raise LLMError(
                f"passthrough returned {response.status_code}: {response.text[:500]}"
            )
        return self._parse(response.json(), requested_model=model)

    def _parse(
        self, data: dict[str, Any], requested_model: str
    ) -> list[CanonicalEvent]:
        content = data.get("content")
        if not isinstance(content, list):
            raise LLMError(f"malformed passthrough response: {str(data)[:300]}")
        produced_by = data.get("model") or requested_model
        out: list[CanonicalEvent] = []
        for block in content:
            kind = block.get("type")
            if kind == "text" and block.get("text"):
                out.append(AssistantText(text=block["text"], model=produced_by))
            elif kind == "tool_use":
                # Provider ids are mapped at the boundary and never used as a
                # join key (A10); the canonical log owns identity.
                out.append(
                    ToolCall(
                        call_id=mint_call_id(),
                        tool=block["name"],
                        arguments=block.get("input") or {},
                    )
                )
            # Vendor-opaque reasoning blocks (thinking, redacted_thinking) are
            # deliberately dropped for now: A10 requires storing them tagged
            # with the producing model and never replaying them outside the
            # session that made them, and the canonical schema has no event
            # for that yet. Nothing requests extended thinking, so none arrive.
        return out

    def close(self) -> None:
        self._http.close()


def _post_with_rate_limit_retry(
    http: httpx.Client, path: str, body: dict[str, Any], sleep: Any
) -> httpx.Response:
    for attempt in (1, 2):
        try:
            response = http.post(path, json=body)
        except httpx.HTTPError as exc:
            raise LLMError(f"LLM request failed: {exc}") from exc
        if response.status_code != RATE_LIMIT_STATUS or attempt == 2:
            return response
        sleep(RATE_LIMIT_RETRY_S)
    raise AssertionError("unreachable")