"""Context trimming, checkpoints, OpenAI client cache, search dedupe, get_state.

**Runtime vs persisted transcript:** In-place trims (``trim_old_large_tool_results``,
``_emergency_trim_messages``, read chunk stubs) mutate the ``messages`` list used
for the **next** LLM request inside ``execute_raw_agent``. Persisted chat rows in
SQLite are written by ``agent_worker`` / ``chat_storage`` from the pre-trim or
post-turn snapshot depending on the append path — when debugging "UI shows full
history but the model erred", compare ``~/.mugi/logs/traces`` (when enabled)
against the DB row for that turn.
"""
from __future__ import annotations

import json
import logging
import os
import re
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from typing import Any, Callable

import httpx
from openai import OpenAI

from agent_loop_config import CHATS_DB_DIR, DEFAULT_AGENT_MAX_TOOL_ROUNDS
from agent_loop_guardrails import _multi_step_user_message
from agent_loop_tool_parse import _parse_xml_tool_calls
from agent_loop_tool_repair import _json_loads_lenient
from persistence import MUGI_DIR, get_current_agent_thread_id
from utils.transcript_roles import (
    is_real_user_message,
    is_summary_carrier_message,
    is_tool_result_carrier_message,
    tool_result_text_for_trim,
)

_log = logging.getLogger(__name__)

def _normalize_search_query(q: str) -> str:
    """Normalize a search query for similarity comparison (lowercase, single spaces)."""
    if not q or not isinstance(q, str):
        return ""
    return re.sub(r"\s+", " ", q.strip().lower())


_SEARCH_STOPWORDS = frozenset(
    "a an the in on at to for of is it and or by with from this that"
    .split()
)


def _search_queries_similar(q1: str, q2: str) -> bool:
    """True if two search queries are effectively the same.

    Containment (n1 in n2) only fires when the extra words are all
    stopwords — adding real content words like a date or location qualifier
    makes the queries different enough to warrant a fresh search.

    Token overlap uses unidirectional >=70% for very short queries (<=3 words)
    and symmetric Jaccard >=0.60 for longer queries so that deliberate
    simplifications (dropping a date, changing a qualifier) are allowed through.
    """
    n1 = _normalize_search_query(q1)
    n2 = _normalize_search_query(q2)
    if not n1 or not n2:
        return False
    if n1 == n2:
        return True
    if n1 in n2 or n2 in n1:
        shorter_s, longer_s = (n1, n2) if len(n1) <= len(n2) else (n2, n1)
        extra = set(longer_s.split()) - set(shorter_s.split())
        if extra <= _SEARCH_STOPWORDS:
            return True
    w1 = set(n1.split())
    w2 = set(n2.split())
    if not w1 or not w2:
        return False
    shorter, longer = (w1, w2) if len(w1) <= len(w2) else (w2, w1)
    if len(shorter) <= 3:
        return len(shorter & longer) / len(shorter) >= 0.70
    jaccard = len(w1 & w2) / len(w1 | w2)
    return jaccard >= 0.60


def _is_nonactionable_search_result(text: str) -> bool:
    """True when a prior search result is itself a guard/warning payload or empty result, not a real result worth reusing."""
    if not text:
        return True
    t = text.strip()
    tl = t.lower()
    if t.startswith("[LOOP WARNING]"):
        return True
    if "you already searched for a very similar query" in tl:
        return True
    if "do not call search_web again for this question" in tl:
        return True
    if "do not call search_web or deep_web_search again" in tl:
        return True
    if "do not call deep_web_search again" in tl:
        return True
    if tl.startswith("no results found"):
        return True
    if "try a shorter or different query" in tl and len(tl) < 300:
        return True
    return False


def _get_recent_similar_search_result(messages: list, current_query: str) -> str | None:
    """If a very similar search query was already run in this conversation, return that tool result content; else None."""
    if not current_query or not messages:
        return None

    # ── Collect tool results from both OpenAI structured format AND XML format ──
    tool_result_by_id: dict[str, str] = {}
    # Track XML-mode results keyed by position so we can pair them with the
    # immediately preceding assistant message that contains the XML tool calls.
    _xml_result_after: dict[int, str] = {}

    for i, m in enumerate(messages):
        if m.get("role") == "tool":
            tid = m.get("tool_call_id")
            if tid:
                tool_result_by_id[tid] = (m.get("content") or "") if isinstance(m.get("content"), str) else str(m.get("content", ""))
        elif m.get("role") == "user" and isinstance(m.get("content"), str) and "<tool_response>" in (m.get("content") or ""):
            _xml_result_after[i] = m["content"]

    for i in range(len(messages) - 1, -1, -1):
        m = messages[i]
        if m.get("role") != "assistant":
            continue
        content = m.get("content") or ""

        # ── Path 1: OpenAI structured tool_calls ──
        for tc in m.get("tool_calls") or []:
            tname = (tc.get("function") or {}).get("name") or tc.get("name") or ""
            if tname not in ("search_web", "search_web_public", "deep_web_search"):
                continue
            try:
                targs = (tc.get("function") or {}).get("arguments") or tc.get("arguments") or "{}"
                if isinstance(targs, str):
                    past = json.loads(targs)
                else:
                    past = targs
                past_query = (past.get("query") or "").strip()
            except Exception:
                past_query = ""
            if not past_query:
                continue
            if _search_queries_similar(current_query, past_query):
                tid = tc.get("id")
                if tid and tid in tool_result_by_id:
                    candidate = tool_result_by_id[tid]
                    if not _is_nonactionable_search_result(candidate):
                        return candidate

        # ── Path 2: XML tool calls embedded in assistant content ──
        if "<tool_call>" in content.lower():
            xml_tcs = _parse_xml_tool_calls(content)
            search_tc = None
            for xtc in xml_tcs:
                if xtc.get("name") in ("search_web", "search_web_public", "deep_web_search"):
                    try:
                        xargs = _json_loads_lenient(xtc.get("arguments", "{}"))
                        past_query = (xargs.get("query") or "").strip()
                    except Exception:
                        past_query = ""
                    if past_query and _search_queries_similar(current_query, past_query):
                        search_tc = xtc
                        break
            if search_tc is not None:
                xml_result = _xml_result_after.get(i + 1)
                if xml_result and not _is_nonactionable_search_result(xml_result):
                    return xml_result

    return None
