"""The tool-calling loop — the agent seam P3's agentic tools and P4's research
subagents run on (docs/V3-PLAN.md §1.6). One provider, one tool set, one
executor callback; the loop owns the tool_use → tool_result echo protocol.

This module is also the MODEL BOUNDARY, and it defends it in layers (the
2026-07 incident: a 4.2M-char ``read_file`` result entered a transcript and
the provider 400'd with ``ContextWindowExceededError``). Server read defaults
stay UI-sized on purpose — it is the loop's job to keep any single tool
result, and the conversation as a whole, inside what a model call can carry:

- every tool_result is capped at :data:`TOOL_RESULT_MAX_CHARS`, truncated
  with a marker that teaches the model to re-request bounded pages;
- before each API call the conversation is sized against
  :data:`MODEL_CONTEXT_CHARS` — first by chars, and anchored by the
  provider's own ``input_tokens`` report from the previous turn (which also
  covers what a char count can't see: tool schemas, JSON-escape inflation);
- an over-budget conversation DEGRADES before it refuses: the oldest
  loop-owned tool results are tombstoned (``[cleared …]``) so the agent can
  finish — subagents re-read from the vault, so old results are cheap to
  shed. Only when nothing evictable remains does the loop raise a clear
  :class:`ContextBudgetError` instead of letting the provider 400 opaquely.

Three more production behaviors, all bounded:

- a final answer cut off at ``max_tokens`` is retried once with a doubled
  ceiling (the truncated attempt's tokens still count) — a truncated findings
  JSON is a silent data loss otherwise;
- an agent repeating the IDENTICAL tool-call set turn after turn is nudged at
  :data:`STUCK_NUDGE_AT` and refused at :data:`STUCK_RAISE_AT` — the
  turn budget alone would let it burn tokens to the wall;
- every exception that leaves the loop carries ``.transcript`` and ``.usage``
  (see :func:`_salvage`) so callers can meter and post-mortem a dead run.

And three control affordances for the callers that host loops (jobs, clients):
``parallel_tools`` runs a multi-call turn concurrently against a thread-safe
executor; ``should_stop`` lets a cancelled job stop paying between calls
(:class:`ToolLoopCancelled`); ``on_turn`` observes progress without being
able to break the loop.

Budgets are module constants, not config: no model-config plumb reaches this
loop today, and they bound the transport, not per-vault taste. (A future
``model.context_chars`` override would land here.)
"""
import json
from concurrent.futures import ThreadPoolExecutor

from ..tokens import CHARS_PER_TOKEN as _CHARS_PER_TOKEN
from .base import ModelProviderError


class ToolLoopLimitError(RuntimeError):
    """The model was still requesting tools at max_turns — or repeated the
    identical call set past :data:`STUCK_RAISE_AT`. A budget, not a bug.
    Carries ``.transcript`` / ``.usage`` (:func:`_salvage`)."""


class ToolLoopCancelled(RuntimeError):
    """``should_stop`` answered True — the caller (a cancelled job, a closing
    client) asked the loop to stop between calls. Deliberately NOT a
    ModelProviderError: cancellation is control flow, not failure. Carries
    ``.transcript`` / ``.usage`` so partial work stays meterable."""


class ContextBudgetError(ModelProviderError):
    """The conversation would exceed the model context budget and nothing
    evictable remains — raised BEFORE the API call, naming the size, instead
    of the provider's opaque 400. Subclasses ModelProviderError so existing
    degrade paths (research agent statuses, ingest's ToolError) handle it
    without new catch sites. Carries ``.transcript`` / ``.usage``."""


# Per-result cap at the model boundary. Server reads stay UI-sized (10MB);
# a model transcript never carries more than this per tool_result.
TOOL_RESULT_MAX_CHARS = 64_000

# ~chars a conversation may reach before the next call must shed old results
# (and, failing that, is refused). Well under the ~1M-token provider window
# so the error is ours, not a 400.
MODEL_CONTEXT_CHARS = 800_000

# Converts the provider's input_tokens report into the char scale above —
# the anchor that catches what a char count under-sees (tool schemas in the
# prompt, unicode/JSON-escape inflation). ~4 chars/token measured in the
# incident (4.2M chars ≈ 1.06M tokens); that measurement is now the shared
# constant in :mod:`alexandria_vault.tokens`, re-exported here so this module's
# readers still find the ratio where the guard uses it.
CHARS_PER_TOKEN = _CHARS_PER_TOKEN

# The stuck-loop guard: identical consecutive tool-call sets get a nudge, then
# a refusal — turn budgets alone let a wedged agent burn tokens to the wall.
STUCK_NUDGE_AT = 3
STUCK_RAISE_AT = 5

