"""Shared constants and helpers for the Decoding Parliament RAG pipeline."""
import json, os, re

DATA = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
CHUNKS = os.path.join(DATA, "chunks.jsonl")
INDEX_DIR = os.path.join(DATA, "index")          # embedding shards + metadata live here

EMBED_MODEL = "BAAI/bge-m3"      # baseline; other candidates in embed_models.py
MAX_SEQ_LEN = 1024            # Burmese ≈ 3.3 chars/token; 1,800-char chunks fit, 512 would truncate 74%
ITEM_TITLE_CAP = 200          # some QA <title>s hold the whole question text

HOUSE_NAMES = {"lower": "Pyithu Hluttaw", "upper": "Amyotha Hluttaw", "union": "Pyidaungsu Hluttaw"}
HOUSE_NAMES_MY = {"lower": "ပြည်သူ့လွှတ်တော်", "upper": "အမျိုးသားလွှတ်တော်", "union": "ပြည်ထောင်စုလွှတ်တော်"}

# health / education / social-welfare screen used to order the embedding run (themed chunks first)
THEME_KEYWORDS = ["ကျန်းမာရေး", "ပညာရေး", "ဆေးရုံ", "ကျောင်း", "ဆရာ", "ဆေး", "လူမှုဖူလုံရေး", "ဆရာဝန်",
                  "ကျောင်းသား", "တက္ကသိုလ်", "ကလေး", "မိခင်", "ပညာသင်", "သူနာပြု"]

def cap_title(t):
    t = t or ""
    return t if len(t) <= ITEM_TITLE_CAP else t[:ITEM_TITLE_CAP] + "…"

def header_of(r):
    """Citation/embedding header: house | date | section type | section title | item title | speaker."""
    return " | ".join(x for x in [HOUSE_NAMES.get(r["house"], r["house"]), r["date"], r["section_type"],
                                  cap_title(r.get("section_title")), cap_title(r.get("item_title")),
                                  r.get("speaker_name") or r.get("role")] if x)

def embed_text_of(r):
    return header_of(r) + "\n" + r["text"]

def is_themed(line_or_rec):
    s = line_or_rec if isinstance(line_or_rec, str) else json.dumps(line_or_rec, ensure_ascii=False)
    return any(k in s for k in THEME_KEYWORDS)

PARTY_SHORT = {"ပြည်ထောင်စုကြံ့ခိုင်ရေးနှင့်ဖွံ့ဖြိုးရေးပါတီ": "USDP", "အမျိုးသားဒီမိုကရေစီအဖွဲ့ချုပ်": "NLD",
               "တိုင်းရင်းသားစည်းလုံးညီညွတ်ရေးပါတီ": "NUP", "ရှမ်းတိုင်းရင်းသားများဒီမိုကရက်တစ်ပါတီ": "SNDP",
               "ရခိုင်တိုင်းရင်းသားများတိုးတက်ရေးပါတီ": "RNDP", "ရခိုင်အမျိုးသားပါတီ": "ANP", "အမျိုးသားဒီမိုကရေစီအင်အားစုပါတီ": "NDF",
               "တစ်သီးပုဂ္ဂလ": "Independent"}
_MP = None
def party_of(r):
    """Party of the speaking MP in the term of the sitting, from the member register (public political affiliation).
    Returns e.g. 'အမျိုးသားဒီမိုကရေစီအဖွဲ့ချုပ် (NLD)' or '' when unknown / not an MP."""
    global _MP
    if r.get("speaker_kind") != "mp" or not r.get("speaker_id"): return ""
    if _MP is None:
        p = os.path.join(DATA, "mp_lookup.json")
        _MP = json.load(open(p, encoding="utf-8")) if os.path.exists(p) else {}
    m = _MP.get(str(r["speaker_id"])) or {}
    party = m.get("party_t2" if r.get("term") == 2 else "party_t1") or ""
    if not party: return ""
    short = PARTY_SHORT.get(party)
    return f"{party} ({short})" if short else party

def citation_of(r):
    """Human-readable citation for answers: house, date, section, speaker (+constituency, party if known)."""
    who = r.get("speaker_name") or {"speaker": "Speaker", "vicespeaker": "Deputy Speaker", "mc": "MC"}.get(r.get("role"), r.get("role"))
    party = party_of(r)
    if who and party: who = f"{who} · {party}"
    title = cap_title(r.get('item_title') or r.get('section_title'))
    parts = [f"{HOUSE_NAMES.get(r['house'], r['house'])}, {r['date']}, {r['section_type']}"]
    if title: parts.append(title)
    if who: parts.append(who)
    return " — ".join(parts)

def meta_of(r):
    """Slim metadata kept beside each embedding (everything needed to display and cite a hit)."""
    return {k: r.get(k) for k in ("chunk_id", "speech_id", "doc_id", "house", "date", "term", "session", "sitting_day",
                                  "section_type", "section_title", "item_title", "role", "speaker_kind", "speaker_id",
                                  "speaker_name", "part", "n_parts", "text")}

def load_embedder(device=None, fp16=True, model="bge"):
    """Load an embedding model from the registry in embed_models.py (default: BGE-M3).
    Returns (SentenceTransformer, device, spec) — spec carries the query/passage prefixes and index dir."""
    import torch
    from sentence_transformers import SentenceTransformer
    from embed_models import spec
    sp = spec(model)
    if device is None:
        device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
    m = SentenceTransformer(sp["hf"], device=device, trust_remote_code=True)
    m.max_seq_length = sp["max_seq"]
    if fp16 and device != "cpu": m.half()
    return m, device, sp