"""Per-turn Knowledge Inspector SSE payload (`knowledge_turn_summary`).

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

from __future__ import annotations

import json
import threading
import uuid
from contextlib import contextmanager
from typing import Any, Iterator

from utils.redaction import cap_str, redact_text

_tls = threading.local()

SCHEMA_VERSION = 1
MAX_PAYLOAD_BYTES = 48 * 1024
MAX_INJECTION_HITS = 8
MAX_TOOL_CALLS = 10
MAX_QUERY_HITS = 6
MAX_SNIPPET_CHARS = 240
MAX_QUERY_PREVIEW_CHARS = 200
MAX_SOURCES = 32
MAX_SOURCE_PREVIEW_CHARS = 200
MAX_CATALOG_LABELS = 16

SOURCE_INJECTION_SEARCH = "injection_search"
SOURCE_INJECTION_TOPIC = "injection_topic"
SOURCE_CATALOG = "catalog"
SOURCE_TOOL_QUERY = "tool_query"
SOURCE_TOOL_READ = "tool_read"
SOURCE_TOOL_COMBINED = "tool_combined"

_DEFAULT_COVERAGE_NOTES = [
    "Catalog lists corpus names only — not synthesis text.",
    "Combined memory+knowledge search appears as one tool row; memory leg is in Memory Inspector.",
    "Library Ask and inline chat citation chips are separate surfaces.",
]


def tls_get_trace() -> "KnowledgeTurnTrace | None":
    return getattr(_tls, "knowledge_trace", None)


def tls_set_trace(t: "KnowledgeTurnTrace") -> None:
    _tls.knowledge_trace = t


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


@contextmanager
def suppress_nested_query_knowledge_trace() -> Iterator[None]:
    """Skip ``query_knowledge`` trace rows while the combined tool invokes it internally."""
    tr = tls_get_trace()
    if tr is None:
        yield
        return
    prev = tr._suppress_nested_query_knowledge
    tr._suppress_nested_query_knowledge = True
    try:
        yield
    finally:
        tr._suppress_nested_query_knowledge = prev


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


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


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


def _sanitize_hit_row(hit: dict[str, Any]) -> dict[str, Any]:
    """Strip filesystem paths; keep corpus slug/label and source name only."""
    out: dict[str, Any] = {
        "corpus_slug": cap_str(str(hit.get("corpus_slug") or ""), 80),
        "corpus_label": cap_str(str(hit.get("corpus_label") or hit.get("corpus_slug") or ""), 120),
        "source_name": cap_str(str(hit.get("source_name") or ""), 120),
        "snippet": _cap_snippet(str(hit.get("snippet") or "")),
    }
    for key in ("line_start", "line_end", "char_start", "char_end"):
        if hit.get(key) is not None:
            try:
                out[key] = int(hit[key])
            except (TypeError, ValueError):
                pass
    if hit.get("score") is not None:
        try:
            out["score"] = float(hit["score"])
        except (TypeError, ValueError):
            pass
    return out


def build_influence_headline(trace: "KnowledgeTurnTrace") -> str:
    if trace.skipped_reason:
        return f"Knowledge trace skipped ({trace.skipped_reason})."

    parts: list[str] = []
    inj = trace.injection
    path = str(inj.get("path") or "absent")
    hits = inj.get("hits") if isinstance(inj.get("hits"), list) else []
    n_hits = len(hits)

    if path == "search" and n_hits:
        parts.append(f"Injected {n_hits} passage(s) from studied material (search)")
    elif path == "topic_fallback" and n_hits:
        parts.append(f"Injected topic/synthesis fallback ({n_hits} corpora)")
    elif path == "topic_fallback":
        parts.append("Injected topic/synthesis fallback")
    elif trace.catalog.get("present"):
        count = int(trace.catalog.get("corpus_count") or 0)
        parts.append(f"No injection; catalog only ({count} corpora)")

    tool_counts: dict[str, int] = {}
    for row in trace.tool_calls:
        name = str(row.get("tool") or "")
        if name:
            tool_counts[name] = tool_counts.get(name, 0) + 1

    if tool_counts.get("search_memory_and_knowledge_corpora"):
        k = tool_counts.pop("search_memory_and_knowledge_corpora")
        parts.append(f"search_memory_and_knowledge_corpora ×{k}")
    for name, count in sorted(tool_counts.items()):
        parts.append(f"{name} ×{count}")

    if not parts:
        return "No knowledge influence this turn."
    return "; ".join(parts) + "."


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

    __slots__ = (
        "thread_id",
        "turn_id",
        "skipped_reason",
        "coverage_notes",
        "injection",
        "catalog",
        "tool_calls",
        "sources",
        "_suppress_nested_query_knowledge",
    )

    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] = list(_DEFAULT_COVERAGE_NOTES)
        self.injection: dict[str, Any] = {"path": "absent", "chars_injected": 0, "hits": []}
        self.catalog: dict[str, Any] = {
            "present": False,
            "corpus_count": 0,
            "chars": 0,
            "corpus_labels": [],
        }
        self.tool_calls: list[dict[str, Any]] = []
        self.sources: list[dict[str, Any]] = []
        self._suppress_nested_query_knowledge = False

    def _append_source(
        self,
        *,
        kind: str,
        label: str,
        preview: str | None = None,
        meta: dict[str, Any] | None = None,
    ) -> None:
        if self.skipped_reason or len(self.sources) >= MAX_SOURCES:
            return
        row: dict[str, Any] = {
            "kind": kind,
            "label": cap_str(label, 120),
        }
        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 record_injection(
        self,
        *,
        path: str,
        trace_rows: list[dict[str, Any]],
        chars_injected: int,
    ) -> None:
        if tls_get_trace() is not self or self.skipped_reason:
            return
        hits = [_sanitize_hit_row(h) for h in trace_rows[:MAX_INJECTION_HITS]]
        self.injection = {
            "path": cap_str(path, 40) or "absent",
            "chars_injected": max(0, int(chars_injected)),
            "hits": hits,
        }
        kind = SOURCE_INJECTION_SEARCH if path == "search" else SOURCE_INJECTION_TOPIC
        if path == "absent" or path == "empty":
            return
        for h in hits:
            label = h.get("corpus_label") or h.get("corpus_slug") or "?"
            src = h.get("source_name") or ""
            if src and src != "synthesis":
                label = f"{label} › {src}"
            self._append_source(
                kind=kind if path == "search" else SOURCE_INJECTION_TOPIC,
                label=str(label),
                preview=h.get("snippet"),
                meta=h,
            )

    def record_catalog(
        self,
        *,
        present: bool,
        corpus_count: int,
        chars: int,
        corpus_labels: list[str] | None = None,
    ) -> None:
        if tls_get_trace() is not self or self.skipped_reason:
            return
        labels = [cap_str(str(x), 120) for x in (corpus_labels or [])[:MAX_CATALOG_LABELS]]
        self.catalog = {
            "present": bool(present),
            "corpus_count": max(0, int(corpus_count)),
            "chars": max(0, int(chars)),
            "corpus_labels": labels,
        }
        if present and labels:
            self._append_source(
                kind=SOURCE_CATALOG,
                label=f"Catalog ({corpus_count} corpora)",
                preview=", ".join(labels[:4]),
                meta={"corpus_count": corpus_count, "chars": chars},
            )

    def record_tool_query(
        self,
        *,
        tool: str,
        query: str,
        corpus_filter: str | None,
        hits: list[dict[str, Any]],
        outcome: str,
    ) -> None:
        if tls_get_trace() is not self or self.skipped_reason:
            return
        if tool == "query_knowledge" and self._suppress_nested_query_knowledge:
            return
        if len(self.tool_calls) >= MAX_TOOL_CALLS:
            return
        hit_rows = [_sanitize_hit_row(h) for h in hits[:MAX_QUERY_HITS]]
        row: dict[str, Any] = {
            "tool": cap_str(tool, 80),
            "query_preview": _cap_query(query),
            "corpus_filter": cap_str(corpus_filter, 80) if corpus_filter else None,
            "hit_count": len(hits),
            "outcome": cap_str(outcome, 40),
            "hits": hit_rows,
        }
        self.tool_calls.append(row)
        preview = hit_rows[0].get("snippet") if hit_rows else _cap_query(query)
        self._append_source(
            kind=SOURCE_TOOL_QUERY,
            label=tool,
            preview=preview,
            meta=row,
        )

    def record_tool_read(
        self,
        *,
        corpus_slug: str,
        source_name: str,
        bytes_read: int,
        truncated: bool,
        offset: int = 0,
    ) -> None:
        if tls_get_trace() is not self or self.skipped_reason:
            return
        if len(self.tool_calls) >= MAX_TOOL_CALLS:
            return
        row: dict[str, Any] = {
            "tool": "read_knowledge_source",
            "corpus_slug": cap_str(corpus_slug, 80),
            "source_name": cap_str(source_name, 120),
            "outcome": "read",
            "bytes_read": max(0, int(bytes_read)),
            "truncated": bool(truncated),
            "offset": max(0, int(offset)),
        }
        self.tool_calls.append(row)
        label = f"{corpus_slug} › {source_name}" if source_name else corpus_slug
        self._append_source(
            kind=SOURCE_TOOL_READ,
            label=f"read_knowledge_source → {label}",
            preview=f"{bytes_read} bytes" + (" (truncated)" if truncated else ""),
            meta=row,
        )

    def record_tool_list_sources(
        self,
        *,
        corpus_slug: str,
        corpus_label: str,
        source_count: int,
    ) -> None:
        if tls_get_trace() is not self or self.skipped_reason:
            return
        if len(self.tool_calls) >= MAX_TOOL_CALLS:
            return
        row: dict[str, Any] = {
            "tool": "read_knowledge_source",
            "corpus_slug": cap_str(corpus_slug, 80),
            "outcome": "list_sources",
            "source_count": max(0, int(source_count)),
        }
        self.tool_calls.append(row)
        self._append_source(
            kind=SOURCE_TOOL_READ,
            label=f"read_knowledge_source → {corpus_label or corpus_slug} (list)",
            preview=f"{source_count} source(s)",
            meta=row,
        )

    def record_tool_combined(
        self,
        *,
        query: str,
        hit_count: int,
        outcome: str,
    ) -> None:
        if tls_get_trace() is not self or self.skipped_reason:
            return
        if len(self.tool_calls) >= MAX_TOOL_CALLS:
            return
        row: dict[str, Any] = {
            "tool": "search_memory_and_knowledge_corpora",
            "query_preview": _cap_query(query),
            "hit_count": max(0, int(hit_count)),
            "outcome": cap_str(outcome, 40),
        }
        self.tool_calls.append(row)
        self._append_source(
            kind=SOURCE_TOOL_COMBINED,
            label="search_memory_and_knowledge_corpora",
            preview=_cap_query(query),
            meta=row,
        )

    def to_payload_dict(self) -> dict[str, Any]:
        headline = build_influence_headline(self)
        out: dict[str, Any] = {
            "type": "knowledge_turn_summary",
            "schema_version": SCHEMA_VERSION,
            "turn_id": self.turn_id,
            "phase": "complete",
            "skipped_reason": self.skipped_reason,
            "coverage_notes": list(self.coverage_notes),
            "injection": dict(self.injection),
            "catalog": dict(self.catalog),
            "tool_calls": list(self.tool_calls),
            "sources": list(self.sources),
            "influence_headline": headline,
        }
        if self.skipped_reason:
            out["injection"] = {"path": "absent", "chars_injected": 0, "hits": []}
            out["tool_calls"] = []
            out["sources"] = []

        raw = json.dumps(out, separators=(",", ":"), default=str).encode("utf-8")
        if len(raw) > MAX_PAYLOAD_BYTES:
            out["truncated"] = True
            inj = out.get("injection")
            if isinstance(inj, dict) and isinstance(inj.get("hits"), list):
                inj["hits"] = inj["hits"][:4]
            out["tool_calls"] = out["tool_calls"][:5]
            for tc in out["tool_calls"]:
                if isinstance(tc.get("hits"), list):
                    tc["hits"] = tc["hits"][:3]
            out["sources"] = out["sources"][:10]
            raw2 = json.dumps(out, separators=(",", ":"), default=str).encode("utf-8")
            if len(raw2) > MAX_PAYLOAD_BYTES:
                out["injection"] = {"path": out["injection"].get("path", "absent"), "chars_injected": 0, "hits": []}
                out["tool_calls"] = out["tool_calls"][:3]
                out["sources"] = out["sources"][:6]
        return out


def knowledge_turn_trace_begin(thread_id: str, summary_ref: list[Any] | None) -> None:
    if summary_ref is None:
        return
    tls_clear_trace()
    tls_set_trace(KnowledgeTurnTrace(thread_id))


def _empty_payload_v1(*, skipped_reason: str | None = None, error: str | None = None) -> dict[str, Any]:
    out: dict[str, Any] = {
        "type": "knowledge_turn_summary",
        "schema_version": SCHEMA_VERSION,
        "phase": "complete",
        "skipped_reason": skipped_reason,
        "injection": {"path": "absent", "chars_injected": 0, "hits": []},
        "catalog": {"present": False, "corpus_count": 0, "chars": 0, "corpus_labels": []},
        "tool_calls": [],
        "sources": [],
        "influence_headline": "No knowledge influence this turn.",
        "coverage_notes": list(_DEFAULT_COVERAGE_NOTES),
    }
    if error:
        out["error"] = error
    return out


def knowledge_turn_trace_finalize(summary_ref: list[Any] | None) -> None:
    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_v1(skipped_reason="internal_no_trace")
        else:
            summary_ref[0] = tr.to_payload_dict()
    except Exception:
        summary_ref[0] = _empty_payload_v1(
            skipped_reason=None,
            error=cap_str(redact_text("serialize_failed"), 120),
        )
    finally:
        tls_clear_trace()
