"""Kernel model client (docs/kernel_spec.md §2).

One streaming interface over any OpenAI-compatible server. Owns native
``tool_calls`` accumulation, final-chunk harvest, one-shot non-stream
recovery, and the single permitted content-XML salvage (per
.cursor/rules/native-tool-calls-invariant.mdc). Nothing else in the kernel
parses model output.

Boundary: stdlib + httpx only. Server lifecycle (MLX restart, model-id
probing) is a peripheral concern — this client surfaces typed errors and the
caller decides.
"""
from __future__ import annotations

import json
import logging
import time
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any

import httpx

from core._salvage import (
    _TOOL_CALL_MARKER,
    ContentSalvage,
    normalize_call_ids,
    resolve_content_salvage,
    salvage_xml_tool_calls,
)
from core._sse import (
    FIRST_CHUNK_TIMEOUT,
    RETRY_ATTEMPTS,
    RETRY_BASE_DELAY,
    RETRY_MAX_DELAY,
    SSE_DATA_PREFIX,
    SSEIdleTimeout,
    is_retryable_error,
    iter_with_timeout,
    scaled_idle_timeout,
)
from core._types import ToolCall

_log = logging.getLogger(__name__)

# ToolCall is re-exported from core._types and the content salvage helpers from
# core._salvage so existing importers (core.loop, core.tools) keep working while
# the salvage logic lives in its own module (kept model.py under the 500-line
# kernel budget and breaks the model<->salvage import cycle).
__all__ = ["ModelClient", "ModelTurn", "ToolCall", "salvage_xml_tool_calls"]


# ─── Types ───────────────────────────────────────────────────────────────────
@dataclass
class ModelTurn:
    """The result of one model turn (streamed or not)."""

    content: str = ""
    tool_calls: list[ToolCall] = field(default_factory=list)
    finish_reason: str | None = None
    # None | "non_stream" | "content_salvage" — how tool_calls were obtained
    # when the primary stream path didn't carry them natively.
    recovered_via: str | None = None


class ModelError(Exception):
    """Base class for model client errors."""

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


class ModelConnectionError(ModelError):
    """Request could not be made or was rejected before streaming began."""

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


class ModelStreamError(ModelError):
    """Stream crashed mid-turn. Carries whatever arrived before the crash."""

    def __init__(
        self,
        message: str,
        *,
        partial_content: str = "",
        partial_tool_calls: list[ToolCall] | None = None,
    ) -> None:
        super().__init__(message)
        self.partial_content = partial_content
        self.partial_tool_calls = partial_tool_calls or []


class StreamIdleTimeout(ModelStreamError):
    """No SSE chunk arrived within the idle deadline."""


# ─── Native tool_calls accumulation / harvest (port, dict-shaped) ────────────
def _tc_field(tc: dict[str, Any], fieldname: str) -> Any:
    fn = tc.get("function") or tc
    return fn.get(fieldname) if isinstance(fn, dict) else None


def resolve_tool_call_index(accum: dict[int, dict[str, str]], tc: dict[str, Any]) -> int:
    """Resolve a streaming tool_call slot when the backend omits ``index``."""
    idx = tc.get("index")
    if idx is not None:
        return int(idx)
    tid = tc.get("id")
    if tid:
        existing = next((k for k, v in accum.items() if v.get("id") == tid), None)
        if existing is not None:
            return existing
        return (max(accum.keys()) + 1) if accum else 0
    return 0 if not accum else min(accum)


def accumulate_tool_call_delta(accum: dict[int, dict[str, str]], tc: dict[str, Any]) -> None:
    """Merge one incremental ``delta.tool_calls`` fragment into the accumulator."""
    idx = resolve_tool_call_index(accum, tc)
    slot = accum.setdefault(idx, {"id": "", "name": "", "arguments": ""})
    tid = tc.get("id")
    tname = _tc_field(tc, "name")
    targs = _tc_field(tc, "arguments")
    if tid:
        slot["id"] = str(tid)
    if tname:
        slot["name"] = str(tname)
    if targs is not None:
        slot["arguments"] += targs if isinstance(targs, str) else str(targs or "")


