"""Kernel agent loop (docs/kernel_spec.md §2, §5, §7).

Call model -> execute tool calls -> append results -> repeat -> persist.
No feature logic lives here: peripherals plug in through
``RunConfig.context_providers`` (system-prompt blocks), the ``ToolRunner``
registry (tools), and ``RunEvents`` (output subscribers).

Error feedback replaces repair (spec §5): a malformed or failing tool call
comes back as a ``role: tool`` message the model acts on; the loop never
guesses arguments and never fabricates assistant content. An empty final
response is surfaced honestly. A confused model that spins on the same call,
or that exceeds the turn budget, is halted with a plain-language notice.

History is always native (``assistant.tool_calls`` + ``role: tool``); the
model client normalizes salvaged calls to ``call_*`` ids before they reach us.

Boundary: stdlib + core/ only.
"""
from __future__ import annotations

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

from core.model import ModelClient, ModelError, ModelTurn, ToolCall
from core.run import RunConfig, RunEvents
from core.store import ConversationStore
from core.tools import ToolRunner

_log = logging.getLogger(__name__)

_EMPTY_RESPONSE_NOTICE = "The model returned an empty response."
_LOOP_REPEAT_NOTICE = (
    "Stopping: the model repeated the same tool call without making progress."
)
_LOOP_BUDGET_NOTICE = (
    "Stopping: reached the maximum number of tool-calling turns for this run."
)
_SAME_TOOL_REPEAT_NOTICE = (
    "Stopping: the model called the same tool repeatedly without making progress."
)
_SAME_TOOL_SYNTHESIS_PROMPT = (
    "Stop calling tools. Answer the user now using the tool results already "
    "in this conversation. Do not call any tools."
)
# Minimum tool payload before we spend one no-tools synthesis turn after a halt.
_SYNTHESIS_MIN_TOOL_CONTENT_CHARS = 60


@dataclass
class RunResult:
    """Outcome of one run.

    ``messages`` is the list of *new* messages this run produced (assistant
    and tool turns) — the prior history stays in the store, untouched.
    """

    ok: bool
    final_text: str = ""
    turns: int = 0
    tool_calls: int = 0
    stopped: bool = False
    error: str | None = None
    messages: list[dict] = field(default_factory=list)


def run_agent(
    config: RunConfig,
    events: RunEvents,
    runner: ToolRunner,
    store: ConversationStore,
    client: ModelClient | None = None,
) -> RunResult:
    """Drive one agent run to a terminal content turn (or an honest stop).

    The caller persists the new user message to ``store`` *before* calling;
    the loop loads that history, contributes context-provider blocks as the
    system message, and persists only the messages it produces.
    """
    owns_client = client is None
    if client is None:
        client = ModelClient(
            config.base_url, config.api_key, content_salvage=config.content_salvage
        )
    try:
        return _drive(config, events, runner, store, client)
    finally:
        if owns_client:
            client.close()


def _invoke_trace(hook: Callable[..., None] | None, *args, **kwargs) -> None:
    if hook is None:
        return
    try:
        hook(*args, **kwargs)
    except Exception:
        pass


def _trace_turn_end(
    config: RunConfig, turn_number: int, final_content: str | None, reason: str
) -> None:
    hooks = config.trace_hooks
    if hooks is None:
        return
    _invoke_trace(hooks.on_turn_end, turn_number, final_content, reason)


_HISTORY_OMITTED = "[Earlier messages omitted]"


def _trim_history_if_needed(messages: list[dict], max_history_chars: int | None) -> None:
    """Drop oldest non-system messages when serialized history exceeds the cap."""
    if not max_history_chars or max_history_chars <= 0:
        return
    tombstone_present = False
    while len(messages) > 2:
        payload = json.dumps(messages, ensure_ascii=False)
        if len(payload) <= max_history_chars:
            return
        drop_idx = next(
            (
                i
                for i, m in enumerate(messages)
                if m.get("role") != "system"
                and (m.get("content") or "").strip() != _HISTORY_OMITTED
            ),
            None,
        )
        if drop_idx is None:
            return
        messages.pop(drop_idx)
        if not tombstone_present:
            insert_at = next(
                (i for i, m in enumerate(messages) if m.get("role") != "system"),
                len(messages),
            )
            messages.insert(insert_at, {"role": "user", "content": _HISTORY_OMITTED})
            tombstone_present = True


