"""Phase 3 — scientific rhetoric ontology: schema, LLM annotation (structured output), deterministic checks, agreement."""
from __future__ import annotations
import json, re, concurrent.futures as cf
from pathlib import Path
import pandas as pd
from . import config, segmentation, llm
from .provenance import stamp, now_iso, Ledger, sha256_text

PROMPT_PATH = Path(__file__).resolve().parents[2] / "prompts" / "annotation" / "document_annotation_v1.md"
ONTOLOGY_VERSION = "1.0"
GENRES = ["empirical", "theoretical", "mixed_theory_empirical", "methodological_econometric", "survey_or_address", "replication_or_correction"]
SUBFIELDS = ["asset_pricing", "corporate_finance", "financial_intermediation_banking", "household_finance", "market_microstructure", "international_finance", "behavioral_finance", "macro_finance", "financial_econometrics", "entrepreneurial_finance_private_equity", "other"]
METHODS = ["reduced_form_regression", "difference_in_differences", "instrumental_variables", "regression_discontinuity", "natural_experiment_event_study", "structural_estimation", "calibration_simulation", "theoretical_model_analytical", "machine_learning_textual", "survey_or_experiment", "other", "none"]
MOVES = ["motivation", "tension_or_puzzle", "literature_gap", "preview_of_approach", "preview_of_results", "contribution_statement", "mechanism_or_theory_sketch", "identification_defense", "robustness_mention", "roadmap", "related_literature_positioning", "policy_or_broader_implication", "other"]
CONTRIB = ["new_evidence", "new_data_or_measure", "new_identification", "new_theory_or_model", "new_method", "resolution_of_puzzle", "replication_or_correction", "synthesis"]
ROLES = ["introduction", "literature", "theory", "data", "empirical_design", "results", "robustness", "conclusion", "appendix", "other"]
ROBUST_TYPES = ["alternative_specification", "placebo", "subsample", "alternative_measure", "alternative_sample_period", "falsification", "sensitivity_to_parameters", "other"]

def _obj(props, req=None): return {"type": "object", "properties": props, "required": req or list(props), "additionalProperties": False}
STRICT_ENUMS = True
def _enum(vals): return {"type": "string", "enum": vals} if STRICT_ENUMS else {"type": "string"}
def _arr(item, maxn=None): d = {"type": "array", "items": item}; return d
def build_schema(strict: bool) -> dict:
    global STRICT_ENUMS; STRICT_ENUMS = strict
    MOVE_FLAG = _obj({"present": {"type": "boolean"}, "sentence_index": {"type": "integer"}})
    return _obj({
    "genre": _enum(GENRES), "subfield": _enum(SUBFIELDS), "subfield_secondary": {"type": "string"}, "primary_method": _enum(METHODS),
    "identification": _obj({"estimand": {"type": "string"}, "strategy": {"type": "string"}, "key_assumptions": _arr({"type": "string"}), "threats_addressed": _arr({"type": "string"}),
                            "causal_language_strength": {"type": "integer"}, "uses_causal_claim": {"type": "boolean"}}),
    "abstract_moves": _obj({k: MOVE_FLAG for k in ["question", "gap", "method", "results", "magnitude", "contribution"]}),
    "introduction": _obj({"moves": _arr(_obj({"paragraph_index": {"type": "integer"}, "move": _enum(MOVES)})), "contribution_paragraph_index": {"type": "integer"},
                          "contribution_types": _arr(_enum(CONTRIB)), "roadmap_present": {"type": "boolean"}, "literature_positioning_location": _enum(["inside_introduction", "separate_section", "both", "none"])}),
    "section_roles": _arr(_obj({"idx": {"type": "integer"}, "role": _enum(ROLES)})),
    "results_presentation": _obj({"economic_significance_discussed": {"type": "boolean"}, "statistical_significance_language": _enum(["stars_and_tstats", "confidence_intervals", "p_values", "mixed", "unclear"]), "magnitude_statements_in_intro": {"type": "boolean"}}),
    "robustness": _obj({"location": _enum(["main_text_section", "within_results", "appendix", "internet_appendix", "mixed", "none_seen"]), "types": _arr(_enum(ROBUST_TYPES))}),
    "tables": _obj({"convention_notes": {"type": "string"}, "units_stated_in_notes": {"type": "boolean"}, "standard_errors_description": _enum(["clustered", "robust", "bootstrap", "newey_west", "mixed", "unclear", "none"])}),
    "figures": _obj({"uncertainty_shown": _enum(["yes", "no", "mixed", "unclear", "no_figures"])}),
    "conclusion": _obj({"bounded_claims": {"type": "boolean"}, "policy_implications": {"type": "boolean"}, "open_questions": {"type": "boolean"}, "restates_contribution": {"type": "boolean"}}),
    "hedging": _obj({"level": _enum(["low", "moderate", "high"]), "typical_devices": _arr({"type": "string"})}),
    "claim_evidence": _arr(_obj({"claim": {"type": "string"}, "evidence_type": _enum(["table", "figure", "equation_or_proof", "prior_literature", "descriptive_statistic", "none_stated"]), "evidence_ref": {"type": "string"}})),
    "structural_summary": {"type": "string"}, "writing_profile_keywords": _arr({"type": "string"}), "confidence": {"type": "number"}, "notes": {"type": "string"},
    })
