"""Per-turn Memory Inspector SSE payload (`memory_turn_summary`).

Runs on the agent worker thread; uses threading.local so tools invoked from the
same thread can append retrieval/write rows without threading.ContextVar quirks.

Schema v2 adds ``sources`` and ``influence_headline`` for answer-centric UI.
"""

from __future__ import annotations

import json
import threading
import uuid
from typing import Any

from utils.redaction import cap_str, redact_text

_tls = threading.local()

SCHEMA_VERSION = 2
MAX_PAYLOAD_BYTES = 56 * 1024
MAX_RETRIEVAL_GROUPS = 12
MAX_RETRIEVAL_HITS_PER_GROUP = 8
MAX_SNIPPET_CHARS = 240
MAX_SOURCES = 32
MAX_SOURCE_PREVIEW_CHARS = 200
_DEFAULT_COVERAGE_NOTE = (
    "Opaque sources (non-memory read_file paths, tools without inspector hooks) may still "
    "shape answers without appearing here."
)

SOURCE_ACTIVE_MEMORY = "active_memory"
SOURCE_MEMORY_PREFETCH = "memory_prefetch"
SOURCE_SESSION_RECALL = "session_recall"
SOURCE_TOOL_SEARCH = "tool_search"
SOURCE_TOOL_WRITE = "tool_write"
SOURCE_PATH_READ = "path_read"


def tls_get_trace() -> "MemoryTurnTrace | None":
    return getattr(_tls, "memory_trace", None)


def tls_set_trace(t: "MemoryTurnTrace") -> None:
    _tls.memory_trace = t


def tls_clear_trace() -> None:
    if hasattr(_tls, "memory_trace"):
        try:
            delattr(_tls, "memory_trace")
        except AttributeError:
            pass


def _cap_preview(text: str | None) -> str:
    if not text:
        return ""
    return cap_str(redact_text(text), MAX_SOURCE_PREVIEW_CHARS)


def policies_defaults_policy_only(policies_echo: dict[str, Any] | None) -> bool:
    """True when tier/policy echo is the usual conservative default (no explicit recall)."""
    pe = policies_echo or {}
    if pe.get("explicit_recall_intent"):
        return False
    if pe.get("cross_thread_signal"):
        return False
    if str(pe.get("recall_preflight_status") or "") not in ("", "skipped"):
        return False
    profile = str(pe.get("memory_retrieval_profile") or "conservative").strip().lower()
    if profile and profile != "conservative":
        return False
    tier = str(pe.get("search_sessions_tier") or "").strip().lower()
    reason = str(pe.get("search_sessions_tier_reason") or "").strip().lower()
    if tier and tier != "recall":
        return False
    if reason and reason not in ("ambiguous_default_recall", "empty_default_recall", ""):
        return False
    return True


def build_influence_headline(trace: "MemoryTurnTrace") -> str:
    """One-line summary of what memory shaped this turn."""
    if trace.skipped_reason:
        return f"Memory trace skipped ({trace.skipped_reason})."

    kinds: list[str] = []
    tool_search = 0
    tool_write = 0
    for src in trace.sources:
        k = str(src.get("kind") or "")
        if k == SOURCE_ACTIVE_MEMORY:
            kinds.append("active_memory")
        elif k == SOURCE_MEMORY_PREFETCH:
            kinds.append("prefetch")
        elif k == SOURCE_SESSION_RECALL:
            kinds.append("session recall")
        elif k == SOURCE_TOOL_SEARCH:
            tool_search += 1
        elif k == SOURCE_TOOL_WRITE:
            tool_write += 1
        elif k == SOURCE_PATH_READ:
            kinds.append("path read")

    parts: list[str] = []
    if "active_memory" in kinds:
        hub = int(trace.preloaded.get("hub_chars") or 0)
        shards = trace.preloaded.get("shards") or []
        if isinstance(shards, list) and shards:
            parts.append(f"Active Memory hub + {len(shards)} shard(s)")
        elif hub > 0:
            parts.append("Active Memory hub")
        elif trace.prewarm_hit:
            parts.append("Active Memory (prewarm cache)")
    if "prefetch" in kinds:
        parts.append("FTS prefetch")
    if "session recall" in kinds:
        parts.append("session recall")
    if tool_search:
        parts.append(f"{tool_search} memory search tool call(s)")
    if tool_write:
        parts.append(f"{tool_write} memory write(s)")
    if "path read" in kinds and not tool_search:
        parts.append("memory path read(s)")

    if not parts:
        if trace.active_memory_used or int(trace.preloaded.get("hub_chars") or 0) > 0:
            return "Active Memory preloaded; no memory tools this turn."
        return "No memory influence recorded this turn."

    tail = ""
    if not tool_search and not tool_write:
        tail = "; no memory tools"
    return ", ".join(parts) + tail + "."


