"""Continue after exposed v2: all old cases become DEV; reserve the last unseen profiles.

The small supplementary holdout cannot replace the broader independent v2
measurement. It is kept profile-disjoint and must be reported with its size.
"""

from __future__ import annotations

import argparse
import copy
import json
import re
from pathlib import Path

from tools.graphrag.evaluation.benchmark_v2 import (
    BASE,
    make_case,
    correct_legacy,
    people_for,
    academics,
    has_tag,
    academic_binding,
    validate_case,
    audit_dataset,
    write_json,
)
from tools.graphrag.evaluation.source_oracle import load_source_oracle, supporting_documents, norm


def reannotate(case, oracle, original):
    if case["id"] in original:
        return correct_legacy(original[case["id"]], oracle)
    category = case["category"]
    graph = case["graph"]
    names = people_for(case, oracle)
    bindings = []
    negative = None
    expected_set = None
    field = None
    count = None
    constraints = {}
    old = graph.get("expected_bindings", [])
    years = [int(y) for y in re.findall(r"\b20\d{2}\b", case["question"])]
    keyword = next((row.get("keyword") for row in old if row.get("keyword")), None)
    if category == "negative_lookup":
        corrected = copy.deepcopy(case)
        corrected["split"] = "dev"
        return corrected
    if category == "shared_keyword_pair":
        if all(has_tag(oracle, name, keyword) for name in names):
            bindings = [{"person": name, "keyword": keyword} for name in names]
        else:
            negative = "No common literal frontmatter tag."
            constraints = {"keyword": keyword, "require_all_people": True}
    elif category == "same_academic_year_pair":
        common = set.intersection(
            *[{r["year"] for r in academics(oracle, name) if r["year"]} for name in names]
        )
        bindings = [{"person": name, "year": year} for year in sorted(common) for name in names]
        if not common:
            negative = "No shared academic year in original sources."
            constraints = {"shared_academic_year": True}
    elif category == "same_start_year_pair":
        bindings = [
            {"person": name, "start_year": oracle["people"][name]["start_year"]} for name in names
        ]
    elif category in {"person_title_year", "person_keyword_title"}:
        records = [r for r in academics(oracle, names[0], year=years[-1]) if r.get("title")]
        if "diplomoveho projektu" in norm(case["question"]):
            records = [r for r in records if r["academic_type"] == "diploma_project"]
        if category == "person_keyword_title" and not has_tag(oracle, names[0], keyword):
            records = []
        bindings = [academic_binding(names[0], row, title=True, keyword=keyword) for row in records]
        if not bindings:
            negative = (
                "Requested year/type/title is absent; a project title is not a future thesis title."
            )
            constraints = {"year": years[-1], **({"keyword": keyword} if keyword else {})}
    elif category == "document_author_category_multi_hop":
        for row in old:
            for path in supporting_documents(oracle, names[0], category=row["category"]):
                author = oracle["documents"][path]["author"]
                if author:
                    bindings.append(
                        {
                            "person": names[0],
                            "category": row["category"],
                            "document_author": author,
                            "document_path": path,
                            "same_document": True,
                        }
                    )
    elif category == "person_multiple_categories":
        expected_set = oracle["people"][names[0]]["categories"]
        field = "category"
        bindings = [{"person": names[0], "category": value} for value in expected_set]
    elif category == "candidate_set_aggregation":
        expected_set = [name for name in names if academics(oracle, name, year=years[-1])]
        count = len(expected_set)
        field = "person"
        bindings = [{"person": name, "year": years[-1]} for name in expected_set]
    elif category == "conjunctive_disambiguation":
        matches = [
            name
            for name in names
            if has_tag(oracle, name, keyword) and academics(oracle, name, year=years[-1])
        ]
        bindings = [{"person": name, "keyword": keyword, "year": years[-1]} for name in matches]
        if not matches:
            negative = "No candidate satisfies both source conditions."
            constraints = {"keyword": keyword, "year": years[-1]}
    else:
        raise ValueError(category)
    result = make_case(
        case["id"],
        "dev",
        category,
        case["question"],
        names,
        oracle,
        bindings,
        no_answer=negative is not None,
        reason=negative,
        negative_constraints=constraints,
        expected_set=expected_set,
        set_field=field,
        count=count,
    )
    result["legacy_role"] = "exposed_v2_case_development_only"
    return result


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, default=BASE / "graphrag_questions_v3.json")
    parser.add_argument("--artifacts", type=Path, default=BASE / "results/v3_preparation")
    args = parser.parse_args()
    oracle = load_source_oracle()
    previous = json.loads((BASE / "graphrag_questions_v2.json").read_text())
    original = {q["id"]: q for q in json.loads((BASE / "graphrag_questions.json").read_text())}
    dev = [reannotate(case, oracle, original) for case in previous]
    used = {p for case in previous for p in case["graph"]["source_profiles"]}
    used.update(
        json.loads((BASE / "results/v2_preparation/partition.json").read_text())[
            "previously_exposed_profiles_excluded_from_test"
        ]
    )
    reserved = sorted(
        [
            name
            for name, person in oracle["people"].items()
            if person["profile"].startswith("pages/students/") and person["profile"] not in used
        ],
        key=norm,
    )
    if len(reserved) != 2:
        raise RuntimeError(
            "Unexpected remaining unseen profile population; review partition policy."
        )
    holdout = []

    def add(category, text, names, bindings, **kwargs):
        holdout.append(
            make_case(
                f"fresh-{len(holdout)+1:03d}",
                "test",
                category,
                text,
                names,
                oracle,
                bindings,
                **kwargs,
            )
        )

    for name in reserved:
        person = oracle["people"][name]
        add(
            "person_multiple_categories",
            f"Vypíš všetky kategórie dokumentov osoby {name}.",
            [name],
            [{"person": name, "category": category} for category in person["categories"]],
            expected_set=person["categories"],
            set_field="category",
        )
        category = person["categories"][0]
        for path in supporting_documents(oracle, name, category=category)[:1]:
            author = oracle["documents"][path]["author"]
            if author:
                add(
                    "document_author_category_multi_hop",
                    f"Ktorý autor dokumentu je pri osobe {name} spojený s kategóriou {category}?",
                    [name],
                    [
                        {
                            "person": name,
                            "category": category,
                            "document_author": author,
                            "document_path": path,
                            "same_document": True,
                        }
                    ],
                )
        add(
            "person_start_year",
            f"Aký rok začiatku štúdia je uvedený pri osobe {name}?",
            [name],
            [{"person": name, "start_year": person["start_year"]}],
        )
        add(
            "negative_lookup",
            f"Aký názov práce má osoba {name}, ak jej dokument obsahuje kľúčové slovo nenajdenyoverovacitag?",
            [name],
            [],
            no_answer=True,
            reason="Tag absent from raw frontmatter.",
            negative_constraints={"keyword": "nenajdenyoverovacitag"},
        )
        add(
            "negative_lookup",
            f"Aký názov práce má osoba {name} v roku 2098?",
            [name],
            [],
            no_answer=True,
            reason="No academic record for requested year.",
            negative_constraints={"year": 2098},
        )
    for name in ["Diana Neznáma Overovacia", "Milan Neprítomný Vzorový"]:
        add(
            "negative_lookup",
            f"Aký názov diplomovej práce má osoba {name}?",
            [name],
            [],
            no_answer=True,
            reason="No profile with requested full name.",
            negative_constraints={"unknown_person": name},
        )
    add(
        "candidate_set_aggregation",
        f'Z kandidátov {", ".join(reserved)} vyber osoby s prácou v roku 2098. Koľko ich je?',
        reserved,
        [],
        expected_set=[],
        set_field="person",
        count=0,
    )
    questions = dev + holdout
    issues = {q["id"]: validate_case(q, oracle) for q in questions if validate_case(q, oracle)}
    if issues:
        raise RuntimeError(f"Independent source audit failed: {issues}")
    dev_profiles = {p for q in dev for p in q["graph"]["source_profiles"]}
    test_profiles = {p for q in holdout for p in q["graph"]["source_profiles"]}
    assert not dev_profiles & test_profiles
    write_json(args.output, questions)
    write_json(args.artifacts / "source_audit.json", audit_dataset(questions, oracle, args.output))
    changes = [
        {
            "id": old["id"],
            "old_bindings": old["graph"]["expected_bindings"],
            "new_bindings": new["graph"]["expected_bindings"],
            "old_should_answer": old["should_answer"],
            "new_should_answer": new["should_answer"],
            "source_evidence": new["source_evidence"],
        }
        for old, new in zip(previous, dev)
        if old["graph"]["expected_bindings"] != new["graph"]["expected_bindings"]
        or old["should_answer"] != new["should_answer"]
    ]
    write_json(
        args.artifacts / "source_adjudications.json",
        {
            "method": "Raw heading project/thesis distinction; no answer-output based relabeling.",
            "changes": changes,
        },
    )
    partition = {
        "dev_count": len(dev),
        "test_count": len(holdout),
        "dev_profiles": sorted(dev_profiles),
        "test_profiles": sorted(test_profiles),
        "source_profile_overlap": 0,
        "limitations": [
            "Only two entirely unseen profiles remain; supplementary holdout has limited category coverage.",
            "The broader v2 independent TEST remains 67/80; v2 cases are now DEV, never fresh TEST.",
        ],
    }
    write_json(args.artifacts / "partition.json", partition)
    print(
        json.dumps(
            {
                "dev": len(dev),
                "test": len(holdout),
                "unseen_profiles": len(reserved),
                "source_adjudications": len(changes),
                "source_issues": 0,
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    main()