# Chats database filenames (CHATS_DB_DIR and APP_PREFS_FILENAME imported from agent_loop_config)
CHATS_DB_FILENAME = "chats_db.json"
CHECKPOINT_DB_FILENAME = "checkpoints.sqlite"
os.makedirs(CHATS_DB_DIR, exist_ok=True)

# ── Multi-round checkpointing ────────────────────────────────────────────────
_CHECKPOINT_DIR = os.path.join(MUGI_DIR, "checkpoints")
_CHECKPOINT_MAX_AGE = 600  # seconds — ignore checkpoints older than 10 minutes


def _save_round_checkpoint(thread_id: str, messages: list, round_: int):
    """Persist lightweight WIP checkpoint so a crashed run can be resumed."""
    try:
        os.makedirs(_CHECKPOINT_DIR, exist_ok=True)
        path = os.path.join(_CHECKPOINT_DIR, f"{thread_id}_wip.json")
        payload = {
            "ts": time.time(),
            "round": round_,
            "messages": messages,
        }
        tmp = path + ".tmp"
        with open(tmp, "w", encoding="utf-8") as f:
            json.dump(payload, f)
        os.replace(tmp, path)
    except Exception:
        _log.debug("Failed to save round checkpoint", exc_info=True)


def _load_round_checkpoint(thread_id: str):
    """Return (messages, round) from a valid WIP checkpoint, or None."""
    path = os.path.join(_CHECKPOINT_DIR, f"{thread_id}_wip.json")
    try:
        if not os.path.isfile(path):
            return None
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
        age = time.time() - data.get("ts", 0)
        if age > _CHECKPOINT_MAX_AGE:
            _clear_round_checkpoint(thread_id)
            return None
        msgs = data.get("messages")
        rnd = data.get("round", 0)
        if not msgs or not isinstance(msgs, list):
            return None
        return msgs, rnd
    except Exception:
        return None


def _clear_round_checkpoint(thread_id: str):
    """Remove WIP checkpoint after successful completion."""
    try:
        path = os.path.join(_CHECKPOINT_DIR, f"{thread_id}_wip.json")
        if os.path.isfile(path):
            os.remove(path)
    except Exception:
        _log.debug("Failed to clear round checkpoint", exc_info=True)

# ---------------------------------------------------------------------------
# Raw OpenAI-compatible agent: tool schema conversion and thread message store
# ---------------------------------------------------------------------------

# Cache OpenAI client by (base_url, api_key) to avoid repeated connection setup per turn.
_openai_client_cache: dict[tuple[str, str], OpenAI] = {}

# Shared thread pool for parallel tool execution (avoids creating/destroying executor per batch).
_tool_executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="mugi_tool")


def _get_openai_client(base_url: str, api_key: str) -> OpenAI:
    """Return a cached or new OpenAI client for the given base_url and api_key."""
    key = (base_url or "", api_key or "")
    if key not in _openai_client_cache:
        _openai_client_cache[key] = OpenAI(
            base_url=base_url,
            api_key=api_key,
            timeout=httpx.Timeout(connect=30.0, read=600.0, write=30.0, pool=30.0),
            max_retries=2,
        )
    return _openai_client_cache[key]
def get_state(config):
    """Return state compatible with app's loadConversationForThread_ / conversationMessageCountForThread_.
    state.values["messages"] is a list of message-like objects with .type ('human'|'ai') and .content."""
    thread_id = (config or {}).get("configurable") or {}
    if isinstance(thread_id, dict):
        thread_id = thread_id.get("thread_id") or "default"
    else:
        thread_id = "default"
    from chat_storage import get_messages as _get_messages
    raw = _get_messages(thread_id)
    messages = []
    for m in raw:
        if not isinstance(m, dict):
            continue
        role = m.get("role")
        content = m.get("content")
        if isinstance(content, list):
            content = " ".join(
                ((b.get("text", str(b)) if isinstance(b, dict) else str(b))
                 for b in content)
            )
        if role == "user":
            messages.append(_Msg(type="human", content=content or ""))
        elif role == "assistant":
            messages.append(_Msg(type="ai", content=content or ""))
    return _State(values={"messages": messages})


class _Msg:
    __slots__ = ("type", "content")
    def __init__(self, type, content):
        self.type = type
        self.content = content