def _drive(
    config: RunConfig,
    events: RunEvents,
    runner: ToolRunner,
    store: ConversationStore,
    client: ModelClient,
) -> RunResult:
    if config.stop_check():
        return RunResult(ok=True, stopped=True)

    system_text = config.system_prompt()
    messages: list[dict] = []
    if system_text:
        messages.append({"role": "system", "content": system_text})
    messages.extend(store.load(config.thread_id))

    produced: list[dict] = []
    prev_signature: tuple | None = None
    repeat_count = 0
    same_tool_name: str | None = None
    same_tool_count = 0
    same_tool_nudged = False
    noop_promise_retries = 0
    total_tool_calls = 0

    for turn_number in range(1, config.max_turns + 1):
        if config.stop_check():
            _trace_turn_end(config, turn_number, None, "stopped")
            return _finish_stopped(store, config, produced)

        hooks = config.trace_hooks
        _invoke_trace(
            hooks.on_turn_start if hooks else None, turn_number, len(messages)
        )

        _trim_history_if_needed(messages, config.max_history_chars)

        # Recompute each turn: a tool (e.g. find_tools) may register new tools
        # into the runner mid-run; they must appear in the next turn's schemas.
        tools = runner.schemas() or None
        create_kwargs = {
            "model": config.model_id,
            "max_tokens": config.max_tokens,
            "temperature": config.temperature,
            "tools": tools,
        }
        _invoke_trace(
            hooks.on_llm_request if hooks else None,
            turn_number,
            messages,
            create_kwargs,
        )
        try:
            turn = client.stream_turn(
                model=config.model_id,
                messages=messages,
                tools=tools,
                on_delta=events.on_chunk,
                temperature=config.temperature,
                max_tokens=config.max_tokens,
                extra_body=config.extra_body,
                stop_check=config.stop_check,
            )
        except ModelError as exc:
            _trace_turn_end(config, turn_number, None, "model_error")
            return _finish_error(store, config, produced, events, str(exc), turn_number)

        tool_call_trace = [
            {"name": tc.name, "id": tc.id, "arguments": tc.arguments}
            for tc in turn.tool_calls
        ]
        _invoke_trace(
            hooks.on_llm_response if hooks else None,
            turn_number,
            turn.content or "",
            tool_call_trace,
        )

        if turn.finish_reason == "stopped":
            _trace_turn_end(config, turn_number, None, "stopped")
            return _finish_stopped(store, config, produced)

        if not turn.tool_calls:
            retry_prompt: str | None = None
            if config.noop_promise_retry is not None and (turn.content or "").strip():
                try:
                    retry_prompt = config.noop_promise_retry(turn.content, messages)
                except Exception:
                    retry_prompt = None
            if (
                retry_prompt
                and noop_promise_retries < config.max_noop_promise_retries
            ):
                noop_promise_retries += 1
                messages.append({"role": "user", "content": retry_prompt})
                continue
            return _finish_content(
                store, config, produced, events, turn, turn_number, total_tool_calls
            )

        assistant_msg = _assistant_tool_message(turn, config)
        messages.append(assistant_msg)
        produced.append(assistant_msg)

        signature = _call_signature(turn.tool_calls)
        if signature == prev_signature:
            repeat_count += 1
        else:
            repeat_count = 0
            prev_signature = signature

        for call in turn.tool_calls:
            total_tool_calls += 1
            events.on_tool_event(
                {"phase": "start", "name": call.name, "call_id": call.id}
            )
            t0 = time.monotonic()
            result = runner.run(call)
            duration = time.monotonic() - t0
            events.on_tool_event(
                {
                    "phase": "end",
                    "name": result.name,
                    "call_id": result.call_id,
                    "ok": result.ok,
                }
            )
            _invoke_trace(
                hooks.on_tool_call if hooks else None,
                turn_number,
                call.name,
                call.arguments,
                result.output,
                duration,
            )
            tool_msg = {
                "role": "tool",
                "tool_call_id": result.call_id,
                "content": result.output,
            }
            messages.append(tool_msg)
            produced.append(tool_msg)

        if repeat_count >= config.max_repeats:
            _trace_turn_end(config, turn_number, None, "loop_repeat")
            return _finish_budget(
                store, config, produced, events, _LOOP_REPEAT_NOTICE,
                turn_number, total_tool_calls,
            )

        if config.max_same_tool_turns >= 1:
            dominant = _dominant_tool_name(turn.tool_calls)
            if dominant is not None:
                if dominant == same_tool_name:
                    same_tool_count += 1
                else:
                    same_tool_name = dominant
                    same_tool_count = 1
                    same_tool_nudged = False
            else:
                same_tool_name = None
                same_tool_count = 0
                same_tool_nudged = False

            if same_tool_count >= config.max_same_tool_turns + 2:
                _trace_turn_end(config, turn_number, None, "same_tool_repeat")
                if _has_tool_content_for_synthesis(produced):
                    synth = _attempt_tool_result_synthesis(
                        config,
                        events,
                        store,
                        client,
                        messages,
                        produced,
                        turn_number,
                        total_tool_calls,
                    )
                    if synth is not None:
                        return synth
                return _finish_budget(
                    store,
                    config,
                    produced,
                    events,
                    _SAME_TOOL_REPEAT_NOTICE,
                    turn_number,
                    total_tool_calls,
                )
            if (
                same_tool_count == config.max_same_tool_turns
                and not same_tool_nudged
                and same_tool_name
            ):
                same_tool_nudged = True
                nudge = (
                    f"You've called `{same_tool_name}` {same_tool_count} turns in a row "
                    "— act on the results or finish."
                )
                messages.append({"role": "user", "content": nudge})
                continue

    _trace_turn_end(config, config.max_turns, None, "loop_budget")
    return _finish_budget(
        store, config, produced, events, _LOOP_BUDGET_NOTICE,
        config.max_turns, total_tool_calls,
    )