def harvest_message_tool_calls(accum: dict[int, dict[str, str]], choice: dict[str, Any]) -> None:
    """Merge complete ``choice.message.tool_calls`` when deltas were omitted.

    mlx-openai-server (among others) may stream content but only attach
    complete tool_calls on the final chunk's ``message``.
    """
    msg = choice.get("message")
    if not isinstance(msg, dict):
        return
    for tc in msg.get("tool_calls") or []:
        if not isinstance(tc, dict):
            continue
        tname = _tc_field(tc, "name")
        if not str(tname or "").strip():
            continue
        idx = resolve_tool_call_index(accum, tc)
        slot = accum.setdefault(idx, {"id": "", "name": "", "arguments": ""})
        tid = tc.get("id")
        targs = _tc_field(tc, "arguments")
        if tid:
            slot["id"] = str(tid)
        slot["name"] = str(tname).strip()
        if targs is not None:
            slot["arguments"] = targs if isinstance(targs, str) else str(targs or "")


def _assemble_tool_calls(accum: dict[int, dict[str, str]]) -> list[ToolCall]:
    calls = [
        ToolCall(
            id=slot.get("id") or "",
            name=slot["name"].strip(),
            arguments=slot.get("arguments") or "{}",
        )
        for _, slot in sorted(accum.items())
        if str(slot.get("name") or "").strip()
    ]
    normalize_call_ids(calls)
    return calls


