"""LLM access (architecture L9).

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. The base URL and token come from the
agent's config (nisse_dev.config), nothing hardcoded.

Model identity is expressed as roles (L9), resolved to concrete model strings
from 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 time
from dataclasses import 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 (L9's caps still bound the
# total).
RATE_LIMIT_STATUS = 429
RATE_LIMIT_RETRY_S = 20.0

# The proxy, or what's behind it, briefly gone: a 502/503/504 or a dropped
# connection. Outages delay, never fail (L5): waited out with backoff before
# the work is stopped. Not a 500 (it may be the request) and not a timeout
# (a hung model has had its request_timeout_s already).
OUTAGE_STATUSES = frozenset({502, 503, 504})
OUTAGE_BACKOFF_S = (5.0, 15.0, 45.0)

# Generous by default. Output is billed by what is produced, not by the
# ceiling, and a ceiling too low does not fail loudly — it truncates a tool
# call mid-argument and hands back something that parses. Measured: Opus
# writing an architecture plan routinely spends 3,000-4,000 output tokens and
# was observed at 4,064 against a 4,096 limit.
DEFAULT_MAX_OUTPUT_TOKENS = 16384

# How long one request may take before it is abandoned (llm.request_timeout_s).
# A safety net for a hung connection, not a budget: a dense local model
# writing a whole module in one turn was measured taking over 120 s, and
# timing it out failed work that was going well.
DEFAULT_REQUEST_TIMEOUT_S = 600.0


class UpstreamUnavailable(LLMError):
    """The proxy or provider could not be reached, or refused for reasons
    unrelated to what was asked: a network blip, a timeout, a 5xx, a rate
    limit that did not clear, an exhausted budget.

    Worth distinguishing because the *work* is not at fault. A delegation
    stopped by one of these has a branch that may be perfectly good, so it is
    preserved and resumed rather than discarded — unlike a delegate that made
    a mess, which is cleaned up.
    """


class BudgetExhausted(UpstreamUnavailable):
    """The proxy refused because a spend limit is gone.

    Distinct from the rest of its family because it will not clear on its
    own within the hour: an exhausted daily budget takes external models away
    until it resets, so retrying wastes time and the operator needs telling.
    L9 keeps the budget itself at the proxy; recognising its refusal is ours.
    """


def _looks_like_budget_refusal(text: str) -> bool:
    # Pattern-matched rather than contracted: LiteLLM's wording is not a
    # documented interface, so this errs toward missing one rather than
    # misreporting an unrelated failure as a budget problem.
    lowered = text.lower()
    return "budget" in lowered and (
        "exceed" in lowered or "exhaust" in lowered or "limit" in lowered
    )


class TruncatedResponse(LLMError):
    """The model was cut off mid-answer.

    Raised rather than returned, because the damage is silent: a truncated
    tool call arrives with required arguments simply absent, and a truncated
    write_file would write half a file and report success. A cut-off turn is
    a failed turn, and every caller already retries or gates on LLMError.
    """


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  # the old loop's (control.py), until M2.8
    planner: str  # the old loop's, until M2.8
    helper: str  # delegations (L9)
    expert: str  # ask_expert (L9)
    driver: str = ""  # drives the new loop; unset, the planner's model
    classifier: str = ""  # reads go-aheads now, actions later (L4); unset, expert's

    ROLES: ClassVar[tuple[str, ...]] = (
        "executor",
        "planner",
        "helper",
        "expert",
        "driver",
        "classifier",
    )
    _FALLBACK: ClassVar[dict[str, str]] = {"driver": "planner", "classifier": "expert"}

    def resolve(self, role: str) -> str:
        if role not in self.ROLES:
            raise LLMError(f"unknown role: {role!r}")
        model = str(getattr(self, role))
        if not model and role in self._FALLBACK:
            return self.resolve(self._FALLBACK[role])
        return model

    def pinned(self) -> dict[str, str]:
        """The full resolved mapping, to record at task start (L9)."""
        return {role: self.resolve(role) for role in self.ROLES}

    @classmethod
    def from_pinned(cls, pinned: dict[str, str], default: "RoleMap") -> "RoleMap":
        """A work item's pinned mapping; a role it predates comes from
        ``default``."""
        return cls(**{r: pinned.get(r) or default.resolve(r) for r in cls.ROLES})


class LiteLLMClient:
    def __init__(
        self,
        base_url: str,
        token: str,
        *,
        timeout_s: float = DEFAULT_REQUEST_TIMEOUT_S,
        transport: httpx.BaseTransport | None = None,
        sleep: Any = time.sleep,
        max_tokens: int | None = None,
    ) -> 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
        self._max_tokens = max_tokens or DEFAULT_MAX_OUTPUT_TOKENS
        # Set after each call so a caller can meter itself (L9 caps);
        # a side channel rather than a return value, so the LLMClient
        # protocol stays the same for every implementation.
        self.last_usage: dict[str, int] = {}
        self.last_model: str | None = None
        self.last_cost: float = 0.0
        # Set even when the call then fails: time spent waiting out a rate
        # limit is the proxy's, not the model's (the benchmark subtracts it).
        self.last_rate_limit_wait_s: float = 0.0

    def complete(
        self,
        events: list[StoredEvent],
        tools: list[dict[str, Any]] | None,
        model: str,
        system: str | None = None,
        *,
        max_tokens: int | None = None,
        require_tool_call: bool = False,
    ) -> list[CanonicalEvent]:
        body = self._renderer.render(events, tools, model, system=system)
        if require_tool_call and tools:
            # Only for models whose server enforces it; see BudgetedClient.
            body["tool_choice"] = "required"
        # A ceiling on this route too: without one a runaway turn is bounded
        # only by the provider's default, and the same truncation check has
        # nothing to check against. A caller that knows the model's limits
        # (the context budget) passes a tighter one.
        body["max_tokens"] = max_tokens or self._max_tokens
        self.last_rate_limit_wait_s = 0.0  # not the last call's, if this one raises
        response, self.last_rate_limit_wait_s = _post_with_rate_limit_retry(
            self._http, "/chat/completions", body, self._sleep
        )
        if response.status_code != 200:
            _raise_for_upstream(response)
            raise LLMError(
                f"LLM request returned {response.status_code}: {response.text[:500]}"
            )
        # LiteLLM prices every call; kept so a caller can meter spending.
        self.last_cost = _response_cost(response)
        return self._parse(response.json(), requested_model=model)

    def _parse(
        self, data: dict[str, Any], requested_model: str
    ) -> list[CanonicalEvent]:
        try:
            choice = data["choices"][0]
            message = choice["message"]
        except (KeyError, IndexError) as exc:
            raise LLMError(f"malformed completion response: {exc}") from exc
        if choice.get("finish_reason") == "length":
            raise TruncatedResponse(
                f"{requested_model} hit its output limit mid-answer; raise "
                "llm.max_output_tokens, or the model's max_output_tokens in LiteLLM"
            )
        produced_by = data.get("model") or requested_model
        # What actually answered, which is not always what was asked for: a
        # proxy fallback answers under another model's name.
        self.last_model = produced_by
        usage = data.get("usage") or {}
        self.last_usage = {
            "input": int(usage.get("prompt_tokens") or 0),
            "output": int(usage.get("completion_tokens") or 0),
        }
        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 (L7): 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 (L9).
            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 (L9). 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 (L9);
    ask_expert and fallbacks use the normalized route, where it does not.
    """

    def __init__(
        self,
        base_url: str,
        token: str,
        *,
        timeout_s: float = DEFAULT_REQUEST_TIMEOUT_S,
        max_tokens: int | None = None,
        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 or DEFAULT_MAX_OUTPUT_TOKENS
        self._sleep = sleep
        self.last_usage: dict[str, int] = {}
        self.last_model: str | None = None
        self.last_cost: float = 0.0
        # Set even when the call then fails: time spent waiting out a rate
        # limit is the proxy's, not the model's (the benchmark subtracts it).
        self.last_rate_limit_wait_s: float = 0.0

    def complete(
        self,
        events: list[StoredEvent],
        tools: list[dict[str, Any]] | None,
        model: str,
        system: str | None = None,
        *,
        max_tokens: int | None = None,
        require_tool_call: bool = False,
    ) -> list[CanonicalEvent]:
        body = self._renderer.render(
            events,
            tools,
            model,
            system=system,
            max_tokens=max_tokens or self._max_tokens,
        )
        if require_tool_call and tools:
            body["tool_choice"] = {"type": "any"}  # Anthropic's name for it
        self.last_rate_limit_wait_s = 0.0  # not the last call's, if this one raises
        response, self.last_rate_limit_wait_s = _post_with_rate_limit_retry(
            self._http, "/v1/messages", body, self._sleep
        )
        if response.status_code != 200:
            _raise_for_upstream(response)
            raise LLMError(
                f"passthrough returned {response.status_code}: {response.text[:500]}"
            )
        # LiteLLM prices every call; kept so a caller can meter spending.
        self.last_cost = _response_cost(response)
        return self._parse(response.json(), requested_model=model)

    def _parse(
        self, data: dict[str, Any], requested_model: str
    ) -> list[CanonicalEvent]:
        if data.get("stop_reason") == "max_tokens":
            raise TruncatedResponse(
                f"{requested_model} hit its output limit mid-answer; raise "
                "llm.max_output_tokens, or the model's max_output_tokens in LiteLLM"
            )
        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
        # What actually answered, which is not always what was asked for: a
        # proxy fallback answers under another model's name.
        self.last_model = produced_by
        usage = data.get("usage") or {}
        self.last_usage = {
            "input": int(usage.get("input_tokens") or 0),
            "output": int(usage.get("output_tokens") or 0),
        }
        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 (L7); 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: L7 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
) -> tuple[httpx.Response, float]:
    """The response, and how long was spent waiting: out a rate limit (once)
    or an outage (with backoff), not on the model."""
    waited = 0.0
    rate_limited = False
    outages = iter(OUTAGE_BACKOFF_S)
    while True:
        try:
            response = http.post(path, json=body)
        except httpx.TimeoutException as exc:
            raise UpstreamUnavailable(f"LLM request failed: {exc}") from exc
        except httpx.HTTPError as exc:
            # Could not reach the proxy at all: nothing about the request was
            # wrong, so callers holding work in progress should keep it.
            pause = next(outages, None)
            if pause is None:
                raise UpstreamUnavailable(f"LLM request failed: {exc}") from exc
            sleep(pause)
            waited += pause
            continue
        if response.status_code == RATE_LIMIT_STATUS and not rate_limited:
            rate_limited = True
            sleep(RATE_LIMIT_RETRY_S)
            waited += RATE_LIMIT_RETRY_S
            continue
        if response.status_code in OUTAGE_STATUSES and not _looks_like_budget_refusal(
            response.text
        ):
            pause = next(outages, None)
            if pause is not None:
                sleep(pause)
                waited += pause
                continue
        return response, waited


# Model families we have a native renderer adapter for. L9: delegations take
# the vendor's passthrough route because a normalising middlebox is lossy
# exactly where a long-running delegation lives.
def _response_cost(response: httpx.Response) -> float:
    try:
        return float(response.headers.get("x-litellm-response-cost") or 0.0)
    except ValueError:
        return 0.0


def _raise_for_upstream(response: httpx.Response) -> None:
    """Classify a non-200 as upstream trouble where it plainly is.

    A 4xx that is not a rate limit means the request itself was wrong, and
    preserving work on the strength of one would hide a real bug — so the
    benefit of the doubt stops at the server's own door.
    """
    if _looks_like_budget_refusal(response.text):
        raise BudgetExhausted(
            f"the proxy refused on budget grounds: {response.text[:300]}"
        )
    if response.status_code >= 500 or response.status_code == RATE_LIMIT_STATUS:
        raise UpstreamUnavailable(
            f"upstream returned {response.status_code}: {response.text[:300]}"
        )