class _State:
    __slots__ = ("values",)
    def __init__(self, values):
        self.values = values

MAX_TOOL_ROUNDS = DEFAULT_AGENT_MAX_TOOL_ROUNDS
REPETITION_WINDOW = 5

# When an OCR (or other) tool returns more than this many chars, we save to file and give the model
# a short summary so context doesn't blow up and the model doesn't "forget" the task (e.g. reply with identity blurb).
MAX_TOOL_RESULT_CHARS = 35_000
# read_file_chunk: cap chars per tool result so multi-chunk reads don't blow up the next request (avoids 500)
MAX_READ_FILE_CHUNK_RESULT_CHARS = 10_000


def _maybe_save_large_tool_result(
    tool_name: str,
    result: str,
    thread_id: str,
    last_output_path_ref: list | None,
    workspace_dir: str | None = None,
    research_result_max_chars: int | None = None,
) -> str:
    """If result is huge (e.g. OCR output), save to ~/Documents/Mugi/OCR and return short message. Uses user-visible path so the user can find the file."""
    if not result:
        return result
    # read_file_chunk: keep batch results under cap so we don't send 5 chunks (17k+ chars) in one message
    if tool_name == "read_file_chunk" and len(result) > MAX_READ_FILE_CHUNK_RESULT_CHARS:
        return (
            result[:MAX_READ_FILE_CHUNK_RESULT_CHARS]
            + f"\n\n[Truncated at {MAX_READ_FILE_CHUNK_RESULT_CHARS:,} chars to avoid overwhelming the server. "
            f"Use read_file_chunk with num_chunks=1 and chunk_index=N to read and edit one chunk at a time.]"
        )
    if len(result) <= MAX_TOOL_RESULT_CHARS:
        return result
    # Research tools and web tools: cap size so many rounds don't blow up the next request (avoids 500)
    max_chars = research_result_max_chars if research_result_max_chars is not None else 15_000
    if tool_name in ("deep_web_search", "read_and_analyze_page", "analyze_multiple_sources"):
        if len(result) > max_chars:
            return (
                result[:max_chars]
                + f"\n\n[Truncated at {max_chars:,} chars to avoid request size limits. "
                f"Summarize key points and use save_research_findings or generate_final_research_report to save the full research.]"
            )
    if tool_name in ("read_webpage", "search_web"):
        if len(result) > max_chars:
            return (
                result[:max_chars]
                + f"\n\n[Truncated at {max_chars:,} chars to fit context. Summarize from the content above or ask the user to open the URL.]"
            )
    # OCR and other long-running tools that return full document text (save to ~/Documents/Mugi/OCR)
    if tool_name not in ("ocr_scanned_pdf", "ocr_complex_document", "ocr_with_kraken_api", "extract_pdf_text"):
        return result
    ext = ".md" if tool_name.startswith("ocr") or "markdown" in result[:200].lower() else ".txt"
    from utils.path_safety import get_ocr_output_dir
    out_dir = get_ocr_output_dir()
    safe_id = (thread_id or "default").replace(os.sep, "_")[:32]
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    out_path = os.path.join(out_dir, f"ocr_result_{safe_id}_{ts}{ext}")
    out_path = os.path.abspath(out_path)
    try:
        with open(out_path, "w", encoding="utf-8") as f:
            f.write(result)
            f.flush()
            try:
                os.fsync(f.fileno())
            except (OSError, AttributeError):
                pass
        if not os.path.isfile(out_path):
            return result[:MAX_TOOL_RESULT_CHARS] + "\n\n[Truncated; save failed: file was not created after write.]"
        size = os.path.getsize(out_path)
        if size < len(result) * 0.9:
            return result[:MAX_TOOL_RESULT_CHARS] + f"\n\n[Truncated; save may be incomplete: wrote {size} bytes, expected ~{len(result)}.]"
        if last_output_path_ref is not None:
            last_output_path_ref[0] = out_path
        return (
            f"The {tool_name} result was too long to include here ({len(result):,} characters). "
            f"It has been saved to:\n{out_path}\n\nTell the user their document is at that path. Do not introduce yourself or repeat your identity—just confirm the file location."
        )
    except OSError as e:
        return result[:MAX_TOOL_RESULT_CHARS] + f"\n\n[Truncated; save failed: {e}]"


_EMERGENCY_KEEP_TAIL = 6  # keep last N messages on emergency trim
_WORKING_BUFFER_INJECT_MAX_CHARS = 3500
# Stub: working buffer with no real entries (only header).
_WORKING_BUFFER_STUB_MAX_LEN = 200


def _is_real_user_query(m: dict) -> bool:
    """True if the message is a genuine user query, not a tool response wrapped in user role."""
    return is_real_user_message(m)