class MemoryTurnTrace:
    """Mutable collector; finalized to JSON-ready dict."""

    __slots__ = (
        "thread_id",
        "turn_id",
        "skipped_reason",
        "coverage_notes",
        "preloaded",
        "active_memory_used",
        "prewarm_hit",
        "retrieval_groups",
        "writes",
        "policies_echo",
        "memory_path_reads",
        "sources",
    )

    def __init__(self, thread_id: str) -> None:
        self.thread_id = (thread_id or "").strip()
        self.turn_id = uuid.uuid4().hex[:16]
        self.skipped_reason: str | None = None
        self.coverage_notes: list[str] = [_DEFAULT_COVERAGE_NOTE]
        self.preloaded: dict[str, Any] = {
            "hub_path": "MEMORY.md",
            "hub_chars": 0,
            "layout": "unknown",
            "shards": [],
            "active_memory_character_count": None,
            "prewarm_hit": False,
            "preload_redacted": False,
        }
        self.active_memory_used: bool = False
        self.prewarm_hit: bool = False
        self.retrieval_groups: list[dict[str, Any]] = []
        self.writes: list[dict[str, Any]] = []
        self.policies_echo: dict[str, Any] = {}
        self.memory_path_reads: list[dict[str, Any]] = []
        self.sources: list[dict[str, Any]] = []

    def _append_source(
        self,
        *,
        kind: str,
        label: str,
        chars: int,
        preview: str | None = None,
        meta: dict[str, Any] | None = None,
    ) -> None:
        if len(self.sources) >= MAX_SOURCES:
            return
        row: dict[str, Any] = {
            "kind": kind,
            "label": cap_str(label, 120),
            "chars": max(0, int(chars)),
        }
        prev = _cap_preview(preview)
        if prev:
            row["preview"] = prev
        if meta:
            row["meta"] = meta
        self.sources.append(row)

    def mark_eligibility(self, *, subagent_mode: bool, is_heartbeat: bool) -> None:
        from utils.active_memory import is_eligible

        if subagent_mode:
            self.skipped_reason = "subagent_mode"
            return
        if is_heartbeat:
            self.skipped_reason = "heartbeat"
            return
        ok, reason = is_eligible(
            thread_id=self.thread_id, subagent_mode=False, is_heartbeat=False,
        )
        if not ok:
            self.skipped_reason = reason

    def attach_preloaded(
        self,
        digest: dict[str, Any] | None,
        *,
        active_memory_used: bool,
        active_block_chars: int,
        prewarm_hit: bool,
        hub_preview: str | None = None,
    ) -> None:
        self.active_memory_used = bool(active_memory_used)
        self.prewarm_hit = bool(prewarm_hit)
        if digest and isinstance(digest, dict):
            self.preloaded.update(digest)
        self.preloaded["active_memory_character_count"] = (
            int(active_block_chars) if active_memory_used else 0
        )
        self.preloaded["prewarm_hit"] = self.prewarm_hit
        if self.skipped_reason:
            self.preloaded["preload_redacted"] = True
            self.preloaded["shards"] = []
            self.preloaded["hub_chars"] = 0
            self.preloaded["layout"] = "redacted"
            return
        if active_memory_used or active_block_chars > 0 or prewarm_hit:
            record_source_active_memory(
                trace=self,
                block_chars=active_block_chars,
                digest=self.preloaded,
                hub_preview=hub_preview,
                prewarm_hit=prewarm_hit,
            )

    def set_policies_echo(self, data: dict[str, Any] | None) -> None:
        if data:
            self.policies_echo = dict(data)

    def record_retrieval_batch(
        self,
        *,
        tool_name: str,
        query: str,
        hit_paths: list[str],
        snippets: list[str] | None = None,
    ) -> None:
        if tls_get_trace() is not self:
            return
        if len(self.retrieval_groups) >= MAX_RETRIEVAL_GROUPS:
            return
        q = cap_str(redact_text(query), min(280, MAX_SNIPPET_CHARS + 40))
        paths_out: list[str] = []
        for p in hit_paths[:MAX_RETRIEVAL_HITS_PER_GROUP]:
            paths_out.append(cap_str(p, 200))
        snip_out: list[str] = []
        if snippets:
            for s in snippets[: min(3, MAX_RETRIEVAL_HITS_PER_GROUP)]:
                snip_out.append(cap_str(redact_text(s), MAX_SNIPPET_CHARS))
        self.retrieval_groups.append(
            {"tool_name": tool_name, "query_preview": q, "paths": paths_out, "snippets": snip_out}
        )
        preview = snip_out[0] if snip_out else (paths_out[0] if paths_out else q)
        self._append_source(
            kind=SOURCE_TOOL_SEARCH,
            label=tool_name,
            chars=sum(len(s) for s in snip_out) or len(q),
            preview=preview,
            meta={"tool_name": tool_name, "paths": paths_out, "query_preview": q},
        )

    def record_write(
        self,
        *,
        tool_name: str,
        path_rel: str,
        outcome: str,
        delta_chars: int | None = None,
        detail: str | None = None,
        promotion_warning: str | None = None,
    ) -> None:
        if tls_get_trace() is not self:
            return
        if len(self.writes) >= 24:
            return
        row: dict[str, Any] = {
            "tool_name": tool_name,
            "path": cap_str(path_rel, 220),
            "outcome": cap_str(outcome, 80),
        }
        if delta_chars is not None:
            row["delta_chars"] = int(delta_chars)
        if detail:
            row["detail"] = cap_str(detail, 300)
        if promotion_warning:
            row["promotion_warning"] = cap_str(promotion_warning, 200)
        self.writes.append(row)
        self._append_source(
            kind=SOURCE_TOOL_WRITE,
            label=f"{tool_name} → {path_rel}",
            chars=int(delta_chars or 0),
            preview=detail or outcome,
            meta=dict(row),
        )

    def record_memory_path_read(
        self,
        *,
        path_rel: str,
        bytes_read: int,
        truncated: bool,
    ) -> None:
        if tls_get_trace() is not self:
            return
        if len(self.memory_path_reads) >= 16:
            return
        row = {
            "path": cap_str(path_rel, 220),
            "bytes_read": max(0, int(bytes_read)),
            "truncated": bool(truncated),
        }
        self.memory_path_reads.append(row)
        self._append_source(
            kind=SOURCE_PATH_READ,
            label=path_rel,
            chars=int(bytes_read),
            preview=f"{bytes_read} bytes" + (" (truncated)" if truncated else ""),
            meta=row,
        )

    def to_payload_dict(self) -> dict[str, Any]:
        pre = dict(self.preloaded)
        if self.skipped_reason:
            pre["preload_redacted"] = True
            pre["shards"] = []
            pre["hub_chars"] = 0
        defaults_only = policies_defaults_policy_only(self.policies_echo)
        headline = build_influence_headline(self)
        out: dict[str, Any] = {
            "type": "memory_turn_summary",
            "schema_version": SCHEMA_VERSION,
            "turn_id": self.turn_id,
            "phase": "complete",
            "skipped_reason": self.skipped_reason,
            "coverage_notes": list(self.coverage_notes),
            "preloaded": pre,
            "retrieval": list(self.retrieval_groups),
            "writes": list(self.writes),
            "policies_echo": dict(self.policies_echo),
            "memory_path_reads": list(self.memory_path_reads),
            "sources": list(self.sources),
            "influence_headline": headline,
            "defaults_policy_only": defaults_only,
        }
        raw = json.dumps(out, separators=(",", ":"), default=str).encode("utf-8")
        if len(raw) > MAX_PAYLOAD_BYTES:
            out["truncated"] = True
            out["retrieval"] = out["retrieval"][:3]
            out["writes"] = out["writes"][:5]
            out["memory_path_reads"] = out["memory_path_reads"][:8]
            out["sources"] = out["sources"][:12]
            if isinstance(out.get("preloaded"), dict) and out["preloaded"].get("shards"):
                out["preloaded"]["shards"] = out["preloaded"]["shards"][:4]
            raw2 = json.dumps(out, separators=(",", ":"), default=str).encode("utf-8")
            if len(raw2) > MAX_PAYLOAD_BYTES:
                out["preloaded"] = {"hub_chars": 0, "shards": [], "truncated_hard": True}
                out["retrieval"] = []
                out["writes"] = []
                out["memory_path_reads"] = []
                out["sources"] = out["sources"][:6]
        return out


