from __future__ import annotations

import math
import uuid
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any

from niannian.memory.adult_redaction import AdultRedactor
from niannian.memory.disclosure import RecallContext, atom_visible_for_recall
from niannian.memory.types import RecallPack
from niannian.personas import BRAINSTEM, CEREBELLUM, HIPPOCAMPUS, LEFT
from niannian.schemas.models import AssistResult, Atom, QueryResult, WikiHit

if TYPE_CHECKING:
    from niannian.brain.cerebellum.rerank import CerebellumRerank, RerankCandidate


def _belief_weight(atom: Atom) -> float:
    status = str((atom.meta or {}).get("belief_status") or "active").lower()
    if status == "retracted":
        return 0.1
    if status == "superseded":
        return 0.3
    return 1.0


def _atom_kind(atom: Atom) -> str:
    return str(atom.meta.get("kind") or "note")


def _atom_scope(atom: Atom) -> dict[str, Any]:
    scope = atom.meta.get("scope")
    return dict(scope) if isinstance(scope, dict) else {}


def _atom_thread(atom: Atom) -> str | None:
    rel = atom.meta.get("relations") or {}
    if isinstance(rel, dict) and rel.get("thread_id"):
        return str(rel["thread_id"])
    scope = _atom_scope(atom)
    if scope.get("loop_id") and scope.get("node_id"):
        return f"{scope['loop_id']}:{scope['node_id']}"
    return None


def _parse_dt(value: str) -> datetime:
    try:
        dt = datetime.fromisoformat(value)
    except ValueError:
        return datetime.now(timezone.utc)
    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=timezone.utc)
    return dt


def _vault_meta(atom: Atom) -> dict[str, Any] | None:
    vault = atom.meta.get("vault")
    if not isinstance(vault, dict):
        return None
    path = vault.get("path")
    vid = vault.get("vault_id")
    if path and vid:
        return vault
    return None


def _text_match(task: str, text: str) -> float:
    if not task or not text:
        return 0.0
    q = task.lower()
    t = text.lower()
    hits = sum(1 for token in q.split() if len(token) > 1 and token in t)
    return min(1.0, hits / max(1, len(q.split())))


def _tag_overlap(task: str, atom: Atom) -> float:
    tags = atom.tags or []
    if not tags:
        return 0.0
    q_tokens = {w for w in task.lower().split() if len(w) > 1}
    tag_tokens = {t.lower() for t in tags}
    if not q_tokens:
        return 0.0
    hits = sum(1 for t in tag_tokens if t in q_tokens or any(t in q or q in t for q in q_tokens))
    return min(1.0, hits / max(1, len(tag_tokens)))


_SCOPE_TIER_SCORES = {"node": 1.0, "loop": 0.85, "project": 0.55, "global": 0.35, "none": 0.1}


def _atom_scope_tier(atom: Atom, req_scope: dict[str, Any]) -> str:
    atom_scope = _atom_scope(atom)
    req_loop = req_scope.get("loop_id")
    req_node = req_scope.get("node_id")
    req_tool = req_scope.get("tool")

    if req_node and atom_scope.get("node_id") == req_node:
        if not req_loop or atom_scope.get("loop_id") == req_loop:
            return "node"
    if req_loop and atom_scope.get("loop_id") == req_loop:
        return "loop"
    if req_tool and atom.project == req_tool:
        return "project"
    if atom.project == "global":
        return "global"
    return "none"


def _scope_ladder_score(atom: Atom, req_scope: dict[str, Any], policy_cfg: dict[str, Any]) -> float:
    tier = _atom_scope_tier(atom, req_scope)
    expand = policy_cfg.get("scope_expand") or ["node", "loop", "project", "global"]
    if tier not in expand:
        return _SCOPE_TIER_SCORES["none"]
    return _SCOPE_TIER_SCORES[tier]


def _thread_summary_keys(store_summaries: dict[str, Any]) -> set[str]:
    out: set[str] = set()
    for key in store_summaries:
        if str(key).startswith("thread:"):
            out.add(str(key).split("thread:", 1)[-1])
    return out