def _adjust_keep_from_openai_tool_pairs(messages: list, keep_from: int) -> int:
    """Move keep_from left (earlier index) so native tool messages are not orphaned."""
    n = len(messages)
    kf = max(1, min(keep_from, n))
    for _ in range(n + 8):
        prev = kf
        for i in range(kf, n):
            if (messages[i].get("role") or "").strip() != "tool":
                continue
            tid = messages[i].get("tool_call_id")
            if not tid:
                continue
            asst_idx = None
            for j in range(i - 1, 0, -1):
                if (messages[j].get("role") or "").strip() != "assistant":
                    continue
                for tc in messages[j].get("tool_calls") or []:
                    if isinstance(tc, dict) and tc.get("id") == tid:
                        asst_idx = j
                        break
                if asst_idx is not None:
                    break
            if asst_idx is not None and asst_idx < kf:
                kf = asst_idx
        j = 1
        while j < n:
            if (messages[j].get("role") or "").strip() != "assistant":
                j += 1
                continue
            tcs = messages[j].get("tool_calls") or []
            if not tcs:
                j += 1
                continue
            k = j + 1
            while k < n and (messages[k].get("role") or "").strip() == "tool":
                k += 1
            block_end = k
            if j < kf < block_end:
                kf = j
            j = block_end
        if kf == prev:
            break
    return kf


def _adjust_keep_from_xml_tool_pairs(messages: list, keep_from: int) -> int:
    """Ensure an assistant XML <tool_call> block is not separated from its <tool_response> user turn."""
    n = len(messages)
    kf = max(1, min(keep_from, n))
    for i in range(kf, n):
        m = messages[i]
        if (m.get("role") or "").strip() != "user":
            continue
        if not is_tool_result_carrier_message(m):
            continue
        for j in range(i - 1, 0, -1):
            prev = messages[j]
            if (prev.get("role") or "").strip() != "assistant":
                continue
            c = prev.get("content")
            if isinstance(c, str) and "<tool_call" in c.lower():
                if j < kf:
                    kf = j
                break
    return kf


def _emergency_keep_from_pair_safe(messages: list, *, tail_target: int) -> int:
    """First conversation index to retain (1-based) with tool/XML pair invariants."""
    n = len(messages)
    if n <= 1:
        return 1
    conv_len = n - 1
    if conv_len <= tail_target:
        return 1
    keep_from = max(1, n - tail_target)
    for _ in range(n + 6):
        k0 = keep_from
        keep_from = _adjust_keep_from_openai_tool_pairs(messages, keep_from)
        keep_from = _adjust_keep_from_xml_tool_pairs(messages, keep_from)
        if keep_from == k0:
            break
    return keep_from


def _emergency_trim_messages(messages: list, on_log=None) -> bool:
    """Aggressively trim messages after a 500 (context overflow).

    Keeps ``messages[0]`` (system) and a suffix of the conversation. The suffix
    size is at least ``_EMERGENCY_KEEP_TAIL`` messages when possible, but expands
    leftward so native ``assistant``/``tool`` pairs and XML ``tool_call``/
    ``tool_response`` pairs are never split.

    Ensures at least one real user message remains when possible (APIs require a user turn).
    Returns True if anything was trimmed.
    """
    n = len(messages)
    if n <= _EMERGENCY_KEEP_TAIL + 1:
        return False
    keep_from = _emergency_keep_from_pair_safe(messages, tail_target=_EMERGENCY_KEEP_TAIL)
    if keep_from <= 1:
        return False
    remove_count = keep_from - 1
    removed = messages[1:keep_from]
    del messages[1:keep_from]
    insert_at = 1
    # Preserve the latest compressed-context summary so the model keeps
    # long-range task context across the trim (trace 6be8ae2b: dropping it
    # left the model with only guard nudges and it redid completed work).
    if not any(is_summary_carrier_message(m) for m in messages[1:]):
        last_summary = None
        for m in reversed(removed):
            if is_summary_carrier_message(m):
                last_summary = m
                break
        if last_summary is not None:
            messages.insert(insert_at, last_summary)
            insert_at += 1
            if on_log:
                on_log("AGENT", "Emergency trim: preserved latest conversation summary across trim.")
    kept_tail = messages[1:]
    if not any(_is_real_user_query(m) for m in kept_tail):
        last_user = None
        for m in reversed(removed):
            if _is_real_user_query(m):
                last_user = m
                break
        if last_user is not None:
            messages.insert(insert_at, last_user)
            if on_log:
                on_log("AGENT", "Emergency trim: re-inserted most recent user message so API receives a user query.")
        else:
            messages.insert(insert_at, {"role": "user", "content": "(Continuing task.)"})
            if on_log:
                on_log("AGENT", "Emergency trim: no real user query survived trimming — inserted synthetic user turn so API does not reject.")
    for m in messages[1:]:
        content = tool_result_text_for_trim(m)
        if isinstance(content, str) and len(content) > 2000:
            clipped = (
                content[:2000]
                + "\n\n[Content truncated to 2000 characters to fit context window. You may ask the user to re-run the tool or summarize what you need.]"
            )
            if (m.get("role") or "").strip() == "tool":
                m["content"] = clipped
            elif is_tool_result_carrier_message(m):
                m["content"] = f"<tool_response>\n{clipped}\n</tool_response>"
    if on_log:
        on_log(
            "WARN",
            f"Emergency trim: removed {remove_count} old messages (pair-safe keep_from={keep_from}, kept system + last {len(messages) - 1} messages).",
        )
    return True


