#!/usr/bin/env python3
"""
load — token-budgeted context pack for a note or topic (vaultq verb; M5).

One call returns a coherent bundle instead of N speculative Reads: the seed
note + its area's folder note + the area-TODO lines that mention it + its most
RELEVANT neighbors (ranked by content similarity to the seed, not mere
adjacency), deduped and trimmed to a token budget. The model reads one pack and
knows the neighborhood — instead of paying a tool round-trip per guess.

Budget allocation (of --budget, default 6000 tokens ≈ 4 chars/tok):
  seed ≤45% · area folder note ≤12% · TODO lines ≤8% · neighbors fill the rest
  (each ≤15% so one giant neighbor can't eat the pack). Parts truncate at line
  boundaries with an explicit marker; unused share rolls forward.

Output: the assembled markdown pack (DEFAULT — this verb's product is prompt
context; JSON-wrapping it would double its cost). `--manifest` prints the JSON
manifest instead: {seed, budget, used, parts:[{rel,tokens,reason,truncated}]}.

Topic mode: if the argument resolves to no note, it's treated as a topic —
`find`'s top hit seeds the pack and the other top hits join as parts.

    python -m alexandria_vault.vault_core.query.vaultq load "Acme Pilot"
    python -m alexandria_vault.vault_core.query.vaultq load "Acme" --budget 3000 --manifest
    python -m alexandria_vault.vault_core.query.vaultq load "stablecoin clearing pricing"   # topic
"""
import os
import re
import sys
import json
import argparse
from pathlib import Path

from ...tokens import est_tokens as _tok  # noqa: E402
from .vaultgraph import find_vault_root, build_graph, split_frontmatter  # noqa: E402
from . import find as _find  # noqa: E402
from . import route as _route  # noqa: E402

BUDGET = 6000
SHARE_SEED = 0.45
SHARE_FOLDER = 0.12
SHARE_TODO = 0.08
SHARE_NEIGHBOR_CAP = 0.15
TOPIC_PARTS = 5
DIVIDER = "─" * 8


def _title(rel):
    return os.path.splitext(os.path.basename(rel))[0]


TRUNC_MARK = "…(truncated by vaultq load)"


def _trim(text, max_tokens):
    """Cut at a line boundary to fit max_tokens INCLUDING the truncation
    marker — 'used ≤ budget' is a hard invariant. → (text, truncated?)."""
    if _tok(text) <= max_tokens:
        return text.rstrip("\n"), False
    room = max_tokens - _tok(TRUNC_MARK) - 1
    out, used = [], 0
    for line in text.splitlines():
        t = _tok(line) + 1
        if used + t > room:
            break
        out.append(line)
        used += t
    return "\n".join(out).rstrip("\n") + "\n" + TRUNC_MARK, True


def _area_folder_note(rel):
    """The seed's container folder note: Areas/<A>/… → Areas/<A>/<A>.md;
    a top-level pillar (People/Knowledge/…) → its same-named folder note."""
    parts = rel.replace(os.sep, "/").split("/")
    if len(parts) < 2:
        return None
    if parts[0] == "Areas" and len(parts) > 2:
        return f"Areas/{parts[1]}/{parts[1]}.md"
    return f"{parts[0]}/{parts[0]}.md"


def _todo_lines(root, rel, names):
    """Lines in the seed's area *TODO.md files that mention the seed
    (title or alias, case-insensitive). → (todo_rel, matched lines) or None."""
    parts = rel.replace(os.sep, "/").split("/")
    if parts[0] != "Areas" or len(parts) < 3:
        return None
    area_dir = root / "Areas" / parts[1]
    pats = [re.compile(re.escape(n), re.I) for n in names if len(n) >= 4]
    for todo in sorted(area_dir.glob("*TODO.md")):
        text = todo.read_text(encoding="utf-8", errors="replace")
        hits = [ln for ln in text.splitlines()
                if any(p.search(ln) for p in pats)]
        if hits:
            return os.path.relpath(todo, root), hits
    return None