# The exact marker appended to a truncated result — it teaches the model to
# page deliberately (the read tools take start/end and max_chars).
TRUNCATION_MARKER = ("[truncated at {kept:,} of {total:,} chars — "
                     "re-request with start/end or max_chars]")

# What replaces an evicted (old, loop-owned) tool result. Same lesson: the
# content is not gone from the world, only from this transcript.
CLEARED_MARKER = ("[cleared {total:,} chars of an older tool result to fit "
                  "the context budget — re-request it if still needed]")

STUCK_MARKER = ("[loop guard: this identical tool-call set has now run {n} "
                "turns in a row — change approach or produce your final answer]")


def _bounded(text):
    """``text`` capped at TOOL_RESULT_MAX_CHARS, the truncation marker appended."""
    total = len(text)
    if total <= TOOL_RESULT_MAX_CHARS:
        return text
    kept = TOOL_RESULT_MAX_CHARS
    return text[:kept] + "\n" + TRUNCATION_MARKER.format(kept=kept, total=total)


def _content_chars(content):
    """~chars one message content contributes: strings count whole; block lists
    recurse (dict blocks by their string/list/dict values, provider SDK echo
    blocks by str()). An estimate — it only has to be the right order of
    magnitude to act before the provider does."""
    if isinstance(content, str):
        return len(content)
    if isinstance(content, (list, tuple)):
        return sum(_content_chars(b) for b in content)
    if isinstance(content, dict):
        return sum(_content_chars(v) for v in content.values()
                   if isinstance(v, (str, list, tuple, dict)))
    return len(str(content))


def _conversation_chars(msgs, system):
    return len(system or "") + sum(_content_chars(m.get("content")) for m in msgs)


def _evict_oldest(msgs, first_owned):
    """Tombstone every real tool_result in the OLDEST loop-owned results
    message (never the final message — the model needs the results it just
    asked for; never caller-supplied input messages — the loop only sheds
    what it created). Returns chars freed; 0 means nothing left to evict."""
    for m in msgs[first_owned:-1]:
        if m.get("role") != "user" or not isinstance(m.get("content"), list):
            continue
        freed = 0
        for b in m["content"]:
            if (isinstance(b, dict) and b.get("type") == "tool_result"
                    and isinstance(b.get("content"), str)
                    and not b["content"].startswith("[cleared ")):
                tomb = CLEARED_MARKER.format(total=len(b["content"]))
                if len(b["content"]) > len(tomb):
                    freed += len(b["content"]) - len(tomb)
                    b["content"] = tomb
        if freed:
            return freed
    return 0


def _salvage(err, msgs, total):
    """Attach the partial transcript + usage to an exception leaving the loop
    — a dead run must still be meterable and post-mortem-able."""
    err.transcript = msgs
    err.usage = dict(total)
    return err


def _add_usage(total, usage):
    for k, v in (usage or {}).items():
        if isinstance(v, int) and v:
            total[k] = total.get(k, 0) + v


def _execute_all(execute, tool_calls, parallel):
    """One turn's tool calls → tool_result blocks, order-preserving. With
    ``parallel``, independent calls run concurrently (the executor must be
    thread-safe — research's Frontier is lock-protected by design); a call
    that raises still becomes its own ``is_error`` block either way."""
    def one(tc):
        try:
            out = execute(tc["name"], tc["input"])
            return {"type": "tool_result", "tool_use_id": tc["id"],
                    "content": _bounded(out if isinstance(out, str)
                                        else json.dumps(out))}
        except Exception as e:  # the model handles tool failure; the loop survives it
            return {"type": "tool_result", "tool_use_id": tc["id"],
                    "content": _bounded(f"tool error: {e}"), "is_error": True}
    if parallel and len(tool_calls) > 1:
        with ThreadPoolExecutor(max_workers=min(8, len(tool_calls))) as pool:
            return list(pool.map(one, tool_calls))
    return [one(tc) for tc in tool_calls]


def _notify(on_turn, turn, reply, total):
    """The observer hears every turn (tool names, stop_reason, running usage)
    — and must never be able to kill the loop it is watching."""
    if on_turn is None:
        return
    try:
        on_turn({"turn": turn, "tool_calls": [tc["name"] for tc in reply.tool_calls],
                 "stop_reason": reply.stop_reason, "usage": dict(total)})
    except Exception:
        pass