def _read_working_buffer_for_inject() -> str | None:
    """Read working buffer file for the current thread (via get_current_agent_thread_id); return truncated content to inject, or None if missing/stub."""
    try:
        from tools.memory_tools import _get_working_buffer_path
        path, err = _get_working_buffer_path()
        if err or not path or not os.path.isfile(path):
            return None
        with open(path, "r", encoding="utf-8") as f:
            content = f.read()
    except OSError:
        return None
    if not content or len(content.strip()) < _WORKING_BUFFER_STUB_MAX_LEN:
        return None
    # Consider it a stub if it's only the header (no section entries).
    stripped = content.strip()
    if "## " not in stripped and "---" not in stripped:
        return None
    if len(content) > _WORKING_BUFFER_INJECT_MAX_CHARS:
        # Tail-keep: for a run log the *newest* entries are what a post-trim
        # model needs; keep the header line plus the most recent tail, cut at
        # an entry boundary so no digest line is split mid-way.
        header, _, _rest = content.partition("\n")
        budget = max(400, _WORKING_BUFFER_INJECT_MAX_CHARS - len(header) - 64)
        tail = content[-budget:]
        nl = tail.find("\n")
        if nl != -1:
            tail = tail[nl + 1:]
        content = header + "\n\n[Older buffer entries truncated.]\n\n" + tail
    return content


# --- Auto-checkpoint run log (Point 4: durable working state under pressure) ---
# The model's accumulated discoveries must survive what the context window
# cannot. The harness collects per-round digests and persists them under this
# section in the per-thread working buffer; compression sites inject them back
# so a post-trim model continues from its own knowledge instead of redoing it.
# Motivating trace: 6be8ae2b (11 rounds of debugging discoveries evaporated in
# one trim and the run died repeating itself).

_AUTO_RUN_LOG_HEADER = "## Auto run log"
AUTO_CHECKPOINT_MAX_ENTRIES = 40
_WB_INJECT_START = "<!-- mugi:working-buffer:start -->"
_WB_INJECT_END = "<!-- mugi:working-buffer:end -->"


def format_auto_checkpoint_entry(round_: int, assistant_preview: str, tool_digests: list[str]) -> str:
    """One digest line per tool round: model's own diagnosis + tool outcomes."""
    preview = " ".join((assistant_preview or "").split())[:140]
    digests = "; ".join(" ".join(d.split())[:220] for d in tool_digests if d)
    body = f'R{int(round_)} "{preview}"' if preview else f"R{int(round_)}"
    if digests:
        body += f" | {digests}"
    return "- " + body


def append_auto_checkpoint_entries(entries: list[str]) -> int:
    """Append digest lines under the Auto run log section of the thread working buffer.

    Rotates to the most recent ``AUTO_CHECKPOINT_MAX_ENTRIES``; skips lines
    already present (idempotent across round-flush + compression-flush).
    Agent-authored sections are preserved verbatim. Never raises.
    """
    entries = [e for e in entries if isinstance(e, str) and e.strip()]
    if not entries:
        return 0
    try:
        from tools.memory_tools import _get_working_buffer_path
        path, err = _get_working_buffer_path()
        if err or not path:
            return 0
        existing = ""
        if os.path.isfile(path):
            with open(path, "r", encoding="utf-8") as f:
                existing = f.read()
        if not existing.strip():
            existing = "# Working Buffer (Danger Zone Log)\n**Status:** ACTIVE\n"
        head_lines: list[str] = []
        auto_lines: list[str] = []
        in_auto = False
        for ln in existing.splitlines():
            if ln.strip() == _AUTO_RUN_LOG_HEADER:
                in_auto = True
                continue
            if in_auto:
                if ln.startswith("- "):
                    auto_lines.append(ln)
                    continue
                if not ln.strip():
                    continue
                in_auto = False
            head_lines.append(ln)
        seen = set(auto_lines)
        new_lines = [e for e in entries if e not in seen]
        if not new_lines:
            return 0
        auto_lines.extend(new_lines)
        auto_lines = auto_lines[-AUTO_CHECKPOINT_MAX_ENTRIES:]
        content = (
            "\n".join(head_lines).rstrip()
            + "\n\n"
            + _AUTO_RUN_LOG_HEADER
            + "\n"
            + "\n".join(auto_lines)
            + "\n"
        )
        parent = os.path.dirname(path)
        if parent:
            os.makedirs(parent, exist_ok=True)
        with open(path, "w", encoding="utf-8") as f:
            f.write(content)
        return len(new_lines)
    except Exception:
        _log.debug("auto checkpoint append failed", exc_info=True)
        return 0


def inject_working_buffer_into_system(messages: list) -> bool:
    """Inject (or refresh) the working-buffer block in the system message.

    Sentinel-wrapped so repeated compressions replace the block instead of
    stacking copies; bounded by ``_WORKING_BUFFER_INJECT_MAX_CHARS``.
    Returns True if a block is present in the system message afterwards.
    """
    try:
        buf_content = _read_working_buffer_for_inject()
        if not buf_content:
            return False
        if not (messages and (messages[0].get("role") or "").strip() == "system"):
            return False
        danger_hint = (
            "\n\nCONTEXT: You are in the danger zone — old messages were trimmed to fit. "
            "The block above summarizes earlier conversation and your own recent tool "
            "rounds (Auto run log). Treat the run log as ground truth for what you "
            "already tried; do not repeat failed approaches. Do NOT re-run reads or "
            "searches the run log shows already succeeded — their results are "
            "summarized there. If you have gathered enough to act, stop gathering "
            "and produce the deliverable now. "
            "Use append_working_buffer after each exchange.\n"
        )
        block = (
            f"\n\n{_WB_INJECT_START}\n---\n\nConversation summary (from working buffer):\n\n"
            + buf_content
            + danger_hint
            + f"{_WB_INJECT_END}"
        )
        sys_content = messages[0].get("content") or ""
        if not isinstance(sys_content, str):
            return False
        start = sys_content.find(_WB_INJECT_START)
        end = sys_content.find(_WB_INJECT_END)
        if start != -1 and end != -1 and end > start:
            sys_content = (sys_content[:start] + sys_content[end + len(_WB_INJECT_END):]).rstrip()
        messages[0]["content"] = sys_content.rstrip() + block
        return True
    except Exception:
        _log.debug("working buffer inject failed", exc_info=True)
        return False