# ─── message shaping ──────────────────────────────────────────────────────────
def _assistant_tool_message(turn: ModelTurn, config: RunConfig) -> dict:
    """Native assistant message carrying the turn's tool calls."""
    content = turn.content or ""
    if config.assistant_tool_content_sanitizer is not None:
        try:
            content = config.assistant_tool_content_sanitizer(content)
        except Exception:
            content = turn.content or ""
    return {
        "role": "assistant",
        "content": content,
        "tool_calls": [
            {
                "id": tc.id,
                "type": "function",
                "function": {"name": tc.name, "arguments": tc.arguments},
            }
            for tc in turn.tool_calls
        ],
    }


def _call_signature(calls: list[ToolCall]) -> tuple:
    return tuple((c.name, c.arguments) for c in calls)


def _dominant_tool_name(calls: list[ToolCall]) -> str | None:
    """Return the tool name when every call in the turn shares it, else None."""
    if not calls:
        return None
    names = {(c.name or "").strip() for c in calls}
    names.discard("")
    if len(names) != 1:
        return None
    return next(iter(names))


def _has_tool_content_for_synthesis(produced: list[dict]) -> bool:
    """True when at least one tool result looks worth synthesizing from."""
    for msg in produced:
        if msg.get("role") != "tool":
            continue
        content = (msg.get("content") or "").strip()
        if not content:
            continue
        if len(content) >= _SYNTHESIS_MIN_TOOL_CONTENT_CHARS:
            return True
        if "Title:" in content or "\nURL:" in content:
            return True
    return False