def record_source_active_memory(
    *,
    trace: MemoryTurnTrace | None = None,
    block_chars: int,
    digest: dict[str, Any] | None,
    hub_preview: str | None = None,
    prewarm_hit: bool = False,
) -> None:
    tr = trace or tls_get_trace()
    if tr is None:
        return
    d = digest or {}
    hub_chars = int(d.get("hub_chars") or 0)
    layout = str(d.get("layout") or "unknown")
    shards = d.get("shards") if isinstance(d.get("shards"), list) else []
    label = "Active Memory"
    if prewarm_hit:
        label = "Active Memory (prewarm cache)"
    elif shards:
        label = f"Active Memory ({len(shards)} shard(s))"
    elif hub_chars > 0:
        label = "Active Memory (MEMORY.md hub)"
    preview = hub_preview or str(d.get("hub_preview") or "")
    tr._append_source(
        kind=SOURCE_ACTIVE_MEMORY,
        label=label,
        chars=max(block_chars, hub_chars),
        preview=preview or None,
        meta={
            "prewarm_hit": bool(prewarm_hit),
            "layout": layout,
            "hub_chars": hub_chars,
            "shards": shards[:8],
        },
    )


def record_source_memory_prefetch(
    *,
    block_chars: int,
    digest: dict[str, Any] | None,
    trace: MemoryTurnTrace | None = None,
) -> None:
    tr = trace or tls_get_trace()
    if tr is None or block_chars <= 0:
        return
    d = digest or {}
    files = d.get("files") if isinstance(d.get("files"), list) else []
    n = len(files)
    label = f"Memory prefetch ({n} file(s))" if n else "Memory prefetch"
    preview = ""
    if files and isinstance(files[0], dict):
        preview = str(files[0].get("preview") or files[0].get("path") or "")
    if not preview and d.get("preview"):
        preview = str(d.get("preview"))
    tr._append_source(
        kind=SOURCE_MEMORY_PREFETCH,
        label=label,
        chars=block_chars,
        preview=preview,
        meta=d,
    )