def _ranked_neighbors(root, seed_rel, notes, inbound, seed_body):
    """1-hop neighbors ranked: linked-both-ways first, then by content
    similarity to the seed (route's vectorizer over just the neighbor set)."""
    outs = set(notes[seed_rel]["out"])
    ins = set(inbound[seed_rel])
    nbrs = sorted(outs | ins)
    if not nbrs:
        return []
    corpus = [{"rel": r, "title": _title(r),
               "body": (root / r).read_text(encoding="utf-8", errors="replace")}
              for r in nbrs]
    qv, nvecs = _route._vectorize(corpus, seed_body)
    cos = {r: sum(qv[t] * v[t] for t in v) for r, v in nvecs.items()}
    return sorted(nbrs, key=lambda r: (-(r in outs and r in ins),
                                       -cos.get(r, 0.0), r))


def load(root, query, budget=BUDGET, manifest_only=False):
    """Assemble the pack. Returns (markdown, manifest); an unresolvable seed in
    note-mode falls back to topic-mode; a topic with no hits → error manifest."""
    root = Path(root)
    notes, by_name, inbound, _ = build_graph(root)

    seed_rel = query if query in notes else by_name.get(query.casefold())
    parts_plan = []      # (rel, reason, text, max_tok)

    if seed_rel:
        raw = (root / seed_rel).read_text(encoding="utf-8", errors="replace")
        _, seed_body = split_frontmatter(raw)
        parts_plan.append((seed_rel, "seed", raw, int(budget * SHARE_SEED)))

        folder = _area_folder_note(seed_rel)
        if folder and folder != seed_rel and folder in notes:
            ftext = (root / folder).read_text(encoding="utf-8", errors="replace")
            parts_plan.append((folder, "area folder note", ftext,
                               int(budget * SHARE_FOLDER)))

        names = [_title(seed_rel)] + [n for n, r in by_name.items()
                                      if r == seed_rel]
        todo = _todo_lines(root, seed_rel, names)
        if todo and todo[0] != seed_rel:
            parts_plan.append((todo[0], "area TODO lines mentioning the seed",
                               "\n".join(todo[1]), int(budget * SHARE_TODO)))

        cap = int(budget * SHARE_NEIGHBOR_CAP)
        for nrel in _ranked_neighbors(root, seed_rel, notes, inbound, seed_body):
            ntext = (root / nrel).read_text(encoding="utf-8", errors="replace")
            parts_plan.append((nrel, "neighbor (relevance-ranked)", ntext, cap))
    else:
        hits = _find.find(root, query, limit=TOPIC_PARTS)["hits"]
        if not hits:
            return None, {"error": f"nothing found for {query!r}",
                          "budget": budget}
        seed_rel = hits[0]["rel"]
        share = [SHARE_SEED] + [SHARE_NEIGHBOR_CAP] * (len(hits) - 1)
        for h, s in zip(hits, share):
            text = (root / h["rel"]).read_text(encoding="utf-8", errors="replace")
            parts_plan.append((h["rel"], f'topic hit (score {h["score"]})',
                               text, int(budget * s)))

    md_parts, manifest_parts = [], []
    seen, used = set(), 0
    for rel, reason, text, max_tok in parts_plan:
        if rel in seen:
            continue
        seen.add(rel)
        room = min(max_tok, budget - used)
        if room <= 50:      # not worth a fragment
            continue
        body, truncated = _trim(text, room)
        t = _tok(body)
        used += t
        md_parts.append(f"{DIVIDER} {rel} · {reason} · ~{t} tok {DIVIDER}\n{body}")
        manifest_parts.append({"rel": rel, "reason": reason, "tokens": t,
                               "truncated": truncated})

    manifest = {"seed": seed_rel, "budget": budget, "used": used,
                "parts": manifest_parts}
    markdown = "\n\n".join(md_parts)
    return markdown, manifest


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("query", help="note name/alias/rel path, or a topic")
    ap.add_argument("--budget", type=int, default=BUDGET, help="token budget")
    ap.add_argument("--manifest", action="store_true",
                    help="print the JSON manifest instead of the pack")
    args = ap.parse_args()
    md, man = load(find_vault_root(), args.query, budget=args.budget)
    if md is None or args.manifest:
        print(json.dumps(man, indent=1, ensure_ascii=False))
    else:
        print(md)


if __name__ == "__main__":
    main()