SCHEMA_FULL = build_schema(True)     # local validation (jsonschema)
SCHEMA = build_schema(False)         # API grammar (enums as free strings: the compiled grammar must stay small)
VOCAB = {"genre": GENRES, "subfield": SUBFIELDS, "primary_method": METHODS, "introduction.moves[].move": MOVES, "introduction.contribution_types[]": CONTRIB, "section_roles[].role": ROLES, "robustness.types[]": ROBUST_TYPES,
         "introduction.literature_positioning_location": ["inside_introduction", "separate_section", "both", "none"], "results_presentation.statistical_significance_language": ["stars_and_tstats", "confidence_intervals", "p_values", "mixed", "unclear"],
         "robustness.location": ["main_text_section", "within_results", "appendix", "internet_appendix", "mixed", "none_seen"], "tables.standard_errors_description": ["clustered", "robust", "bootstrap", "newey_west", "mixed", "unclear", "none"],
         "figures.uncertainty_shown": ["yes", "no", "mixed", "unclear", "no_figures"], "hedging.level": ["low", "moderate", "high"], "claim_evidence[].evidence_type": ["table", "figure", "equation_or_proof", "prior_literature", "descriptive_statistic", "none_stated"]}

def validate_full(ann: dict) -> list[str]:
    import jsonschema
    return [e.message[:160] + " @ " + "/".join(str(x) for x in e.absolute_path) for e in jsonschema.Draft202012Validator(SCHEMA_FULL).iter_errors(ann)]

def prompt_text() -> tuple[str, str]:
    raw = PROMPT_PATH.read_text(); m = re.search(r'version:\s*"([^"]+)"', raw)
    body = raw.split("---", 2)[2].strip()
    body += "\n\nAllowed values (use exactly these strings):\n" + "\n".join(f"- {k}: {', '.join(v)}" for k, v in VOCAB.items())
    body += "\n\nOutput: a single JSON object (no code fence, no prose) that validates against this JSON Schema:\n" + json.dumps(SCHEMA_FULL, separators=(",", ":"))
    return body, (m.group(1) if m else "0")

def deterministic_checks(d: dict, payload: dict) -> dict:
    """Rule-based signals computed from the parsed text, independent of the LLM."""
    intro = "\n".join(payload["introduction_paragraphs"]); abs_ = payload["abstract"] or ""
    checks = {
        "roadmap_regex": bool(re.search(r"(remainder|rest) of (the|this) (paper|article)|paper (proceeds|is organized)|Section [IVX]+ (presents|describes|discusses|concludes|develops|introduces|reports)", intro)),
        "contribution_regex_in_intro": bool(re.search(r"\b(contribut|this paper (is the first|adds|extends|provides|offers|shows))", intro, re.I)),
        "magnitude_in_abstract": bool(re.search(r"\d+(\.\d+)?\s?(%|percent|basis points|bps|percentage points|standard deviation)", abs_, re.I)),
        "abstract_n_sentences": len([s for s in re.split(r"(?<=[.!?])\s+", abs_) if s.strip()]) if abs_ else 0,
        "intro_n_paragraphs": len(payload["introduction_paragraphs"]), "intro_n_words": sum(len(p.split()) for p in payload["introduction_paragraphs"]),
        "causal_words_in_abstract": len(re.findall(r"\b(causal|causes?|effect of|impact of|leads? to|induces?|exogenous|plausibly)\b", abs_, re.I)),
        "hedge_words_in_abstract": len(re.findall(r"\b(suggest|may|might|could|appear|consistent with|likely|tend to|indicate)\b", abs_, re.I)),
        "stars_mentioned": any(re.search(r"\*\*\*|\*\*|significance at", (t.get("notes_excerpt") or "")) for t in payload["tables"]),
        "clustered_se_mentioned": any(re.search(r"clustered", (t.get("notes_excerpt") or ""), re.I) for t in payload["tables"]),
        "internet_appendix_mentioned": bool(re.search(r"Internet Appendix", intro + " " + json.dumps(payload["sections"]))),
        "section_roles_rule": {s["idx"]: s["role_rule"] for s in payload["sections"]},
    }
    return checks