def record_source_session_recall(
    *,
    status: str,
    reason: str,
    result_chars: int,
    thread_count: int,
    truncated: bool,
    block_chars: int,
    trace: MemoryTurnTrace | None = None,
) -> None:
    tr = trace or tls_get_trace()
    if tr is None:
        return
    st = (status or "skipped").strip()
    label = f"Session recall ({st})"
    if result_chars > 0:
        label = f"Session recall — {result_chars} chars from {thread_count} thread(s)"
    tr._append_source(
        kind=SOURCE_SESSION_RECALL,
        label=label,
        chars=max(block_chars, result_chars),
        preview=cap_str(reason, MAX_SOURCE_PREVIEW_CHARS) if reason else None,
        meta={
            "status": st,
            "reason": cap_str(reason, 200),
            "result_chars": int(result_chars),
            "thread_count": int(thread_count),
            "truncated": bool(truncated),
        },
    )


def memory_turn_trace_begin(thread_id: str, summary_ref: list[Any] | None) -> None:
    """Install a fresh trace for this worker thread."""
    if summary_ref is None:
        return
    tls_clear_trace()
    t = MemoryTurnTrace(thread_id)
    tls_set_trace(t)


def _empty_payload_v2(*, skipped_reason: str | None = None, error: str | None = None) -> dict[str, Any]:
    out: dict[str, Any] = {
        "type": "memory_turn_summary",
        "schema_version": SCHEMA_VERSION,
        "phase": "complete",
        "skipped_reason": skipped_reason,
        "preloaded": {"hub_chars": 0, "shards": []},
        "retrieval": [],
        "writes": [],
        "policies_echo": {},
        "memory_path_reads": [],
        "sources": [],
        "influence_headline": "No memory influence recorded this turn.",
        "defaults_policy_only": True,
        "coverage_notes": [_DEFAULT_COVERAGE_NOTE],
    }
    if error:
        out["error"] = error
    return out