# OCR word-by-word cleanup only needs the current chunk; keep 1.
_READ_FILE_CHUNK_CONTEXT_KEEP_OCR = 1
# Research / synthesis / review need accumulated content; keep more.
_READ_FILE_CHUNK_CONTEXT_KEEP_GENERAL = 6
# Legacy alias used by imports elsewhere.
_READ_FILE_CHUNK_CONTEXT_KEEP = _READ_FILE_CHUNK_CONTEXT_KEEP_OCR


def _is_ocr_session(messages: list) -> bool:
    """Return True if the conversation contains OCR correction tool calls."""
    for msg in messages:
        if not isinstance(msg, dict):
            continue
        tcs = msg.get("tool_calls")
        if isinstance(tcs, list):
            for tc in tcs:
                fn = tc if isinstance(tc, dict) else {}
                name = (fn.get("function") or fn).get("name", "") if isinstance(fn.get("function"), dict) else fn.get("name", "")
                if name == "submit_ocr_chunk_corrections":
                    return True
        content = msg.get("content") or ""
        if isinstance(content, str) and "submit_ocr_chunk_corrections" in content:
            if msg.get("role") in ("assistant", "tool"):
                return True
    return False


def _trim_old_read_file_chunk_results(messages: list, keep_last_n: int | None = None) -> list[tuple[int, int, int]]:
    """Replace old read_file_chunk tool results with short placeholders so each
    API request doesn't grow unbounded.

    Behavior adapts based on whether the session is an OCR correction flow:
    - OCR: keep only 1 chunk (word-for-word cleanup only needs the current one).
    - General (research/synthesis/review): keep 6 so the model retains gathered
      content for synthesis. Uses a neutral placeholder that doesn't mislead the
      model into thinking it's doing OCR.

    Returns list of (index, old_content_len, new_content_len) for incremental
    token estimate.
    """
    ocr_mode = _is_ocr_session(messages)
    if keep_last_n is None:
        keep_last_n = _READ_FILE_CHUNK_CONTEXT_KEEP_OCR if ocr_mode else _READ_FILE_CHUNK_CONTEXT_KEEP_GENERAL

    deltas: list[tuple[int, int, int]] = []
    if keep_last_n <= 0:
        return deltas
    indices = []
    for i, msg in enumerate(messages):
        if not isinstance(msg, dict):
            continue
        content = tool_result_text_for_trim(msg)
        if not isinstance(content, str) or len(content) < 2000:
            continue
        if "Chunk" in content and " of " in content and ("chars " in content or "Chunk " in content[:100]):
            indices.append(i)
    for i in indices[:-keep_last_n]:
        orig = tool_result_text_for_trim(messages[i])
        if len(orig) < 2000:
            continue
        first_line = orig.split("\n")[0] if orig else ""
        if ocr_mode:
            new_content = (
                f"[Chunk already corrected and saved. Continue with next chunk: read_file_chunk(chunk_index=N, num_chunks=1), "
                f"then submit_ocr_chunk_corrections with the corrected text. {first_line[:60]}.]"
            )
        else:
            new_content = (
                f"[Earlier chunk content trimmed from context. Re-read with "
                f"read_file_chunk(chunk_index=N) if you need this content again. {first_line[:80]}]"
            )
        if (messages[i].get("role") or "").strip() == "tool":
            messages[i] = {**messages[i], "content": new_content}
        else:
            messages[i] = {**messages[i], "content": f"<tool_response>\n{new_content}\n</tool_response>"}
        deltas.append((i, len(orig), len(new_content)))
    return deltas


def _content_chars(messages: list, start: int, end: int) -> int:
    """Sum message text + native tool_calls payload chars for token pressure."""
    total = 0
    for m in messages[start:end]:
        content = m.get("content")
        if isinstance(content, str):
            total += len(content)
        tool_calls = m.get("tool_calls")
        if tool_calls:
            try:
                total += len(json.dumps(tool_calls, ensure_ascii=False))
            except Exception:
                total += 200
    return total


def _delta_tc_field(tc, field, sub="function"):
    """Get a field from a tool-call delta (dict or Pydantic). Handles nested function.name / function.arguments."""
    if isinstance(tc, dict):
        fn = tc.get(sub) or tc
        return fn.get(field) if isinstance(fn, dict) else None
    fn = getattr(tc, sub, None)
    if fn is not None:
        return getattr(fn, field, None)
    return getattr(tc, field, None)


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