def repair(ann: dict, schema: dict | None = None, removed: list | None = None, path: str = "") -> dict:
    """Deterministic repair: drop unknown keys where the schema forbids additional properties; hoist top-level keys that the model
    nested one level too deep. Every removal is recorded so the repair is auditable."""
    schema = schema or SCHEMA_FULL; removed = removed if removed is not None else []
    if not isinstance(ann, dict): return ann
    props = schema.get("properties", {})
    if path == "":  # hoist misplaced top-level keys
        for k in list(ann.keys()):
            if isinstance(ann[k], dict):
                for kk in list(ann[k].keys()):
                    if kk in props and kk not in ann and kk not in props.get(k, {}).get("properties", {}):
                        ann[kk] = ann[k].pop(kk); removed.append(f"hoisted {k}.{kk} -> {kk}")
    for k in list(ann.keys()):
        if k not in props and schema.get("additionalProperties") is False:
            removed.append(f"dropped {path}{k}"); ann.pop(k); continue
        sub = props.get(k, {})
        if isinstance(ann[k], dict) and sub.get("type") == "object": repair(ann[k], sub, removed, f"{path}{k}.")
        elif isinstance(ann[k], list) and sub.get("items", {}).get("type") == "object":
            for i, it in enumerate(ann[k]): repair(it, sub["items"], removed, f"{path}{k}[{i}].")
    return ann

def annotate_document(doc_id: str, force: bool = False) -> dict:
    outdir = config.ROOT / "data" / "annotations" / "documents"; outdir.mkdir(parents=True, exist_ok=True)
    outp = outdir / f"{doc_id}.json"
    if outp.exists() and not force: return json.loads(outp.read_text())
    d = segmentation.load_parsed(doc_id); payload = segmentation.annotation_payload(d)
    system, pv = prompt_text(); cfg = config.load("models")
    user = "Annotate this article.\n\n```json\n" + json.dumps({k: v for k, v in payload.items() if k != "payload_chars"}, ensure_ascii=False) + "\n```"
    res = llm.call(system, user, model=cfg["annotation_model"], schema=None, parse_json=True, max_tokens=cfg["max_tokens_annotation"], effort=cfg["effort_annotation"],
                   purpose="annotation", prompt_id="document_annotation", prompt_version=pv, meta={"doc_id": doc_id})
    ann = res["json"]; removed = []
    if ann: ann = repair(ann, removed=removed)
    verrs = validate_full(ann) if ann else ["no json"]
    if verrs and force:  # one uncached retry with the validation errors echoed back
        user2 = user + "\n\nYour previous answer failed strict validation: " + "; ".join(verrs[:5]) + ". Return only the JSON object, with exactly the schema's keys and allowed values."
        res = llm.call(system, user2, model=cfg["annotation_model"], schema=None, parse_json=True, max_tokens=cfg["max_tokens_annotation"], effort=cfg["effort_annotation"],
                       purpose="annotation_retry", prompt_id="document_annotation", prompt_version=pv, use_cache=False, meta={"doc_id": doc_id})
        ann = res["json"]; removed = []
        if ann: ann = repair(ann, removed=removed)
        verrs = validate_full(ann) if ann else ["no json"]
    rec = {"validation_errors": verrs, "repairs": removed, "doc_id": doc_id, "ontology_version": ONTOLOGY_VERSION, "prompt_version": pv, "model": res["model"], "effort": res["effort"], "annotated_at": now_iso(),
           "source_sha256": d["source_sha256"], "parser": d["parser"], "payload_chars": payload["payload_chars"], "usage": res["usage"], "cost_usd": res["cost_usd"], "request_id": res["request_id"],
           "annotation": ann, "deterministic_checks": deterministic_checks(d, payload), "schema_valid": ann is not None and not verrs, **{"code_version": stamp()["code_version"]}}
    if ann is None: rec["raw_text"] = res["text"][:2000]
    outp.write_text(json.dumps(rec, ensure_ascii=False, indent=1))
    return rec

