"""Phase 5 — private hybrid index. Units = structural annotations, summaries, headings and short excerpts (≤ 400 chars, rights.yaml).
Backend: BM25 (rank_bm25) + local sentence-transformers embeddings stored as .npy; metadata in parquet. No full texts are indexed."""
from __future__ import annotations
import json, re
from pathlib import Path
import numpy as np, pandas as pd
from . import config, segmentation
from .provenance import stamp, now_iso

INDEX_DIR = config.ROOT / "data" / "derived" / "index"
EXCERPT = 400

def units_for_doc(doc_id: str, ann: dict | None) -> list[dict]:
    d = segmentation.load_parsed(doc_id); m = d["metadata"]; base = {"doc_id": doc_id, "sha256": d["source_sha256"], "year": m["year"], "title": m["title"], "article_type": m["article_type"], "doi": m["doi"]}
    a = (ann or {}).get("annotation") or {}
    out = []
    if a:
        out.append({**base, "unit": "document_summary", "section": None, "page": 1, "role": "document", "genre": a["genre"], "subfield": a["subfield"], "method": a["primary_method"],
                    "text": f"{m['title']}. Genre: {a['genre']}; subfield: {a['subfield']}; method: {a['primary_method']}. {a['structural_summary']} Keywords: {', '.join(a['writing_profile_keywords'])}", "moves": None, "confidence": a["confidence"]})
        out.append({**base, "unit": "identification", "section": None, "page": 1, "role": "empirical_design", "genre": a["genre"], "subfield": a["subfield"], "method": a["primary_method"],
                    "text": f"Estimand: {a['identification']['estimand']}. Strategy: {a['identification']['strategy']}. Assumptions: {'; '.join(a['identification']['key_assumptions'])}. Threats addressed: {'; '.join(a['identification']['threats_addressed'])}. Causal language strength {a['identification']['causal_language_strength']}/5.", "moves": None, "confidence": a["confidence"]})
        for ce in a["claim_evidence"]:
            out.append({**base, "unit": "claim_evidence", "section": "introduction", "page": 1, "role": "introduction", "genre": a["genre"], "subfield": a["subfield"], "method": a["primary_method"],
                        "text": f"Claim: {ce['claim']} Evidence type: {ce['evidence_type']} ({ce['evidence_ref']}).", "moves": None, "confidence": a["confidence"]})
        moves_by_para = {}
        for mv in a["introduction"]["moves"]: moves_by_para.setdefault(mv["paragraph_index"], []).append(mv["move"])
        intro = next((s for s in d["sections"] if s.get("is_intro")), None)
        if intro:
            for i, p in enumerate(segmentation.paragraphs(intro)[:12]):
                if i in moves_by_para:
                    out.append({**base, "unit": "intro_move", "section": "introduction", "page": p["page"], "role": "introduction", "genre": a["genre"], "subfield": a["subfield"], "method": a["primary_method"],
                                "text": f"[{', '.join(moves_by_para[i])}] " + segmentation.clip(p["text"], EXCERPT), "moves": "|".join(moves_by_para[i]), "confidence": a["confidence"]})
        roles = {x["idx"]: x["role"] for x in a["section_roles"]}
    else:
        roles = {}
    for i, s in enumerate(d["sections"]):
        ps = segmentation.paragraphs(s)
        if not ps: continue
        role = roles.get(i) or segmentation.section_role(s["heading"], s.get("is_appendix", False), s.get("is_intro", False))
        out.append({**base, "unit": "section_opening", "section": s["heading"], "page": s["page"], "role": role, "genre": a.get("genre"), "subfield": a.get("subfield"), "method": a.get("primary_method"),
                    "text": f"Section '{s['heading']}' ({role}, {s['n_words']} words; subsections: {', '.join(x['heading'] for x in s['subsections'][:6]) or 'none'}). Opening: " + segmentation.clip(ps[0]["text"], EXCERPT), "moves": None, "confidence": d["confidence"]["overall"]})
    for t in d["tables"]:
        if t.get("notes_excerpt"):
            out.append({**base, "unit": "table_notes", "section": t["label"], "page": t["page"], "role": "results", "genre": a.get("genre"), "subfield": a.get("subfield"), "method": a.get("primary_method"),
                        "text": f"{t['label']}: {t['title'] or ''}. Notes: " + segmentation.clip(t["notes_excerpt"], EXCERPT), "moves": None, "confidence": 0.9})
    for f in d["figures"][:8]:
        out.append({**base, "unit": "figure_caption", "section": f["label"], "page": f["page"], "role": "results", "genre": a.get("genre"), "subfield": a.get("subfield"), "method": a.get("primary_method"),
                    "text": segmentation.clip(f["caption"], EXCERPT), "moves": None, "confidence": 0.9})
    return out

def build_index(model_name: str | None = None) -> dict:
    from sentence_transformers import SentenceTransformer
    INDEX_DIR.mkdir(parents=True, exist_ok=True)
    man = pd.read_parquet(config.manifests_dir() / "corpus_manifest.parquet"); ids = man[man.decision == "include"].doc_id.tolist()
    adir = config.ROOT / "data" / "annotations" / "documents"; units = []
    for i in ids:
        ap = adir / f"{i}.json"; ann = json.loads(ap.read_text()) if ap.exists() else None
        if ann and not ann.get("schema_valid"): ann = None
        units += units_for_doc(i, ann)
    df = pd.DataFrame(units); df["unit_id"] = range(len(df))
    sp = json.loads((config.ROOT / "data" / "splits" / "heldout_v1.json").read_text()); df["split"] = np.where(df.doc_id.isin(sp["heldout"]), "heldout", "train")
    model_name = model_name or config.load("models")["embedding_model"]
    model = SentenceTransformer(model_name, device="cpu")
    emb = model.encode(df["text"].tolist(), batch_size=64, normalize_embeddings=True, show_progress_bar=False).astype(np.float32)
    np.save(INDEX_DIR / "embeddings.npy", emb); df.to_parquet(INDEX_DIR / "units.parquet", index=False)
    meta = {"built": stamp(), "n_units": int(len(df)), "n_docs": int(df.doc_id.nunique()), "n_docs_with_annotation": int(df[df.unit == "document_summary"].doc_id.nunique()), "embedding_model": model_name, "dim": int(emb.shape[1]), "units_by_type": {k: int(v) for k, v in df.unit.value_counts().items()}, "excerpt_cap_chars": EXCERPT}
    (INDEX_DIR / "index_meta.json").write_text(json.dumps(meta, indent=1)); return meta