def _accumulate_stream_tool_call_delta(
    tool_calls_accum: dict[int, dict[str, str]],
    tc: Any,
) -> None:
    """Merge one incremental ``delta.tool_calls`` fragment into the accumulator."""
    idx = _resolve_stream_tool_call_index(tool_calls_accum, tc)
    if idx not in tool_calls_accum:
        tool_calls_accum[idx] = {"id": "", "name": "", "arguments": ""}
    tid = tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None)
    tname = _delta_tc_field(tc, "name")
    targs = _delta_tc_field(tc, "arguments")
    if tid:
        tool_calls_accum[idx]["id"] = tid
    if tname:
        tool_calls_accum[idx]["name"] = tname
    if targs is not None:
        tool_calls_accum[idx]["arguments"] = (
            tool_calls_accum[idx].get("arguments", "")
            + (targs if isinstance(targs, str) else str(targs or ""))
        )


def _normalize_tool_calls_for_mlx_native_history(tool_calls_list: list[dict[str, Any]]) -> None:
    """Rewrite xmltc_* ids so the next MLX turn sees OpenAI-style tool_call ids."""
    import secrets
    import time

    for tc in tool_calls_list:
        if not isinstance(tc, dict):
            continue
        tid = str(tc.get("id") or "").strip()
        if tid.startswith("call_"):
            continue
        tc["id"] = f"call_{int(time.time() * 1_000_000)}_{secrets.token_hex(4)}"


def try_mlx_non_stream_native_tool_calls(
    client: Any,
    *,
    model: str,
    messages: list[dict[str, Any]],
    tools: list[dict[str, Any]],
    max_tokens: int,
    temperature: float,
    extra_body: dict[str, Any] | None = None,
    tool_choice: str = "auto",
) -> list[dict[str, str]]:
    """One-shot non-stream completion when streaming omitted native tool_calls."""
    try:
        kwargs: dict[str, Any] = dict(
            model=model,
            messages=messages,
            tools=tools,
            tool_choice=tool_choice,
            stream=False,
            max_tokens=max(256, min(int(max_tokens), 8192)),
            temperature=temperature,
        )
        if extra_body:
            kwargs["extra_body"] = extra_body
        resp = client.chat.completions.create(**kwargs)
        ch0 = resp.choices[0] if getattr(resp, "choices", None) else None
        msg = getattr(ch0, "message", None) if ch0 else None
        tcs = getattr(msg, "tool_calls", None) or []
        out: list[dict[str, str]] = []
        for tc in tcs:
            name = _delta_tc_field(tc, "name")
            if not str(name or "").strip():
                continue
            tid = tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None)
            args = _delta_tc_field(tc, "arguments")
            args_str = args if isinstance(args, str) else (json.dumps(args) if args is not None else "{}")
            out.append(
                {
                    "id": str(tid or ""),
                    "name": str(name).strip(),
                    "arguments": args_str or "{}",
                }
            )
        return out
    except Exception:
        _log.debug("MLX non-stream native tool_calls recovery failed", exc_info=True)
        return []


def _harvest_stream_choice_native_tool_calls(
    tool_calls_accum: dict[int, dict[str, str]],
    choice: Any,
) -> None:
    """Merge ``choice.message.tool_calls`` when the server omits incremental deltas.

    mlx-openai-server (and some other OpenAI-compatible stacks) may stream assistant
    content but only attach complete ``tool_calls`` on the final chunk's ``message``.
  """
    msg = getattr(choice, "message", None)
    if msg is None:
        return
    tcs = getattr(msg, "tool_calls", None) or []
    for tc in tcs:
        idx = _resolve_stream_tool_call_index(tool_calls_accum, tc)
        tid = tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None)
        tname = _delta_tc_field(tc, "name")
        targs = _delta_tc_field(tc, "arguments")
        if not str(tname or "").strip():
            continue
        if idx not in tool_calls_accum:
            tool_calls_accum[idx] = {"id": "", "name": "", "arguments": ""}
        if tid:
            tool_calls_accum[idx]["id"] = str(tid)
        tool_calls_accum[idx]["name"] = str(tname).strip()
        if targs is not None:
            tool_calls_accum[idx]["arguments"] = (
                targs if isinstance(targs, str) else str(targs or "")
            )


# Proactive nudge: remind the model to offer suggestions or use memory when it hasn't for a while.
PROACTIVE_NUDGE_MIN_TURNS = 5
PROACTIVE_NUDGE_PROBABILITY = 0.4
PROACTIVE_NUDGE_RECENT_ROUNDS = 2  # don't nudge if agent used proactive tools in last N assistant tool-call rounds


def _agent_was_proactive_recently(messages, last_n_tool_rounds=PROACTIVE_NUDGE_RECENT_ROUNDS):
    """True if in the last `last_n_tool_rounds` assistant tool-calling turns the model used proactive tools."""
    _PROACTIVE_TOOLS = {
        "append_to_memory",
        "write_session_state",
        "append_working_buffer",
        "resolve_drive",
        "create_drive",
    }
    seen = 0
    for m in reversed(messages):
        if m.get("role") != "assistant":
            continue
        tool_calls = m.get("tool_calls") or []
        if not tool_calls:
            content = m.get("content") or ""
            if isinstance(content, str) and "<tool_call" in content.lower():
                try:
                    tool_calls = _parse_xml_tool_calls(content)
                except Exception:
                    tool_calls = []
        if not tool_calls:
            continue
        seen += 1
        if seen > last_n_tool_rounds:
            break
        for tc in tool_calls:
            fn = tc.get("function") if isinstance(tc, dict) else None
            name = ""
            if isinstance(fn, dict):
                name = fn.get("name") or ""
                args = fn.get("arguments") or ""
            else:
                name = (tc.get("name") or "") if isinstance(tc, dict) else ""
                args = (tc.get("arguments") or "") if isinstance(tc, dict) else ""
            if name in _PROACTIVE_TOOLS:
                return True
            if name == "write_text_file":
                if "USER.md" in args or "user.md" in args:
                    return True
    return False