# ─── Client ──────────────────────────────────────────────────────────────────
class ModelClient:
    """Streaming chat-completions client for OpenAI-compatible servers."""

    def __init__(
        self,
        base_url: str,
        api_key: str = "",
        *,
        connect_timeout: float = 30.0,
        read_timeout: float = 600.0,
        content_salvage: "ContentSalvage | None" = None,
    ) -> None:
        self.base_url = (base_url or "").rstrip("/")
        # Optional richer content salvage (peripheral-injected). When None the
        # built-in Hermes-only XML salvage applies. See core/_salvage.py.
        self._content_salvage = content_salvage
        headers = {"Content-Type": "application/json"}
        if api_key:
            headers["Authorization"] = f"Bearer {api_key}"
        self._client = httpx.Client(
            base_url=self.base_url,
            headers=headers,
            timeout=httpx.Timeout(
                connect=connect_timeout, read=read_timeout, write=30.0, pool=30.0
            ),
        )

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

    def __enter__(self) -> "ModelClient":
        return self

    def __exit__(self, *_exc: Any) -> None:
        self.close()

    # ── payloads ──
    @staticmethod
    def _payload(
        *,
        model: str,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None,
        stream: bool,
        temperature: float | None,
        max_tokens: int | None,
        extra_body: dict[str, Any] | None,
        tool_choice: str,
    ) -> dict[str, Any]:
        payload: dict[str, Any] = {"model": model, "messages": messages, "stream": stream}
        if tools:
            payload["tools"] = tools
            payload["tool_choice"] = tool_choice
        if temperature is not None:
            payload["temperature"] = temperature
        if max_tokens is not None:
            payload["max_tokens"] = max_tokens
        if extra_body:
            payload.update(extra_body)
        return payload

    # ── public API ──
    def list_models(self) -> list[str]:
        try:
            resp = self._client.get("/models")
            resp.raise_for_status()
            data = resp.json()
        except Exception as exc:
            raise ModelConnectionError(
                f"GET /models failed: {exc}", retryable=is_retryable_error(exc)
            ) from exc
        return [
            str(m.get("id"))
            for m in (data.get("data") or [])
            if isinstance(m, dict) and m.get("id")
        ]

    def complete(
        self,
        *,
        model: str,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None = None,
        temperature: float | None = None,
        max_tokens: int | None = None,
        extra_body: dict[str, Any] | None = None,
        tool_choice: str = "auto",
    ) -> ModelTurn:
        """One non-stream chat completion (also the recovery path)."""
        payload = self._payload(
            model=model, messages=messages, tools=tools, stream=False,
            temperature=temperature, max_tokens=max_tokens,
            extra_body=extra_body, tool_choice=tool_choice,
        )
        last_error: Exception | None = None
        for attempt in range(RETRY_ATTEMPTS):
            try:
                resp = self._client.post("/chat/completions", json=payload)
                if resp.status_code != 200:
                    raise ModelConnectionError(
                        f"chat/completions returned {resp.status_code}: {resp.text[:300]}",
                        status_code=resp.status_code,
                        retryable=resp.status_code == 429 or resp.status_code >= 500,
                    )
                return self._turn_from_completion(resp.json())
            except ModelConnectionError as exc:
                last_error = exc
                if not exc.retryable or attempt >= RETRY_ATTEMPTS - 1:
                    raise
            except httpx.HTTPError as exc:
                last_error = exc
                if not is_retryable_error(exc) or attempt >= RETRY_ATTEMPTS - 1:
                    raise ModelConnectionError(
                        f"chat/completions request failed: {exc}",
                        retryable=is_retryable_error(exc),
                    ) from exc
            time.sleep(min(RETRY_BASE_DELAY * (2 ** attempt), RETRY_MAX_DELAY))
        raise ModelConnectionError(f"chat/completions failed: {last_error}")  # pragma: no cover

    def stream_turn(
        self,
        *,
        model: str,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None = None,
        on_delta: Callable[[str], None] | None = None,
        temperature: float | None = None,
        max_tokens: int | None = None,
        extra_body: dict[str, Any] | None = None,
        stop_check: Callable[[], bool] | None = None,
        estimated_tokens: int = 0,
    ) -> ModelTurn:
        """Stream one model turn; recover via non-stream / salvage when needed."""
        payload = self._payload(
            model=model, messages=messages, tools=tools, stream=True,
            temperature=temperature, max_tokens=max_tokens,
            extra_body=extra_body, tool_choice="auto",
        )
        idle = scaled_idle_timeout(estimated_tokens)
        first = min(idle, FIRST_CHUNK_TIMEOUT)

        last_error: Exception | None = None
        for attempt in range(RETRY_ATTEMPTS):
            try:
                turn = self._stream_once(
                    payload,
                    on_delta=on_delta,
                    stop_check=stop_check,
                    idle_timeout=idle,
                    first_chunk_timeout=first,
                )
                return self._finalize_turn(
                    turn,
                    model=model, messages=messages, tools=tools,
                    temperature=temperature, max_tokens=max_tokens,
                    extra_body=extra_body,
                )
            except ModelStreamError as exc:
                # Mid-stream crash: one non-stream recovery, then surface honestly.
                recovered = self._try_non_stream_recovery(
                    model=model, messages=messages, tools=tools,
                    temperature=temperature, max_tokens=max_tokens,
                    extra_body=extra_body,
                )
                if recovered is not None:
                    return recovered
                raise
            except ModelConnectionError as exc:
                last_error = exc
                if not exc.retryable or attempt >= RETRY_ATTEMPTS - 1:
                    raise
            time.sleep(min(RETRY_BASE_DELAY * (2 ** attempt), RETRY_MAX_DELAY))
        raise ModelConnectionError(f"stream failed: {last_error}")  # pragma: no cover

    # ── internals ──
    def _stream_once(
        self,
        payload: dict[str, Any],
        *,
        on_delta: Callable[[str], None] | None,
        stop_check: Callable[[], bool] | None,
        idle_timeout: float,
        first_chunk_timeout: float,
    ) -> ModelTurn:
        content_parts: list[str] = []
        accum: dict[int, dict[str, str]] = {}
        finish_reason: str | None = None
        got_first_chunk = False
        saw_done = False
        try:
            with self._client.stream("POST", "/chat/completions", json=payload) as resp:
                if resp.status_code != 200:
                    resp.read()
                    raise ModelConnectionError(
                        f"chat/completions returned {resp.status_code}: {resp.text[:300]}",
                        status_code=resp.status_code,
                        retryable=resp.status_code == 429 or resp.status_code >= 500,
                    )
                lines = iter_with_timeout(
                    iter(resp.iter_lines()),
                    idle_timeout=idle_timeout,
                    first_chunk_timeout=first_chunk_timeout,
                )
                for line in lines:
                    if stop_check is not None and stop_check():
                        finish_reason = "stopped"
                        break
                    if not line or not line.startswith(SSE_DATA_PREFIX):
                        continue
                    data = line[len(SSE_DATA_PREFIX):].strip()
                    if data == "[DONE]":
                        saw_done = True
                        break
                    try:
                        chunk = json.loads(data)
                    except json.JSONDecodeError:
                        continue  # truncated keep-alive noise; not a model payload
                    got_first_chunk = True
                    choices = chunk.get("choices") or []
                    if not choices:
                        continue
                    choice = choices[0] if isinstance(choices[0], dict) else {}
                    delta = choice.get("delta") or {}
                    piece = delta.get("content")
                    if piece:
                        content_parts.append(piece)
                        if on_delta is not None:
                            on_delta(piece)
                    for tc in delta.get("tool_calls") or []:
                        if isinstance(tc, dict):
                            accumulate_tool_call_delta(accum, tc)
                    harvest_message_tool_calls(accum, choice)
                    if choice.get("finish_reason"):
                        finish_reason = str(choice["finish_reason"])
        except (ModelConnectionError, ModelStreamError):
            raise
        except SSEIdleTimeout as exc:
            raise StreamIdleTimeout(
                str(exc),
                partial_content="".join(content_parts),
                partial_tool_calls=_assemble_tool_calls(accum),
            ) from exc
        except httpx.HTTPError as exc:
            if not got_first_chunk and is_retryable_error(exc):
                raise ModelConnectionError(
                    f"stream connect failed: {exc}", retryable=True
                ) from exc
            raise ModelStreamError(
                f"stream crashed: {exc}",
                partial_content="".join(content_parts),
                partial_tool_calls=_assemble_tool_calls(accum),
            ) from exc
        except Exception as exc:
            raise ModelStreamError(
                f"stream crashed: {exc}",
                partial_content="".join(content_parts),
                partial_tool_calls=_assemble_tool_calls(accum),
            ) from exc

        # An OpenAI-compatible stream always terminates with a finish_reason
        # and/or [DONE]. Clean EOF without either means the connection was
        # dropped mid-turn (servers and proxies often close without an error).
        if finish_reason is None and not saw_done:
            raise ModelStreamError(
                "stream ended without finish_reason or [DONE] (truncated)",
                partial_content="".join(content_parts),
                partial_tool_calls=_assemble_tool_calls(accum),
            )
        return ModelTurn(
            content="".join(content_parts),
            tool_calls=_assemble_tool_calls(accum),
            finish_reason=finish_reason,
        )

    def _finalize_turn(
        self,
        turn: ModelTurn,
        *,
        model: str,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None,
        temperature: float | None,
        max_tokens: int | None,
        extra_body: dict[str, Any] | None,
    ) -> ModelTurn:
        """Apply salvage / non-stream recovery when the stream carried no native calls."""
        if turn.tool_calls or turn.finish_reason == "stopped":
            return turn
        salvaged, cleaned = resolve_content_salvage(turn.content, self._content_salvage)
        if salvaged:
            return ModelTurn(
                content=cleaned,
                tool_calls=salvaged,
                finish_reason=turn.finish_reason,
                recovered_via="content_salvage",
            )
        # Markup present but unparseable: the model clearly intended a tool
        # call the stream lost — one non-stream attempt before giving up.
        if _TOOL_CALL_MARKER in turn.content.lower() and tools:
            recovered = self._try_non_stream_recovery(
                model=model, messages=messages, tools=tools,
                temperature=temperature, max_tokens=max_tokens,
                extra_body=extra_body,
            )
            if recovered is not None:
                return recovered
        return turn

    def _try_non_stream_recovery(
        self,
        *,
        model: str,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None,
        temperature: float | None,
        max_tokens: int | None,
        extra_body: dict[str, Any] | None,
    ) -> ModelTurn | None:
        """One-shot non-stream completion (port of try_mlx_non_stream_native_tool_calls)."""
        try:
            turn = self.complete(
                model=model, messages=messages, tools=tools,
                temperature=temperature,
                max_tokens=max(256, min(int(max_tokens or 4096), 8192)),
                extra_body=extra_body,
            )
        except ModelError:
            _log.debug("non-stream recovery failed", exc_info=True)
            return None
        if not turn.content and not turn.tool_calls:
            return None
        turn.recovered_via = "non_stream"
        return turn

    def _turn_from_completion(self, body: dict[str, Any]) -> ModelTurn:
        choices = body.get("choices") or []
        msg = (choices[0].get("message") or {}) if choices and isinstance(choices[0], dict) else {}
        finish = choices[0].get("finish_reason") if choices and isinstance(choices[0], dict) else None
        accum: dict[int, dict[str, str]] = {}
        harvest_message_tool_calls(accum, {"message": msg})
        content = msg.get("content") or ""
        calls = _assemble_tool_calls(accum)
        recovered_via = None
        if not calls:
            salvaged, cleaned = resolve_content_salvage(content, self._content_salvage)
            if salvaged:
                calls, content = salvaged, cleaned
                recovered_via = "content_salvage"
        return ModelTurn(
            content=content,
            tool_calls=calls,
            finish_reason=str(finish) if finish else None,
            recovered_via=recovered_via,
        )