def _attempt_tool_result_synthesis(
    config: RunConfig,
    events: RunEvents,
    store: ConversationStore,
    client: ModelClient,
    messages: list[dict],
    produced: list[dict],
    turn_number: int,
    total_tool_calls: int,
) -> RunResult | None:
    """One final no-tools model turn after same-tool halt; None if it fails."""
    if config.stop_check():
        return None

    synth_turn = turn_number + 1
    messages.append({"role": "user", "content": _SAME_TOOL_SYNTHESIS_PROMPT})

    hooks = config.trace_hooks
    _invoke_trace(
        hooks.on_turn_start if hooks else None, synth_turn, len(messages)
    )
    _trim_history_if_needed(messages, config.max_history_chars)

    create_kwargs = {
        "model": config.model_id,
        "max_tokens": config.max_tokens,
        "temperature": config.temperature,
        "tools": None,
    }
    _invoke_trace(
        hooks.on_llm_request if hooks else None,
        synth_turn,
        messages,
        create_kwargs,
    )
    try:
        turn = client.stream_turn(
            model=config.model_id,
            messages=messages,
            tools=None,
            on_delta=events.on_chunk,
            temperature=config.temperature,
            max_tokens=config.max_tokens,
            extra_body=config.extra_body,
            stop_check=config.stop_check,
        )
    except ModelError:
        return None

    _invoke_trace(
        hooks.on_llm_response if hooks else None,
        synth_turn,
        turn.content or "",
        [],
    )

    if turn.finish_reason == "stopped" or turn.tool_calls:
        return None

    content = (turn.content or "").strip()
    if not content:
        return None

    if config.text_sanitizer is not None:
        try:
            content = config.text_sanitizer(content)
        except Exception:
            content = (turn.content or "").strip()
    if not content.strip():
        return None

    assistant_msg = {"role": "assistant", "content": content}
    produced.append(assistant_msg)
    _persist(store, config, produced)
    _trace_turn_end(config, synth_turn, content, "same_tool_repeat_synthesis")
    events.on_finished(content)
    return RunResult(
        ok=True,
        final_text=content,
        turns=synth_turn,
        tool_calls=total_tool_calls,
        messages=produced,
    )


# ─── terminal outcomes ──────────────────────────────────────────────────────────
def _persist(store: ConversationStore, config: RunConfig, produced: list[dict]) -> None:
    if produced:
        store.append(config.thread_id, produced)


def _finish_content(
    store: ConversationStore,
    config: RunConfig,
    produced: list[dict],
    events: RunEvents,
    turn: ModelTurn,
    turns: int,
    tool_calls: int,
) -> RunResult:
    content = turn.content or ""
    # Sanitize the final reply (reasoning-channel leakage, tool markup) before
    # it is persisted or shown. Injected by the wiring layer; fail-open so a
    # broken sanitizer never swallows a real reply.
    if config.text_sanitizer is not None:
        try:
            content = config.text_sanitizer(content)
        except Exception:
            content = turn.content or ""
    if not content.strip():
        # Honest empty response (spec §5): no synthetic content, no retry here
        # (the model client already did its one non-stream recovery).
        _trace_turn_end(config, turns, None, "empty_response")
        events.on_error(_EMPTY_RESPONSE_NOTICE)
        _persist(store, config, produced)
        return RunResult(
            ok=False, final_text="", turns=turns, tool_calls=tool_calls,
            error=_EMPTY_RESPONSE_NOTICE, messages=produced,
        )
    assistant_msg = {"role": "assistant", "content": content}
    produced.append(assistant_msg)
    _persist(store, config, produced)
    _trace_turn_end(config, turns, content, "final_reply")
    events.on_finished(content)
    return RunResult(
        ok=True, final_text=content, turns=turns, tool_calls=tool_calls,
        messages=produced,
    )


def _finish_error(
    store: ConversationStore,
    config: RunConfig,
    produced: list[dict],
    events: RunEvents,
    message: str,
    turns: int,
) -> RunResult:
    events.on_error(message)
    _persist(store, config, produced)
    return RunResult(ok=False, error=message, turns=turns, messages=produced)


def _finish_stopped(
    store: ConversationStore, config: RunConfig, produced: list[dict]
) -> RunResult:
    _persist(store, config, produced)
    return RunResult(ok=True, stopped=True, messages=produced)


def _finish_budget(
    store: ConversationStore,
    config: RunConfig,
    produced: list[dict],
    events: RunEvents,
    notice: str,
    turns: int,
    tool_calls: int,
) -> RunResult:
    events.on_error(notice)
    _persist(store, config, produced)
    return RunResult(
        ok=False, final_text="", turns=turns, tool_calls=tool_calls,
        error=notice, messages=produced,
    )
