#!/usr/bin/env python3
"""echo_recall.py — H1: hybrid lexical (BM25) + graph recall.  [v1.0 — wired]

Replaces the keyword-only recall (search/simple + fixed one-hop) with a ranked, local
retriever that finds a note even when the query is phrased differently than the note
was stored, then expands along the graph with a per-hop decay so the *connected*
context is surfaced in relevance order.

DESIGN (bounded by the project's pure-stdlib, no-deps constraint):
  * Lexical layer  : BM25 over note bodies. No embeddings / external API — BM25 is a
                     strong, dependency-free baseline that handles vocabulary mismatch
                     far better than substring search.
  * Graph layer    : from the top lexical hits, walk body wikilinks + source_notes up
                     to MAX_HOPS, scoring each node `parent * GRAPH_DECAY**hop` (max over
                     paths) so near, strongly-cued neighbours rank above distant ones.
  * Persistence    : the BM25 postings + doc stats are a DERIVED artifact, so they live
                     in the vault's machine index dir (`_agent/index/recall-index.json`),
                     maintained on `capture` (update_note) and fully rebuilt by `sweep.py`
                     — exactly how `entities.json` is maintained. Rebuildable => portable.
  * Resilience     : if the index can't be built (vault unreachable), recall falls back
                     to the server's /search/simple so it degrades instead of dying.

Corpus = the indexable ENTITY notes (idx_mod.kind_for_path is not None): people,
companies, concepts, references, meetings, projects, areas, decisions, semantic/
episodic memory, skills — the durable memory graph. (Sessions/journal are a future
add; tracked in ROADMAP-1.0.md.)
"""
from __future__ import annotations

import json
import math
import re
import sys
import urllib.parse
from collections import Counter, deque
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import echo                      # noqa: E402
import echo_index as idx_mod     # noqa: E402
import echo_links as links       # noqa: E402

RECALL_INDEX_PATH = "_agent/index/recall-index.json"

# Tuning knobs. Defaults are standard BM25 (k1/b) + a 0.6 hop decay over <=2 hops.
K1 = 1.5
B = 0.75
MAX_HOPS = 2
GRAPH_DECAY = 0.6
MIN_SCORE = 0.0  # relevance floor; raise to trim weak lexical hits

_WORD = re.compile(r"[a-z0-9]+")
# Minimal English stoplist — keeps BM25 idf from being dominated by glue words.
_STOP = frozenset(
    "the a an and or of to in is are was were be been for on at by with as it this that "
    "from into over under not no yes do does did has have had will would can could".split()
)
_TEMPLATE_RE = re.compile(r"(^|/)(templates/|.*-template\.md$)")
_SKIP_BASENAMES = {"README.md"}


def tokenize(text: str) -> list[str]:
    """Lowercase word tokens, stopworded. Frontmatter should be stripped by the caller."""
    return [t for t in _WORD.findall(text.lower()) if t not in _STOP and len(t) > 1]


def strip_frontmatter(text: str) -> str:
    if not text.startswith("---"):
        return text
    end = text.find("\n---", 3)
    return text[end + 4:] if end != -1 else text


# --------------------------------------------------------------------- BM25 core
class Bm25Index:
    """A compact, serializable BM25 index. add()/remove() are idempotent per path."""

    def __init__(self) -> None:
        self.df: Counter[str] = Counter()              # term -> #docs containing it
        self.postings: dict[str, dict[str, int]] = {}  # term -> {doc_path: tf}
        self.length: dict[str, int] = {}               # doc_path -> token count
        self.n_docs = 0
        self.avg_len = 0.0

    def _recompute(self) -> None:
        self.n_docs = len(self.length)
        self.avg_len = (sum(self.length.values()) / self.n_docs) if self.n_docs else 0.0

    def add(self, path: str, text: str) -> None:
        if path in self.length:        # re-index in place (idempotent)
            self.remove(path)
        toks = tokenize(text)
        self.length[path] = len(toks)
        for term, tf in Counter(toks).items():
            self.postings.setdefault(term, {})[path] = tf
            self.df[term] += 1
        self._recompute()

    def remove(self, path: str) -> None:
        if path not in self.length:
            return
        for term in list(self.postings.keys()):
            if path in self.postings[term]:
                del self.postings[term][path]
                self.df[term] -= 1
                if not self.postings[term]:
                    del self.postings[term]
                    del self.df[term]
        del self.length[path]
        self._recompute()

    def score(self, query: str, limit: int = 10) -> list[tuple[str, float]]:
        scores: Counter[str] = Counter()
        for term in set(tokenize(query)):
            posting = self.postings.get(term)
            if not posting:
                continue
            idf = math.log(1 + (self.n_docs - self.df[term] + 0.5) / (self.df[term] + 0.5))
            for path, tf in posting.items():
                dl = self.length.get(path, 0)
                denom = tf + K1 * (1 - B + B * (dl / self.avg_len if self.avg_len else 1))
                scores[path] += idf * (tf * (K1 + 1) / denom if denom else 0)
        return [(p, s) for p, s in scores.most_common(limit) if s > MIN_SCORE]

    # ---- serialization (compact; rebuildable, so the format can change freely) ----
    def to_json(self) -> dict:
        return {"schema": 1, "n_docs": self.n_docs, "avg_len": self.avg_len,
                "length": self.length, "postings": self.postings}

    @classmethod
    def from_json(cls, d: dict) -> "Bm25Index":
        ix = cls()
        ix.length = d.get("length", {})
        ix.postings = d.get("postings", {})
        ix.n_docs = d.get("n_docs", len(ix.length))
        ix.avg_len = d.get("avg_len", 0.0)
        ix.df = Counter({term: len(posting) for term, posting in ix.postings.items()})
        return ix