def annotate_corpus(limit=None, force=False, only=None, workers=4) -> dict:
    man = pd.read_parquet(config.manifests_dir() / "corpus_manifest.parquet"); inc = man[man.decision == "include"]
    ids = [only] if only else inc.doc_id.tolist()
    if limit: ids = ids[:limit]
    ok = err = 0; errors = []
    with cf.ThreadPoolExecutor(max_workers=workers) as ex:
        futs = {ex.submit(annotate_document, i, force): i for i in ids}
        for f in cf.as_completed(futs):
            try:
                r = f.result(); ok += int(r["schema_valid"]); err += int(not r["schema_valid"])
            except llm.BudgetExceeded as e:
                errors.append(str(e)); break
            except Exception as e:
                err += 1; errors.append(f"{futs[f]}: {repr(e)[:200]}")
    build_table()
    return {"ok": ok, "errors": err, "error_samples": errors[:5], "spent_usd_total": llm.spent_usd()}

def build_table() -> pd.DataFrame:
    rows = []
    for p in sorted((config.ROOT / "data" / "annotations" / "documents").glob("*.json")):
        r = json.loads(p.read_text()); a = r.get("annotation") or {}; c = r["deterministic_checks"]
        if not a or not r.get("schema_valid"): rows.append({"doc_id": r["doc_id"], "schema_valid": False, "validation_errors": "; ".join(r.get("validation_errors", []))[:300]}); continue
        am = a["abstract_moves"]; intro = a["introduction"]
        moves = [m["move"] for m in sorted(intro["moves"], key=lambda m: m["paragraph_index"])]
        rows.append({"doc_id": r["doc_id"], "schema_valid": True, "genre": a["genre"], "subfield": a["subfield"], "primary_method": a["primary_method"],
                     "causal_strength": a["identification"]["causal_language_strength"], "uses_causal_claim": a["identification"]["uses_causal_claim"],
                     **{f"abs_{k}": v["present"] for k, v in am.items()}, **{f"abs_{k}_idx": v["sentence_index"] for k, v in am.items()},
                     "intro_moves_seq": "|".join(moves), "intro_n_moves": len(moves), "contribution_para_idx": intro["contribution_paragraph_index"], "contribution_types": "|".join(intro["contribution_types"]),
                     "roadmap_llm": intro["roadmap_present"], "roadmap_regex": c["roadmap_regex"], "lit_positioning": intro["literature_positioning_location"],
                     "econ_signif": a["results_presentation"]["economic_significance_discussed"], "stat_lang": a["results_presentation"]["statistical_significance_language"], "magnitude_in_intro": a["results_presentation"]["magnitude_statements_in_intro"],
                     "robust_location": a["robustness"]["location"], "robust_types": "|".join(a["robustness"]["types"]), "se_desc": a["tables"]["standard_errors_description"], "units_in_notes": a["tables"]["units_stated_in_notes"],
                     "fig_uncertainty": a["figures"]["uncertainty_shown"], "concl_bounded": a["conclusion"]["bounded_claims"], "concl_policy": a["conclusion"]["policy_implications"], "concl_open": a["conclusion"]["open_questions"],
                     "hedging": a["hedging"]["level"], "n_claims": len(a["claim_evidence"]), "claims_with_evidence": sum(1 for x in a["claim_evidence"] if x["evidence_type"] != "none_stated"),
                     "llm_confidence": a["confidence"], "abstract_n_sent": c["abstract_n_sentences"], "intro_n_paragraphs": c["intro_n_paragraphs"], "intro_n_words": c["intro_n_words"],
                     "magnitude_regex_abstract": c["magnitude_in_abstract"], "contribution_regex_intro": c["contribution_regex_in_intro"], "clustered_se_regex": c["clustered_se_mentioned"], "stars_regex": c["stars_mentioned"],
                     "role_agreement": role_agreement(a, c), "cost_usd": r["cost_usd"], "model": r["model"]})
    df = pd.DataFrame(rows); out = config.ROOT / "data" / "annotations"
    df.to_parquet(out / "annotations_table.parquet", index=False); df.to_csv(out / "annotations_table.csv", index=False)
    return df

def role_agreement(a: dict, c: dict) -> float | None:
    rule = {int(k): v for k, v in c["section_roles_rule"].items()}; llm_roles = {x["idx"]: x["role"] for x in a["section_roles"]}
    common = [i for i in rule if i in llm_roles and rule[i] != "other"]
    return round(sum(rule[i] == llm_roles[i] for i in common) / len(common), 3) if common else None