def memory_turn_trace_finalize(summary_ref: list[Any] | None) -> None:
    """Serialize to summary_ref[0] and clear TLS for instrumented turns only."""
    if summary_ref is None:
        return
    try:
        if summary_ref and summary_ref[0] is not None:
            return
        tr = tls_get_trace()
        if tr is None:
            summary_ref[0] = _empty_payload_v2(skipped_reason="internal_no_trace")
        else:
            summary_ref[0] = tr.to_payload_dict()
    except Exception:
        summary_ref[0] = _empty_payload_v2(
            skipped_reason=None,
            error=cap_str(redact_text("serialize_failed"), 120),
        )
    finally:
        tls_clear_trace()


def record_memory_search_results(
    *,
    tool_name: str,
    query: str,
    fused_snippets: list[str],
    max_hits: int = MAX_RETRIEVAL_HITS_PER_GROUP,
) -> None:
    """Derive coarse paths/snippets from fused search lines (bounded)."""
    tr = tls_get_trace()
    if tr is None or not fused_snippets:
        return
    paths: list[str] = []
    snips: list[str] = []
    import re as _re

    path_pat = _re.compile(
        r"(?:^|\n)\s*\*?\*?((?:MEMORY\.md|memory/[^\\\n:]+?\.md))(?:\*\*)?\s*[:\n]",
        _re.IGNORECASE | _re.MULTILINE,
    )
    for block in fused_snippets[:max_hits]:
        m = path_pat.search(block)
        rel = ""
        if m:
            rel = (m.group(1) or "").strip().replace("\\", "/")
        if rel:
            paths.append(rel)
        body = block.strip()
        if body:
            snips.append(body)

    tr.record_retrieval_batch(
        tool_name=tool_name,
        query=query,
        hit_paths=paths or [],
        snippets=snips or None,
    )


def workspace_relative_path(workspace_root: str, abs_path: str) -> str | None:
    """Return posix workspace-relative path, or None if abs_path is outside workspace."""
    import os

    try:
        ws = os.path.abspath(workspace_root)
        ap = os.path.abspath(abs_path)
        if os.path.commonpath([ws, ap]) != ws:
            return None
        rel = os.path.relpath(ap, ws).replace("\\", "/")
        return rel if rel and rel != "." else None
    except (ValueError, OSError):
        return None


def is_memory_inspector_workspace_rel(rel: str) -> bool:
    """True for hub + shard/archive markdown paths we surface in the Memory Inspector."""
    r = (rel or "").replace("\\", "/").strip().lower().lstrip("./")
    if r == "memory.md":
        return True
    return r.startswith("memory/shards/") or r.startswith("memory/archives/")


def record_memory_path_read_file(*, path_rel: str, bytes_read: int, truncated: bool) -> None:
    """Record a bounded read_file touch against memory paths for the active turn trace."""
    tr = tls_get_trace()
    if tr is None:
        return
    try:
        tr.record_memory_path_read(
            path_rel=path_rel,
            bytes_read=bytes_read,
            truncated=truncated,
        )
    except Exception:
        pass