def run_tool_loop(provider, messages, tools, execute, *, system=None, model=None,
                  max_tokens=4096, max_turns=8, parallel_tools=False,
                  should_stop=None, on_turn=None):
    """Drive ``provider.chat`` until the model stops calling tools.

    ``execute(name, input) -> str|json-able`` runs one tool; an exception
    becomes an ``is_error`` tool_result (the model sees the failure and can
    recover). Returns ``(final_reply, transcript, usage_total)`` — transcript is
    the full message list (input list is not mutated, though the loop's own
    appended results may be tombstoned by eviction), usage_total sums every
    turn (this is what the gateway meters, and what research jobs report).

    ``parallel_tools`` runs a multi-call turn concurrently (thread-safe
    executors only). ``should_stop()`` is polled before each API call — True
    raises :class:`ToolLoopCancelled` (with salvage) so a cancelled job stops
    paying immediately. ``on_turn(info)`` observes every turn and cannot
    break the loop.

    Model-boundary behavior (module docstring): per-result caps, evict-then-
    refuse context budgeting, one max_tokens growth retry, the stuck guard,
    and salvage attributes on every raised exception.
    """
    msgs = list(messages)
    total = {"input_tokens": 0, "output_tokens": 0}
    first_owned = len(msgs)     # eviction may touch only messages the loop appends
    last_input_tokens = 0       # the provider's own size report for the last call
    chars_since_call = 0        # ~chars appended after that report
    cur_max_tokens = max_tokens
    grew_max_tokens = False
    prev_sig, streak = None, 0
    # A loop is the textbook prompt-cache shape (fixed system+tools, growing
    # prefix) — providers that advertise the cap get breakpoints on every call.
    try:
        chat_extra = {"cache": True} if provider.caps().get("prompt_cache") else {}
    except Exception:
        chat_extra = {}

    for turn in range(1, max_turns + 1):
        if should_stop is not None and should_stop():
            raise _salvage(ToolLoopCancelled(
                f"cancelled by caller before turn {turn}"), msgs, total)
        # ---- pre-flight: size, then shed, then refuse -----------------------------
        size = _conversation_chars(msgs, system)
        # int(): the shared ratio is a float, and this number reaches the user
        # in the ContextBudgetError below, where ':,' would render it "800,000.0".
        anchored = int(CHARS_PER_TOKEN * last_input_tokens) + chars_since_call
        while max(size, anchored) > MODEL_CONTEXT_CHARS:
            freed = _evict_oldest(msgs, first_owned)
            if not freed:
                raise _salvage(ContextBudgetError(
                    f"conversation is ~{max(size, anchored):,} chars, over the "
                    f"{MODEL_CONTEXT_CHARS:,}-char model context budget and nothing "
                    "more is evictable — page tool results (start/end, max_chars) "
                    "or start a fresh conversation"), msgs, total)
            size -= freed
            anchored -= freed

        try:
            reply = provider.chat(msgs, system=system, model=model,
                                  max_tokens=cur_max_tokens, tools=tools,
                                  **chat_extra)
        except ModelProviderError as e:
            raise _salvage(e, msgs, total)
        _add_usage(total, reply.usage)
        last_input_tokens = reply.usage.get("input_tokens", 0) or 0
        chars_since_call = 0

        if not reply.tool_calls:
            if reply.stop_reason == "max_tokens" and not grew_max_tokens:
                # a truncated FINAL answer is silent data loss (a cut-off
                # findings JSON parses as nothing) — pay for one more attempt
                # with real room. The truncated attempt's tokens stay counted.
                grew_max_tokens = True
                cur_max_tokens = max_tokens * 2
                continue
            _notify(on_turn, turn, reply, total)
            return reply, msgs, total

        msgs.append({"role": "assistant", "content": reply.raw_content})
        chars_since_call += _content_chars(reply.raw_content)

        # ---- the stuck guard: identical call set, turn after turn -----------------
        sig = tuple(sorted(
            (tc["name"], json.dumps(tc["input"], sort_keys=True, default=str))
            for tc in reply.tool_calls))
        streak = streak + 1 if sig == prev_sig else 1
        prev_sig = sig
        if streak >= STUCK_RAISE_AT:
            raise _salvage(ToolLoopLimitError(
                f"stuck: the identical tool-call set ran {streak} consecutive "
                "turns — refusing to keep paying for it"), msgs, total)

        results = _execute_all(execute, reply.tool_calls, parallel_tools)
        if streak >= STUCK_NUDGE_AT:
            results.append({"type": "text", "text": STUCK_MARKER.format(n=streak)})
        msgs.append({"role": "user", "content": results})
        chars_since_call += _content_chars(results)
        _notify(on_turn, turn, reply, total)

    raise _salvage(ToolLoopLimitError(
        f"model still requesting tools after {max_turns} turns"), msgs, total)