# ------------------------------------------------------------------- vault I/O
def _get(path: str) -> str | None:
    status, body = echo.request("GET", echo.vault_url(path))
    if status == 404:
        return None
    echo.check(status, body, f"recall get {path}")
    return body.decode(errors="replace")


def _list_dir(path: str):
    p = "" if path in ("", "/") else (path if path.endswith("/") else path + "/")
    body = _get(p)
    if body is None:
        return [], []
    try:
        j = json.loads(body)
    except json.JSONDecodeError:
        return [], []
    entries = list(j.get("files", [])) + list(j.get("folders", []))
    files = [e for e in entries if not e.endswith("/")]
    folders = [e[:-1] for e in entries if e.endswith("/")]
    return files, folders


def _walk(prefix: str = ""):
    files, folders = _list_dir(prefix)
    for f in files:
        yield prefix + f
    for d in folders:
        yield from _walk(f"{prefix}{d}/")


def _indexable(path: str) -> bool:
    base = path.rsplit("/", 1)[-1]
    return (path.endswith(".md") and idx_mod.kind_for_path(path) is not None
            and base not in _SKIP_BASENAMES and not _TEMPLATE_RE.search(path))


# ------------------------------------------------------------------- persistence
def load_index() -> Bm25Index:
    status, body = echo.request("GET", echo.vault_url(RECALL_INDEX_PATH))
    if status == 404:
        return Bm25Index()
    echo.check(status, body, "recall-index load")
    try:
        return Bm25Index.from_json(json.loads(body))
    except (json.JSONDecodeError, KeyError, TypeError):
        return Bm25Index()


def save_index(ix: Bm25Index) -> None:
    body = json.dumps(ix.to_json(), ensure_ascii=False).encode("utf-8")
    status, b = echo.request("PUT", echo.vault_url(RECALL_INDEX_PATH), data=body,
                             headers={"Content-Type": "application/json"})
    echo.check(status, b, "recall-index save")


def rebuild(prefetched: dict | None = None) -> Bm25Index:
    """Full rebuild from every indexable entity note. Called by sweep.py and as the
    cold-cache path inside recall(). Saves and returns the index.

    `prefetched` (a {path: text} map, e.g. from echo.read_many) lets sweep.py reuse the
    content it already pulled instead of walking and re-fetching the whole vault again."""
    ix = Bm25Index()
    if prefetched is not None:
        items = ((p, t) for p, t in prefetched.items() if _indexable(p))
    else:
        items = ((p, _get(p)) for p in _walk() if _indexable(p))
    for path, text in items:
        if text is not None:
            ix.add(path, strip_frontmatter(text))
    save_index(ix)
    return ix


def update_note(path: str, text: str) -> None:
    """Best-effort incremental upkeep for a single note (mirrors how entities.json is
    maintained on capture). Never raises into the caller. The load->save is wrapped in the
    advisory lock with a fresh re-read (H3), so concurrent captures can't clobber the
    recall index either."""
    try:
        if not _indexable(path):
            return  # only entity notes are part of the recall corpus
        import echo_concurrency
        with echo_concurrency.vault_lock():
            ix = load_index()                      # fresh read inside the lock
            ix.add(path, strip_frontmatter(text))
            save_index(ix)
    except Exception as exc:  # noqa: BLE001 — upkeep must never fail a capture
        print(f"echo_recall: index update skipped ({exc})", file=sys.stderr)


# ------------------------------------------------------------------- graph fusion
def _resolve_target(target: str, nmap: dict) -> str | None:
    target = target.strip()
    if not target:
        return None
    if "/" in target:
        return target if target.endswith(".md") else target + ".md"
    return nmap.get(idx_mod.slugify(target))