class RecallEngine:
    """③ 回忆 Recall — policy 驱动检索 + RecallPack。"""

    def __init__(
        self,
        left,
        cerebellum,
        hippocampus,
        right,
        config: dict[str, Any],
        *,
        rerank: CerebellumRerank | None = None,
        vault_store=None,
        embed_recall=None,
    ) -> None:
        self._left = left
        self._cerebellum = cerebellum
        self._hippocampus = hippocampus
        self._right = right
        self._config = config
        self._rerank = rerank
        self._vault_store = vault_store
        self._embed_recall = embed_recall
        recall_cfg = config.get("recall", {})
        self._default_policy = recall_cfg.get("default_policy", "assist_default")
        self._half_life = float(recall_cfg.get("half_life_days", 14))
        self._policies = recall_cfg.get("policies", {})
        self._rerank_defaults = recall_cfg.get("rerank") or {}
        self._redactor = AdultRedactor.from_app_config(config)

    def _decode_on_export(self) -> bool:
        cfg = self._config.get("adult_redaction") or {}
        return bool(cfg.get("decode_on_export", True))

    def _export_text(self, text: str, *, project: str, for_llm: bool) -> str:
        if for_llm or not self._decode_on_export():
            return text
        return self._redactor.decode_for_display(text, project=project)

    def recall(
        self,
        task: str,
        *,
        context: str = "",
        scope: dict[str, Any] | None = None,
        policy: str | None = None,
        project: str | None = None,
        return_format: str = "pack",
        filters: dict[str, Any] | None = None,
        recall_context: RecallContext | None = None,
    ) -> RecallPack | str | list[dict[str, Any]]:
        policy_name = policy or self._default_policy
        req_scope = dict(scope or {})
        proj = project or req_scope.get("tool") or "global"
        trace: list[str] = []
        ctx = recall_context or RecallContext(mode="subject")

        policy_cfg = dict(self._policies.get(policy_name, {}))
        if filters:
            policy_cfg = {**policy_cfg, **{k: v for k, v in filters.items() if v is not None}}

        atoms = self._left.list_atoms(project=proj, limit=2000)
        if proj != "global":
            atoms += [a for a in self._left.list_atoms(project="global", limit=500) if a not in atoms]
        atoms = [a for a in atoms if atom_visible_for_recall(a, ctx)]
        trace.append(f"[{LEFT}] disclosure scope={ctx.mode} object={ctx.object_id or '-'} atoms={len(atoms)}")

        candidates = self._filter_atoms(atoms, policy_name, req_scope, policy_cfg, task)
        scored = self._score_atoms(candidates, task, req_scope, policy_cfg)
        scored.sort(key=lambda x: x[1], reverse=True)
        scored = self._maybe_embed_scored_atoms(task, scored, policy_cfg, trace)
        scored = self._maybe_rerank_scored_atoms(task, scored, policy_cfg, trace)

        max_atoms = int(policy_cfg.get("max_atoms", 20))
        if self._layer_enabled("atom", policy_cfg):
            selected = [a for a, _ in scored[:max_atoms]]
        else:
            selected = []
            trace.append(f"[{LEFT}] atoms skipped: policy excludes atom layer")

        if self._layer_enabled("wiki", policy_cfg):
            wiki_hits = self._cerebellum.recall(
                f"{task}\n{context}",
                project=proj,
                limit=int(policy_cfg.get("max_wiki", 5)),
                namespace=ctx.mode,
                object_id=ctx.object_id if ctx.mode == "object" else None,
            )
            bfs_ids = [h.page_id for h in wiki_hits if h.source == "bfs" and h.page_id]
            if bfs_ids:
                trace.append(f"[{CEREBELLUM}] Wiki bfs: {','.join(bfs_ids[:8])}")
            trace.append(f"[{CEREBELLUM}] Wiki {len(wiki_hits)}")
        else:
            wiki_hits = []
            trace.append(f"[{CEREBELLUM}] Wiki skipped: policy excludes wiki layer")

        if self._layer_enabled("summary", policy_cfg):
            summaries = self._collect_summaries(task, proj, req_scope, policy_cfg, selected)
            trace.append(f"[{HIPPOCAMPUS}] summaries {len(summaries)}")
        else:
            summaries = []
            trace.append(f"[{HIPPOCAMPUS}] summaries skipped: policy excludes summary layer")

        trace.append(f"[{LEFT}] atoms {len(atoms)} → {len(selected)} after policy")

        if self._layer_enabled("thread", policy_cfg):
            threads = self._build_threads(selected, req_scope, policy_cfg)
        else:
            threads = []
            trace.append(f"[{LEFT}] threads skipped: policy excludes thread layer")
        vault_hits = self._collect_vault_hits(task, atoms, policy_cfg, trace, recall_context=ctx)
        wiki_hits, vault_hits = self._merge_vault_wiki_rerank_candidates(
            task,
            wiki_hits,
            vault_hits,
            policy_cfg,
            trace,
        )
        vault_hits = self._maybe_agentic_vault_search(task, vault_hits, policy_cfg, trace)
        citations = [
            {
                "atom_id": a.id,
                "kind": _atom_kind(a),
                "snippet": a.raw_text[:120],
                "reason": "scope_match",
            }
            for a in selected[:10]
        ]

        context_block = self._build_context_block(
            task,
            selected,
            wiki_hits,
            summaries,
            threads,
            context,
            policy_cfg,
            vault_hits=vault_hits,
        )

        pack = RecallPack(
            pack_id=f"rp-{uuid.uuid4().hex[:12]}",
            policy=policy_name,
            context_block=context_block,
            layers={
                "summaries": summaries,
                "threads": threads,
                "atoms": [
                    {
                        "id": a.id,
                        "kind": _atom_kind(a),
                        "snippet": a.raw_text[:160],
                        "scope": _atom_scope(a),
                    }
                    for a in selected
                ],
                "wiki": [self._wiki_dict(h) for h in wiki_hits],
                "vault": vault_hits,
            },
            citations=citations,
            recall_trace=trace,
        )
        self._maybe_sample_recall_trace(pack, policy_name)

        if return_format == "context_only":
            return pack.context_block
        if return_format == "atoms":
            return pack.layers.get("atoms", [])
        if return_format == "minimal":
            pack.layers = {
                "summaries": summaries,
                "threads": [t for t in threads if t.get("atom_ids")],
            }
            return pack
        return pack

    def _filter_atoms(
        self,
        atoms: list[Atom],
        policy: str,
        req_scope: dict[str, Any],
        policy_cfg: dict[str, Any],
        task: str,
    ) -> list[Atom]:
        allowed_kinds = set(policy_cfg.get("kinds") or [])
        exclude_kinds = set(policy_cfg.get("exclude_kinds") or [])
        req_loop = req_scope.get("loop_id")
        req_node = req_scope.get("node_id")
        req_thread = f"{req_loop}:{req_node}" if req_loop and req_node else None

        out: list[Atom] = []
        for atom in atoms:
            kind = _atom_kind(atom)
            if allowed_kinds and kind not in allowed_kinds:
                continue
            if kind in exclude_kinds:
                continue
            atom_scope = _atom_scope(atom)
            atom_loop = atom_scope.get("loop_id")
            atom_node = atom_scope.get("node_id")
            thread = _atom_thread(atom)

            if policy == "loop_continue":
                if kind == "preference":
                    out.append(atom)
                    continue
                if not req_loop:
                    out.append(atom)
                    continue
                if kind in ("adopt", "review") and atom_loop == req_loop:
                    out.append(atom)
                    continue
                if kind == "write":
                    if atom_loop != req_loop:
                        continue
                    if req_node and atom_node == req_node:
                        out.append(atom)
                    elif atom_loop == req_loop and not req_node:
                        out.append(atom)
                    continue
                if kind == "fact" and atom.project in (req_scope.get("tool"), "global"):
                    out.append(atom)
                continue

            if policy == "node_retry":
                if req_thread and thread == req_thread and kind in ("review", "reject", "write"):
                    out.append(atom)
                continue

            if policy == "audit_gate":
                if kind == "adopt" and req_loop and atom_loop == req_loop:
                    out.append(atom)
                    continue
                if req_thread and thread == req_thread and kind in ("write", "review"):
                    out.append(atom)
                continue

            if policy in ("assist_default", "memory_query", "global_habits", "tool_resume"):
                if policy == "global_habits" and kind not in ("preference", "behavior"):
                    continue
                if policy == "assist_default":
                    if kind in ("preference", "fact", "adopt", "note", "behavior"):
                        out.append(atom)
                    elif _text_match(task, atom.raw_text) > 0:
                        out.append(atom)
                    continue
                out.append(atom)
                continue

            out.append(atom)

        if policy == "loop_continue" and req_loop:
            reviews = [a for a in out if _atom_kind(a) == "review" and _atom_scope(a).get("loop_id") == req_loop]
            reviews.sort(key=lambda a: a.created_at, reverse=True)
            review_ids = {a.id for a in reviews[:5]}
            filtered: list[Atom] = []
            for a in out:
                if _atom_kind(a) == "review" and a.id not in review_ids:
                    continue
                filtered.append(a)
            out = filtered

        if policy == "node_retry" and req_thread:
            writes = [a for a in out if _atom_kind(a) == "write" and _atom_thread(a) == req_thread]
            if writes:
                latest_write = max(writes, key=lambda a: a.created_at)
                out = [a for a in out if _atom_kind(a) != "write" or a.id == latest_write.id]

        custom_kinds = policy_cfg.get("kinds")
        if custom_kinds and policy not in ("loop_continue", "node_retry", "audit_gate"):
            out = [a for a in out if _atom_kind(a) in custom_kinds]

        if policy_cfg.get("loop_id"):
            lid = policy_cfg["loop_id"]
            out = [a for a in out if _atom_scope(a).get("loop_id") == lid]

        return out

    def _score_atoms(
        self,
        atoms: list[Atom],
        task: str,
        req_scope: dict[str, Any],
        policy_cfg: dict[str, Any],
    ) -> list[tuple[Atom, float]]:
        kind_weights = policy_cfg.get("kind_weights") or {}
        kinds_list = list(policy_cfg.get("kinds") or [])
        half_life = float(policy_cfg.get("half_life_days", self._half_life))
        suppress_raw = bool(policy_cfg.get("suppress_raw_when_summary", False))
        summary_threads: set[str] = set()
        if suppress_raw:
            summary_threads = _thread_summary_keys(self._hippocampus.store.read_summaries())

        now = datetime.now(timezone.utc)
        scored: list[tuple[Atom, float]] = []

        for atom in atoms:
            kind = _atom_kind(atom)
            if kind in kind_weights:
                w_kind = float(kind_weights[kind])
            elif kind in kinds_list:
                w_kind = max(0.4, 1.0 - kinds_list.index(kind) * 0.08)
            else:
                w_kind = 0.5
            if kind == "reject":
                w_kind = max(w_kind, 1.0)

            scope_score = _scope_ladder_score(atom, req_scope, policy_cfg)
            age_days = max(0.0, (now - _parse_dt(atom.created_at)).total_seconds() / 86400)
            recency = math.exp(-age_days / max(1.0, half_life))
            text_score = _text_match(task, atom.raw_text)
            tag_score = _tag_overlap(task, atom)
            status_w = 0.4 if atom.status == "raw" else 0.8

            score = (
                scope_score * 0.38
                + w_kind * 0.28
                + recency * 0.14
                + text_score * 0.1
                + tag_score * 0.05
                + status_w * 0.05
            )

            thread = _atom_thread(atom)
            if suppress_raw and thread and thread in summary_threads:
                score *= 0.25

            score *= _belief_weight(atom)
            scored.append((atom, score))
        return scored

    def _embed_policy_cfg(self, policy_cfg: dict[str, Any]) -> dict[str, Any]:
        emb_root = self._config.get("recall", {}).get("embedding") or {}
        policy_emb = policy_cfg.get("embedding") if isinstance(policy_cfg.get("embedding"), dict) else {}
        merged = dict(emb_root)
        merged.update(policy_emb)
        return merged

    def _maybe_embed_scored_atoms(
        self,
        task: str,
        scored: list[tuple[Atom, float]],
        policy_cfg: dict[str, Any],
        trace: list[str],
    ) -> list[tuple[Atom, float]]:
        emb_cfg = self._embed_policy_cfg(policy_cfg)
        if not emb_cfg.get("enabled"):
            return scored
        if self._embed_recall is None:
            trace.append(f"[{CEREBELLUM}] embed skipped: 未注入 CerebellumEmbedRecall")
            return scored
        blend = emb_cfg.get("blend") if isinstance(emb_cfg.get("blend"), dict) else {}
        return self._embed_recall.blend_scored_atoms(
            task,
            scored,
            max_coarse=int(emb_cfg.get("max_coarse", 40)),
            w_embed=float(blend.get("w_embed", 0.35)),
            w_rule=float(blend.get("w_rule", 0.65)),
            trace=trace,
        )

    def _rerank_policy_cfg(self, policy_cfg: dict[str, Any]) -> dict[str, Any]:
        policy_rr = policy_cfg.get("rerank") if isinstance(policy_cfg.get("rerank"), dict) else {}
        merged = dict(self._rerank_defaults)
        merged.update(policy_rr)
        blend = dict((self._rerank_defaults.get("blend") or {}))
        blend.update((policy_rr.get("blend") or {}))
        merged["blend"] = blend
        return merged

    def _maybe_rerank_scored_atoms(
        self,
        task: str,
        scored: list[tuple[Atom, float]],
        policy_cfg: dict[str, Any],
        trace: list[str],
    ) -> list[tuple[Atom, float]]:
        rr_cfg = self._rerank_policy_cfg(policy_cfg)
        if not rr_cfg.get("enabled"):
            return scored
        if self._rerank is None:
            trace.append(f"[{CEREBELLUM}] rerank skipped: 未注入 CerebellumRerank")
            return scored

        min_candidates = int(rr_cfg.get("min_candidates", 3))
        max_candidates = int(rr_cfg.get("max_candidates", 80))
        pool = scored[:max_candidates]
        if len(pool) < min_candidates:
            trace.append(f"[{CEREBELLUM}] rerank skipped: candidates<{min_candidates}")
            return scored

        from niannian.brain.cerebellum.rerank import RerankCandidate

        candidates = [
            RerankCandidate(
                id=a.id,
                channel="A",
                text=(a.raw_text or "")[:512],
                meta={"rule_score": rule_score},
            )
            for a, rule_score in pool
        ]
        ranked = self._rerank.rerank_candidates(
            task,
            candidates,
            top_n=len(pool),
            trace=trace,
        )
        if not ranked:
            return scored

        rerank_by_id = {c.id: score for c, score in ranked}
        max_rr = max(rerank_by_id.values()) if rerank_by_id else 1.0
        max_rule = max(rule for _, rule in pool) or 1.0
        blend = rr_cfg.get("blend") or {}
        w_rerank = float(blend.get("w_rerank", 0.75))
        w_rule = float(blend.get("w_rule", 0.25))

        def _blended(atom: Atom, rule_score: float) -> float:
            rr = rerank_by_id.get(atom.id, 0.0)
            rr_norm = rr / max_rr if max_rr else 0.0
            rule_norm = rule_score / max_rule if max_rule else 0.0
            return w_rerank * rr_norm + w_rule * rule_norm

        reranked_pool = [(a, _blended(a, rule)) for a, rule in pool]
        reranked_pool.sort(key=lambda x: x[1], reverse=True)
        pool_ids = {a.id for a, _ in pool}
        tail = [(a, s) for a, s in scored if a.id not in pool_ids]
        trace.append(f"[{CEREBELLUM}] rerank {len(pool)} → {len(reranked_pool)}")
        return reranked_pool + tail

    @staticmethod
    def _snippet_dedup_key(text: str, *, n: int = 64) -> str:
        collapsed = " ".join(str(text or "").lower().split())
        return collapsed[:n]

    def _merge_vault_wiki_rerank_candidates(
        self,
        task: str,
        wiki_hits: list[WikiHit],
        vault_hits: list[dict[str, Any]],
        policy_cfg: dict[str, Any],
        trace: list[str],
    ) -> tuple[list[WikiHit], list[dict[str, Any]]]:
        """vault chunk 与 wiki page 同一 rerank 批次，Assembly 去重相似段落。"""
        rr_cfg = self._rerank_policy_cfg(policy_cfg)
        if not rr_cfg.get("enabled"):
            return wiki_hits, vault_hits
        if self._rerank is None:
            return wiki_hits, vault_hits
        if not wiki_hits and not vault_hits:
            return wiki_hits, vault_hits

        from niannian.brain.cerebellum.rerank import RerankCandidate

        candidates: list[RerankCandidate] = []
        for idx, hit in enumerate(wiki_hits):
            text = f"{hit.concept} {hit.summary}".strip()
            candidates.append(
                RerankCandidate(
                    id=f"wiki:{hit.page_id or hit.concept or idx}",
                    channel="wiki",
                    text=text[:512],
                    meta={"kind": "wiki", "index": idx},
                )
            )
        for idx, hit in enumerate(vault_hits):
            snippet = str(hit.get("snippet") or hit.get("title") or "")
            candidates.append(
                RerankCandidate(
                    id=f"vault:{hit.get('vault_id', idx)}",
                    channel="vault",
                    text=snippet[:512],
                    meta={"kind": "vault", "index": idx},
                )
            )

        top_n = int(rr_cfg.get("top_n", len(candidates)))
        ranked = self._rerank.rerank_candidates(task, candidates, top_n=top_n, trace=trace)
        if not ranked:
            return wiki_hits, vault_hits

        out_wiki: list[WikiHit] = []
        out_vault: list[dict[str, Any]] = []
        seen_keys: set[str] = set()
        for cand, _score in ranked:
            key = self._snippet_dedup_key(cand.text)
            if not key:
                continue
            if key in seen_keys:
                continue
            seen_keys.add(key)
            if cand.channel == "wiki":
                out_wiki.append(wiki_hits[int(cand.meta["index"])])
            elif cand.channel == "vault":
                out_vault.append(vault_hits[int(cand.meta["index"])])

        trace.append(
            f"[{CEREBELLUM}] V+W rerank {len(candidates)} → wiki {len(out_wiki)} vault {len(out_vault)}"
        )
        return out_wiki, out_vault

    def _layer_enabled(self, layer: str, policy_cfg: dict[str, Any]) -> bool:
        exclude = set(policy_cfg.get("exclude_layers") or [])
        if layer in exclude:
            return False
        layers = policy_cfg.get("layers")
        if isinstance(layers, list) and layers:
            return layer in layers
        return True

    def _layers_vault_enabled(self, policy_cfg: dict[str, Any]) -> bool:
        return self._layer_enabled("vault", policy_cfg)

    def _layers_wiki_enabled(self, policy_cfg: dict[str, Any]) -> bool:
        return self._layer_enabled("wiki", policy_cfg)

    def _collect_vault_hits(
        self,
        task: str,
        atoms: list[Atom],
        policy_cfg: dict[str, Any],
        trace: list[str],
        *,
        recall_context: RecallContext | None = None,
    ) -> list[dict[str, Any]]:
        if not self._layers_vault_enabled(policy_cfg):
            trace.append("[vault] skipped: policy excludes vault layer")
            return []

        pointers = [a for a in atoms if _vault_meta(a)]
        if not pointers:
            trace.append("[vault] pointers 0 → 0 snippets")
            return []

        max_vault = int(policy_cfg.get("max_vault", 5))
        scored: list[tuple[Atom, float]] = []
        for atom in pointers:
            text_score = _text_match(task, atom.raw_text)
            tag_score = _tag_overlap(task, atom)
            score = text_score + tag_score * 0.5
            if task.strip() and score <= 0:
                continue
            scored.append((atom, score if score > 0 else 0.35))

        scored.sort(key=lambda x: x[1], reverse=True)
        pool = scored if task.strip() else [(a, 0.35) for a in pointers]

        hits: list[dict[str, Any]] = []
        overview_hit = self._try_overview_hit(task, trace)
        if overview_hit:
            hits.append(overview_hit)

        for atom, rule_score in pool[:max_vault]:
            vm = _vault_meta(atom)
            if vm is None:
                continue
            rel_path = str(vm["path"])
            vault_id = str(vm["vault_id"])
            title = atom.raw_text[:80]
            snippet = atom.raw_text[:120]
            score = max(75.0, round(rule_score * 100, 1))

            if self._vault_store is not None:
                try:
                    from niannian.vault.chunking import best_chunk_for_query, split_md_by_heading

                    doc = self._vault_store.read_document(rel_path)
                    fm = doc.get("frontmatter") or {}
                    body = str(doc.get("body") or "")
                    title = str(fm.get("title") or title)
                    chunks = split_md_by_heading(body)
                    if chunks:
                        embed_mode = str(((self._config.get("vault") or {}).get("embed") or {}).get("mode") or "off")
                        picked = None
                        if embed_mode == "lazy" and self._embed_recall is not None:
                            from niannian.vault.embed_index import VaultEmbedIndex

                            index = VaultEmbedIndex(self._vault_store._root)  # noqa: SLF001
                            ranked = self._embed_recall.rank_vault_chunks(
                                task,
                                chunks,
                                vault_id=vault_id,
                                index=index,
                                trace=trace,
                            )
                            if ranked:
                                picked = ranked[0][0]
                        if picked is None:
                            picked = best_chunk_for_query(chunks, task)
                        if picked:
                            heading = str(picked.get("heading") or "")
                            chunk_text = str(picked.get("text") or "")
                            snippet = f"## {heading}\n{chunk_text}" if heading else chunk_text
                            snippet = snippet[:400]
                    elif body:
                        snippet = body[:400]
                except (FileNotFoundError, OSError):
                    pass

            hits.append(
                {
                    "vault_id": vault_id,
                    "title": title,
                    "rel_path": rel_path,
                    "snippet": snippet,
                    "score": score,
                }
            )

        trace.append(f"[vault] pointers {len(pointers)} → {len(hits)} snippets")
        return self._filter_guest_vault_hits(hits, atoms)

    def _filter_guest_vault_hits(self, hits: list[dict[str, Any]], atoms: list[Atom]) -> list[dict[str, Any]]:
        guest = any(
            str((a.meta.get("provenance") or {}).get("trust_tier") or "") == "guest_reference"
            for a in atoms
            if _vault_meta(a)
        )
        if not guest:
            return hits
        return [
            {**h, "rel_path": "", "snippet": str(h.get("snippet", ""))[:120]}
            for h in hits
            if not str(h.get("rel_path") or "").startswith("user/")
        ]

    def _maybe_agentic_vault_search(
        self,
        task: str,
        vault_hits: list[dict[str, Any]],
        policy_cfg: dict[str, Any],
        trace: list[str],
    ) -> list[dict[str, Any]]:
        if not policy_cfg.get("agentic_vault_search"):
            return vault_hits
        if self._vault_store is None or not task.strip():
            return vault_hits
        from niannian.vault.search import VaultSearchService

        extra = VaultSearchService(self._vault_store).search(task, limit=3)
        if not extra:
            trace.append("[vault] agentic search 0")
            return vault_hits
        trace.append(f"[vault] agentic search {len(extra)}")
        seen = {str(h.get("vault_id")) for h in vault_hits}
        out = list(vault_hits)
        for row in extra:
            vid = str(row.get("vault_id") or "")
            if vid and vid not in seen:
                out.append(
                    {
                        "vault_id": vid,
                        "title": row.get("title"),
                        "rel_path": row.get("rel_path"),
                        "snippet": row.get("snippet"),
                        "score": 70.0,
                        "source": "agentic",
                    }
                )
                seen.add(vid)
        return out

    def _maybe_sample_recall_trace(self, pack: RecallPack, policy: str) -> None:
        rate = float((self._config.get("recall") or {}).get("trace_sample_rate") or 0)
        if rate <= 0:
            return
        store = getattr(self._left, "_store", None)
        home = getattr(store, "_home", None) if store is not None else None
        if home is None:
            home = getattr(self._left, "_home", None) or getattr(self._left, "home", None)
        if home is None:
            return
        from pathlib import Path

        from niannian.memory.recall_trace import maybe_write_recall_trace

        maybe_write_recall_trace(Path(home), pack=pack, policy=policy, sample_rate=rate)

    def _try_overview_hit(self, task: str, trace: list[str]) -> dict[str, Any] | None:
        if self._vault_store is None or not task.strip():
            return None
        from niannian.vault.overview import overview_rel_path

        q = task.lower()
        seen_subs: set[str] = set()
        for entry in self._vault_store.load_index().get("entries") or []:
            if not isinstance(entry, dict):
                continue
            sub = str(entry.get("subproject") or "")
            cat = str(entry.get("category") or "")
            ns = str(entry.get("namespace") or "user")
            key = f"{ns}:{cat}:{sub}"
            if not sub or key in seen_subs or sub.lower() not in q:
                continue
            seen_subs.add(key)
            rel = overview_rel_path(namespace=ns, category=cat, subproject=sub)
            try:
                doc = self._vault_store.read_document(rel)
            except (FileNotFoundError, OSError):
                continue
            body = str(doc.get("body") or "")
            fm = doc.get("frontmatter") or {}
            trace.append(f"[vault] overview hit {sub}")
            return {
                "vault_id": f"overview:{ns}:{cat}:{sub}",
                "title": str(fm.get("title") or f"{sub} 总览"),
                "rel_path": rel,
                "snippet": body[:400],
                "score": 88.0,
                "source": "overview",
            }
        return None

    def _collect_summaries(
        self,
        task: str,
        project: str,
        req_scope: dict[str, Any],
        policy_cfg: dict[str, Any],
        selected: list[Atom],
    ) -> list[dict[str, Any]]:
        prefer_thread = bool(policy_cfg.get("prefer_thread_summary", True))
        all_summaries = self._hippocampus.store.read_summaries()
        thread_out: list[dict[str, Any]] = []
        tag_out: list[dict[str, Any]] = []

        req_loop = req_scope.get("loop_id")
        req_node = req_scope.get("node_id")
        req_thread = f"{req_loop}:{req_node}" if req_loop and req_node else None

        for key, item in all_summaries.items():
            text = str(item.get("text", "")).strip()
            if not text:
                continue
            if str(key).startswith("thread:"):
                thread_id = str(key).split("thread:", 1)[-1]
                if req_loop and not thread_id.startswith(str(req_loop)):
                    continue
                priority = 0 if req_thread and thread_id == req_thread else 1
                thread_out.append(
                    {
                        "text": text,
                        "scope": thread_id,
                        "source": "hippocampus_thread",
                        "priority": priority,
                    }
                )
            else:
                tag_out.append(
                    {
                        "text": f"【{key}】{text}",
                        "scope": project,
                        "source": "hippocampus",
                        "priority": 2,
                    }
                )

        if not tag_out:
            for s in self._hippocampus.get_summaries(task, limit=5):
                tag_out.append(
                    {"text": s, "scope": project, "source": "hippocampus", "priority": 2}
                )

        thread_out.sort(key=lambda s: s.get("priority", 9))
        combined = (thread_out + tag_out) if prefer_thread else (tag_out + thread_out)
        return combined[:8]

    def _build_threads(
        self,
        selected: list[Atom],
        req_scope: dict[str, Any],
        policy_cfg: dict[str, Any],
    ) -> list[dict[str, Any]]:
        by_thread: dict[str, list[Atom]] = {}
        for atom in selected:
            tid = _atom_thread(atom)
            if tid:
                by_thread.setdefault(tid, []).append(atom)

        threads: list[dict[str, Any]] = []
        for tid, group in by_thread.items():
            adopts = [a for a in group if _atom_kind(a) == "adopt"]
            reviews = [a for a in group if _atom_kind(a) == "review"]
            narrative_parts = [a.raw_text[:100] for a in adopts[:1] + reviews[:2]]
            threads.append(
                {
                    "thread_id": tid,
                    "narrative": "；".join(narrative_parts) or tid,
                    "atom_ids": [a.id for a in group],
                }
            )
        return threads

    def _build_context_block(
        self,
        task: str,
        atoms: list[Atom],
        wiki_hits: list[WikiHit],
        summaries: list[dict[str, Any]],
        threads: list[dict[str, Any]],
        context: str,
        policy_cfg: dict[str, Any],
        *,
        vault_hits: list[dict[str, Any]] | None = None,
    ) -> str:
        suppress_raw = bool(policy_cfg.get("suppress_raw_when_summary", False))
        summary_threads = _thread_summary_keys(self._hippocampus.store.read_summaries()) if suppress_raw else set()

        thread_summaries = [s for s in summaries if s.get("source") == "hippocampus_thread"]
        tag_texts = [s.get("text", "") for s in summaries if s.get("source") != "hippocampus_thread" and s.get("text")]

        display_atoms = list(atoms)
        if suppress_raw and summary_threads:
            display_atoms = [
                a
                for a in atoms
                if not (
                    _atom_thread(a)
                    and _atom_thread(a) in summary_threads
                    and _atom_kind(a) in ("write", "review")
                )
            ]

        block = self._right.build_context_block(
            task,
            wiki_hits=wiki_hits,
            summaries=tag_texts,
            atoms=display_atoms,
        )

        if thread_summaries:
            thread_lines = ["### 线程提炼（小海 · 优先）"]
            for s in thread_summaries[:5]:
                scope = s.get("scope", "")
                prefix = f"[{scope}] " if scope else ""
                thread_lines.append(f"- {prefix}{s.get('text', '')[:240]}")
            title, sep, rest = block.partition("\n\n")
            block = f"{title}\n\n" + "\n".join(thread_lines) + f"\n\n{rest}" if rest else f"{title}\n\n" + "\n".join(thread_lines)

        if threads and not thread_summaries:
            thread_lines = ["", "### 线程摘要（小海）"]
            for t in threads[:5]:
                thread_lines.append(f"- [{t['thread_id']}] {t.get('narrative', '')[:160]}")
            block = f"{block}\n" + "\n".join(thread_lines)

        if display_atoms:
            mrrr_lines = ["", "### 结构化碎片（MRRR）"]
            for a in display_atoms[:8]:
                kind = _atom_kind(a)
                scope = _atom_scope(a)
                node = scope.get("node_id", "")
                limit = 80 if suppress_raw else 120
                mrrr_lines.append(f"- [{kind}@{node}] {a.raw_text[:limit]}")
            block = f"{block}\n" + "\n".join(mrrr_lines)

        if vault_hits:
            vault_lines = ["", "### 资料库摘要（vault）"]
            for v in vault_hits[:5]:
                title = v.get("title", "")
                vault_lines.append(f"- [{title}] {str(v.get('snippet', ''))[:200]}")
            block = f"{block}\n" + "\n".join(vault_lines)

        if context:
            block = f"{block}\n\n### 宿主附加上下文\n{context.strip()}"
        return block.strip()

    @staticmethod
    def _wiki_dict(hit: WikiHit) -> dict[str, Any]:
        return {
            "concept": hit.concept,
            "summary": hit.summary,
            "links": hit.links,
        }

    def recall_wiki(self, query: str, *, project: str | None = None, limit: int = 5) -> list[dict[str, Any]]:
        hits = self._cerebellum.recall(query, project=project or "global", limit=limit)
        return [
            {"concept": h.concept, "summary": h.summary, "links": h.links, "score": h.score}
            for h in hits
        ]

    def query(
        self,
        question: str,
        *,
        project: str | None = None,
        llm=None,
        for_llm: bool = False,
        recall_context: RecallContext | None = None,
    ) -> QueryResult:
        """联想问答 — recall(policy=memory_query) + 小思 answer。"""
        pack = self.recall(question, project=project, policy="memory_query", recall_context=recall_context)
        assert isinstance(pack, RecallPack)
        result = self.to_query_result(pack, question, llm=llm)
        proj = project or "global"
        result.answer = self._export_text(result.answer, project=proj, for_llm=for_llm)
        if not for_llm and self._decode_on_export():
            result.citations = self._redactor.decode_citations(result.citations, project=proj)
        return result

    def assist(
        self,
        task: str,
        *,
        context: str = "",
        project: str | None = None,
        for_llm: bool = False,
        recall_context: RecallContext | None = None,
    ) -> AssistResult:
        """碎片化思考辅助 — recall(policy=assist_default, format=context_only)。"""
        proj = project or "global"
        pack = self.recall(
            task,
            context=context,
            project=proj,
            policy="assist_default",
            recall_context=recall_context,
        )
        assert isinstance(pack, RecallPack)
        data = self.to_assist_result(pack, task)
        block = self._export_text(data["context_block"], project=proj, for_llm=for_llm)
        fragments = [
            self._export_text(str(f), project=proj, for_llm=for_llm)
            for f in data["fragments"]
        ]
        citations = list(pack.citations)
        if not for_llm and self._decode_on_export():
            citations = self._redactor.decode_citations(citations, project=proj)
        return AssistResult(
            context_block=block,
            fragments=fragments,
            wiki_links=data["wiki_links"],
            citations=citations,
            trace=data["trace"],
        )

    def to_query_result(
        self,
        pack: RecallPack,
        question: str,
        *,
        llm=None,
    ) -> QueryResult:
        wiki_hits = [
            WikiHit(
                concept=str(w.get("concept", "")),
                summary=str(w.get("summary", "")),
                links=list(w.get("links") or []),
            )
            for w in (pack.layers.get("wiki") or [])
        ]
        atoms: list[Atom] = []
        for item in pack.layers.get("atoms") or []:
            aid = item.get("id")
            if aid:
                atom = self._left.get(str(aid))
                if atom is not None:
                    atoms.append(atom)

        summaries: list[str] = []
        for s in pack.layers.get("summaries") or []:
            if isinstance(s, dict):
                text = s.get("text", "")
                if text:
                    summaries.append(str(text))
            elif s:
                summaries.append(str(s))

        answer = self._right.answer(
            question,
            wiki_hits=wiki_hits,
            summaries=summaries,
            atoms=atoms,
            llm=llm,
        )
        trace = [f"[{BRAINSTEM}] recall policy={pack.policy}"] + list(pack.recall_trace)
        trace.append(self._right.trace_line("query 完成"))
        return QueryResult(
            answer=answer,
            wiki_hits=wiki_hits,
            summary_snippets=summaries,
            citations=list(pack.citations),
            trace=trace,
        )

    def to_assist_result(self, pack: RecallPack, task: str) -> dict[str, Any]:
        atoms_layer = pack.layers.get("atoms") or []
        wiki = pack.layers.get("wiki") or []
        return {
            "context_block": pack.context_block,
            "fragments": [a.get("snippet", "") for a in atoms_layer[:3]],
            "wiki_links": [w.get("concept", "") for w in wiki],
            "trace": [f"[{BRAINSTEM}] recall policy={pack.policy}"] + pack.recall_trace,
        }