def _is_synthetic_startup_init(messages: list[dict[str, Any]] | None) -> bool:
    """True when messages are exactly one synthetic ``init`` user turn.

    Used for tests and diagnostics. ``run_raw_agent`` no longer switches to a
    tool-less mini-prompt for this case — startup uses the full harness so
    personality and MEMORY stay in context.
    """
    conversation_only = [
        m for m in (messages or [])
        if isinstance(m, dict) and m.get("role") != "system"
    ]
    if len(conversation_only) != 1:
        return False
    msg = conversation_only[0]
    if msg.get("role") != "user":
        return False
    content = msg.get("content")
    return isinstance(content, str) and content.strip().lower() == "init"


_PLAIN_CHAT_TOOL_TRIGGER_RE = re.compile(
    r"\b("
    r"search|browse|google|look\s+up|web|online|current|latest|today|now|"
    r"weather|temperature|news|stock|price|score|"
    r"read|write|edit|create|delete|save|append|open|move|copy|rename|"
    r"file|folder|directory|pdf|image|screenshot|ocr|clipboard|"
    r"run|execute|shell|terminal|python|code|script|command|"
    r"email|mail|calendar|remind|schedule|timer|"
    r"research|investigate|compare|find|check|"
    r"remember|recall|memory|memories|previous|past|earlier|"
    r"conversation|conversations|"
    r"met\s+before|(?:we(?:'|’)ve|we\s+have)\s+met|know\s+me|"
    r"commit|push|git|plugin|workflow|skill|mesh|peer"
    r")\b",
    re.IGNORECASE,
)


def _latest_real_user_text(messages: list[dict[str, Any]] | None) -> str:
    for msg in reversed(messages or []):
        if not is_real_user_message(msg):
            continue
        content = msg.get("content")
        if isinstance(content, str) and content.strip():
            return content.strip()
        if isinstance(content, list):
            for part in content:
                if isinstance(part, dict) and part.get("type") == "text":
                    text = (part.get("text") or "").strip()
                    if text:
                        return text
    return ""




def _has_explicit_recall_intent(text: str | None) -> bool:
    """Conservative detector for user questions about continuity/past chats."""
    from utils.recall_signals import has_explicit_recall_intent

    return has_explicit_recall_intent(text)


def _escape_untrusted_system_context(text: str) -> str:
    return (
        (text or "")
        .replace("&", "&amp;")
        .replace("<", "&lt;")
        .replace(">", "&gt;")
    )


def _thread_disallows_recall_preflight(thread_id: str | None) -> bool:
    tid = (thread_id or "").strip().lower()
    if not tid:
        return False
    return tid.startswith((
        "mesh",
        "group",
        "peer",
        "btw",
        "heartbeat",
        "background_heartbeat",
        "drive_action",
    ))


_UNSUPPORTED_NO_RECALL_CLAIM_RE = re.compile(
    r"("
    r"\b(?:no\s+way|can't|cannot|can\s+not|don't|do\s+not)\b.{0,100}\b(?:look\s+back|access|search|recall|remember|memory|memories|previous|past|conversation|history)\b|"
    r"\b(?:no|not)\s+(?:access|ability)\b.{0,80}\b(?:memory|memories|previous|past|conversation|history)\b|"
    r"\bi\s+(?:only|just)\s+have\s+(?:this|the current)\s+(?:chat|conversation|thread)\b"
    r")",
    re.IGNORECASE | re.DOTALL,
)


def _looks_like_unsupported_no_recall_claim(text: str | None) -> bool:
    return bool(_UNSUPPORTED_NO_RECALL_CLAIM_RE.search(text or ""))


def _looks_like_plain_conversational_turn(
    messages: list[dict[str, Any]] | None,
    *,
    has_media_attachment: bool,
    research_mode: bool,
    include_cli: bool,
    quick_panel: bool,
    subagent_mode: bool,
    policy_context: Any | None,
) -> bool:
    """Best-effort latency path for short turns that clearly do not need tools.

    This intentionally errs toward the full agent path. It exists to keep simple
    chat from paying local-LLM prefill costs for the entire tool harness.
    """
    if subagent_mode or has_media_attachment or research_mode:
        return False
    if bool(getattr(policy_context, "idle_tick", False)):
        return False
    if bool(getattr(policy_context, "messaging", False)) or bool(getattr(policy_context, "mesh_sourced", False)):
        return False
    text = _latest_real_user_text(messages)
    if not text:
        return False
    if len(text) > 500:
        return False
    if "\n" in text and len(text) > 180:
        return False
    lowered = text.lower()
    if "```" in lowered or "<tool" in lowered:
        return False
    if _PLAIN_CHAT_TOOL_TRIGGER_RE.search(lowered):
        return False
    try:
        if _multi_step_user_message(text):
            return False
    except Exception:
        pass
    return True