def expand_graph(seeds, nmap: dict, base_scores: dict, max_hops: int = MAX_HOPS):
    """BFS from seeds over body wikilinks + source_notes, scoring each neighbour
    `parent_score * GRAPH_DECAY**hop` (kept as the max over reaching paths). Returns a
    score-sorted list of (path, (score, via)) excluding the seeds themselves."""
    score_of = dict(base_scores)
    seen = set(seeds)
    results: dict[str, tuple[float, str]] = {}
    # Breadth-first by HOP so each frontier's note bodies can be fetched concurrently
    # (echo.read_many) instead of one blocking GET per node. Scoring is unchanged: a
    # neighbour keeps the max decayed score over the paths that reach it, and non-seed
    # parents score 1.0 exactly as the serial deque version did.
    frontier = list(seeds)
    for hop in range(max_hops):
        if not frontier:
            break
        texts = echo.read_many(frontier)
        next_frontier: list[str] = []
        for path in frontier:
            text = texts.get(path)
            if text is None:
                continue
            body = strip_frontmatter(text)
            targets = set(links.all_wikilinks(body)) | set(links.source_notes(text))
            parent = score_of.get(path, 1.0)
            decayed = parent * (GRAPH_DECAY ** (hop + 1))
            for t in targets:
                tp = _resolve_target(t, nmap)
                if not tp or tp in seeds:
                    continue
                if decayed > results.get(tp, (0.0, ""))[0]:
                    results[tp] = (decayed, path)
                if tp not in seen:
                    seen.add(tp)
                    next_frontier.append(tp)
        frontier = next_frontier
    return sorted(results.items(), key=lambda kv: -kv[1][0])


# ------------------------------------------------------------------- presentation
def _brief(path: str, score: float | None = None, via: str | None = None) -> None:
    text = links.get_text(path)
    if text is None:
        return
    head = f"\n### {path}"
    if score is not None:
        head += f"   (score {score:.2f})"
    if via:
        head += f"   (via {via})"
    print(head)
    mtype = re.search(r"(?m)^type:\s*(.+)$", text)
    if mtype:
        print(f"_type: {mtype.group(1).strip()}_")
    body = strip_frontmatter(text)
    status = re.search(r"(?ms)^##\s+Status\s*$(.*?)(?=^##\s|\Z)", body)
    if status and status.group(1).strip():
        print(status.group(1).strip()[:400])
        return
    shown = 0
    for line in body.splitlines():
        if not line.strip() or line.startswith("# "):
            continue
        print(line[:200])
        shown += 1
        if shown >= 6:
            break


# ------------------------------------------------------------------- entrypoint
def recall(query, limit: int = 8) -> int:
    q = " ".join(query) if isinstance(query, list) else query
    index = idx_mod.load()
    nmap = idx_mod.name_map(index)

    # --- lexical layer (BM25); rebuild the cache on a cold miss --------------
    lexical: list[tuple[str, float]] = []
    try:
        ix = load_index()
        if ix.n_docs == 0:
            ix = rebuild()
        lexical = ix.score(q, limit=limit)
    except echo.EchoError as exc:
        print(f"recall: BM25 index unavailable ({exc}); falling back to server search",
              file=sys.stderr)

    base: dict[str, float] = {p: s for p, s in lexical}
    hits: list[str] = [p for p, _ in lexical]

    # --- entity-index seed (alias-aware exact match) -------------------------
    _, e = idx_mod.resolve(index, q)
    if e and e.get("path") and e["path"] not in base:
        hits.insert(0, e["path"])
        base[e["path"]] = max(base.values(), default=1.0)

    # --- resilience fallback: server search when lexical found nothing -------
    if not hits:
        try:
            status, body = echo.request(
                "POST", f"{echo.BASE}/search/simple/?query={urllib.parse.quote_plus(q)}")
            if status == 200:
                for r in (json.loads(body) or [])[:limit]:
                    fn = r.get("filename") or r.get("path")
                    if fn and fn not in base:
                        hits.append(fn)
                        base[fn] = 0.0
        except Exception:  # noqa: BLE001 — fallback only; stay quiet
            pass

    # --- graph layer ----------------------------------------------------------
    neighbours = expand_graph(hits, nmap, base, max_hops=MAX_HOPS)

    print(f"== recall: {q} ==")
    print(f"\n# Primary hits ({len(hits)})")
    for p in hits:
        _brief(p, score=base.get(p))
    print(f"\n# Linked context ({len(neighbours)})")
    for p, (sc, via) in neighbours[: 2 * limit]:
        _brief(p, score=sc, via=via)
    return 0