"""Section-role mapping and annotation-unit construction (respects config/rights.yaml excerpt limits)."""
from __future__ import annotations
import re, json
from pathlib import Path
from . import config

ROLE_RULES = [
    ("introduction", r"^introduction$|^intro"), ("literature", r"literature|related (work|research)|prior (work|research)|background"),
    ("theory", r"model|theor|framework|hypothes|equilibrium|setup|setting|conceptual|propositions?|proof"),
    ("data", r"^data|data and|sample|measurement|variables?|descriptive|institutional (background|setting|detail)"),
    ("empirical_design", r"methodolog|empirical (strategy|design|approach|specification|framework|methodology)|identification|econometric|estimation|research design"),
    ("results", r"result|evidence|findings|estimates|main (analysis|tests)|baseline"),
    ("robustness", r"robust|additional (tests|analys|evidence)|alternative|placebo|sensitivity|extensions?|further (tests|evidence|analys)|heterogeneity|mechanism|channel"),
    ("conclusion", r"conclu|concluding|final remarks|summary|discussion"),
    ("appendix", r"^appendix|^internet appendix|^supporting information"),
]

def section_role(heading: str, is_appendix: bool = False, is_intro: bool = False) -> str:
    if is_intro: return "introduction"
    if is_appendix: return "appendix"
    h = (heading or "").lower().strip()
    for role, pat in ROLE_RULES:
        if re.search(pat, h): return role
    return "other"

def load_parsed(doc_id: str) -> dict:
    return json.loads((config.derived_dir() / "parsed" / f"{doc_id}.json").read_text())

def paragraphs(sec: dict) -> list[dict]:
    return [p for p in sec["paragraphs"] if p.get("kind") != "subheading"]

def clip(text: str, n: int) -> str:
    return text if len(text) <= n else text[:n].rsplit(" ", 1)[0] + " […]"

def annotation_payload(d: dict) -> dict:
    """Build the (bounded) material sent to the annotator: abstract, intro, conclusion, first paragraph of every section,
    table/figure captions, headings with word counts. Total size is capped by rights.yaml (max_excerpt_chars_per_call)."""
    cap = config.load("rights")["external_api_excerpt_policy"]["max_excerpt_chars_per_call"]
    secs = []
    for i, s in enumerate(d["sections"]):
        ps = paragraphs(s); role = section_role(s["heading"], s.get("is_appendix", False), s.get("is_intro", False))
        secs.append({"idx": i, "number": s.get("number"), "heading": s["heading"], "role_rule": role, "n_words": s["n_words"], "n_paragraphs": len(ps), "page": s["page"],
                     "subsections": [x["heading"] for x in s["subsections"]][:12], "first_paragraph": clip(ps[0]["text"], 700) if ps else ""})
    intro = next((s for s in d["sections"] if s.get("is_intro")), None)
    intro_paras = [clip(p["text"], 1500) for p in paragraphs(intro)] if intro else []
    concl = [s for s in d["sections"] if section_role(s["heading"], s.get("is_appendix", False)) == "conclusion"]
    concl_paras = [clip(p["text"], 1200) for p in paragraphs(concl[-1])] if concl else []
    payload = {"doc_id": d["doc_id"], "title": d["metadata"]["title"], "year": d["metadata"]["year"], "article_type": d["metadata"]["article_type"],
               "abstract": d["abstract"], "sections": secs, "introduction_paragraphs": intro_paras, "conclusion_paragraphs": concl_paras,
               "tables": [{"label": t["label"], "title": t["title"], "notes_excerpt": clip(t["notes_excerpt"] or "", 350)} for t in d["tables"]],
               "figures": [{"label": f["label"], "caption": clip(f["caption"], 300)} for f in d["figures"]],
               "stats": d["stats"], "n_equations_by_section": {}}
    for e in d["equations"]: payload["n_equations_by_section"][e["section"]] = payload["n_equations_by_section"].get(e["section"], 0) + 1
    # enforce cap by trimming intro/conclusion paragraphs from the middle
    def size(): return len(json.dumps(payload, ensure_ascii=False))
    while size() > cap and len(payload["introduction_paragraphs"]) > 4:
        mid = len(payload["introduction_paragraphs"]) // 2; payload["introduction_paragraphs"].pop(mid)
    while size() > cap and len(payload["conclusion_paragraphs"]) > 2:
        payload["conclusion_paragraphs"].pop(len(payload["conclusion_paragraphs"]) // 2)
    while size() > cap and len(payload["sections"]) > 0:
        for s in payload["sections"]: s["first_paragraph"] = clip(s["first_paragraph"], 300)
        break
    payload["payload_chars"] = size()
    return payload