"""Predeclare a broad source-grounded holdout, without reading retrieval outputs.

All profiles have already been exposed in v1-v3. This is a NEW QUERY-GROUP
holdout over the same corpus, explicitly not an unseen-person experiment.
Historic cases remain DEV; paraphrases of the same logical task stay together.
"""

from __future__ import annotations

import argparse
import copy
import json
import random
import re
from collections import Counter
from pathlib import Path

from .benchmark_v2 import (
    BASE,
    academic_binding,
    academics,
    audit_dataset,
    fingerprint,
    independent_pool,
    make_case,
    validate_case,
    write_json,
)
from .source_oracle import load_source_oracle, norm

SEED = 20261025
QUOTAS = {
    "shared_keyword_pair": 40,
    "same_academic_year_pair": 40,
    "same_start_year_pair": 32,
    "person_title_year": 24,
    "person_keyword_title": 32,
    "document_author_category_multi_hop": 24,
    "person_multiple_categories": 24,
    "candidate_set_aggregation": 32,
    "conjunctive_disambiguation": 40,
    "negative_lookup": 32,
}


def logical_group(case: dict) -> str:
    """Normalize task/subjects/constraints; never make paraphrases independent."""
    graph = case["graph"]
    category = case["category"]
    names = sorted(norm(row["value"]) for row in graph["start_entities"] if row["type"] == "person")
    bindings = graph.get("expected_bindings", [])
    keywords = sorted({norm(row["keyword"]) for row in bindings if row.get("keyword")})
    years = sorted({int(value) for value in re.findall(r"\b\d{4}\b", case["question"])})
    # A shared-year query has no requested year; its complete answer is derived.
    selectors = {"keywords": keywords, "requested_years": years}
    if category == "document_author_category_multi_hop":
        selectors["categories"] = sorted(
            {norm(row["category"]) for row in bindings if row.get("category")}
        )
    if graph.get("expected_no_answer"):
        selectors["negative_constraints"] = {
            key: norm(value) if isinstance(value, str) else value
            for key, value in graph.get("negative_constraints", {}).items()
        }
    return fingerprint({"category": category, "subjects": names, "selectors": selectors})


def paraphrase(case: dict, variant: int) -> str:
    names = [e["value"] for e in case["graph"]["start_entities"]]
    category = case["category"]
    rows = case["graph"]["expected_bindings"]
    year = next(iter(re.findall(r"\b20\d{2}\b", case["question"])), "")
    tag = next((r["keyword"] for r in rows if r.get("keyword")), "")
    if variant == 0 or category == "negative_lookup":
        return case["question"]
    if category == "shared_keyword_pair":
        return (
            f"Majú osoby {names[0]} a {names[1]} spoločné kľúčové slovo {tag}?"
            if variant == 1
            else f"Pre osoby {names[0]} a {names[1]} over rovnaké kľúčové slovo {tag}."
        )
    if category == "same_academic_year_pair":
        return (
            f"Ktoré spoločné roky prác majú osoby {names[0]} a {names[1]}?"
            if variant == 1
            else f"Porovnaj rovnaké roky práce pri osobách {names[0]} a {names[1]}."
        )
    if category == "same_start_year_pair":
        return f"Aký spoločný rok začiatku štúdia majú osoby {names[0]} a {names[1]}?"
    if category == "person_multiple_categories":
        return f"Aké všetky kategórie majú dokumenty osoby {names[0]}?"
    if category == "document_author_category_multi_hop":
        return f'Ktorý autor dokumentu je pri osobe {names[0]} spojený s kategóriou {rows[0]["category"]}?'
    if category == "person_title_year":
        # Preserve an explicit project selector, if present in the source query.
        label = (
            "diplomového projektu" if "diplomoveho projektu" in norm(case["question"]) else "práce"
        )
        return f"Uveď názov {label} osoby {names[0]} a rok práce pre rok {year}."
    if category == "person_keyword_title":
        return f"Uveď názov práce osoby {names[0]} v roku {year} s kľúčovým slovom {tag}."
    if category == "candidate_set_aggregation":
        return (
            f'Z kandidátov {", ".join(names)} vyber osoby s prácou v roku {year}. Aký je ich počet?'
        )
    if category == "conjunctive_disambiguation":
        return f"Ktorá z osôb {names[0]} a {names[1]} spĺňa kľúčové slovo {tag} a rok práce {year}?"
    return case["question"]


def build_pool(oracle: dict) -> dict[str, list[dict]]:
    names = sorted(oracle["people"], key=norm)
    pool = independent_pool(oracle, names)
    negatives = []
    # Realistic absent combinations as well as unknown entities. Positive cases
    # make up 90% of TEST; no easy-negative inflation of the headline score.
    for index, name in enumerate(names):
        other_tags = sorted(
            {
                tag
                for p in oracle["people"].values()
                for tag in p["tags"]
                if norm(tag) not in {norm(t) for t in oracle["people"][name]["tags"]}
            },
            key=norm,
        )
        if other_tags:
            tag = other_tags[index % len(other_tags)]
            negatives.append(
                make_case(
                    "",
                    "test",
                    "negative_lookup",
                    f"Aký názov práce má osoba {name} s kľúčovým slovom {tag}?",
                    [name],
                    oracle,
                    [],
                    no_answer=True,
                    reason="Requested literal tag is absent from this source document family.",
                    negative_constraints={"keyword": tag},
                )
            )
        observed = {r["year"] for r in academics(oracle, name)}
        absent_year = next(y for y in range(2018, 2027) if y not in observed)
        negatives.append(
            make_case(
                "",
                "test",
                "negative_lookup",
                f"Aký názov práce má osoba {name} v roku {absent_year}?",
                [name],
                oracle,
                [],
                no_answer=True,
                reason="Requested academic year is absent from the original profile.",
                negative_constraints={"year": absent_year},
            )
        )
    for first, last in [
        ("Jozef", "Neznámy Overovací"),
        ("Petra", "Neznáma Overovacia"),
        ("Juraj", "Nevyskytujúci Skúšobný"),
        ("Lucia", "Chýbajúca Vzorová"),
        ("Adam", "Nedostupný Kontrolný"),
        ("Eva", "Neprítomná Kontrolná"),
    ]:
        name = f"{first} {last}"
        negatives.append(
            make_case(
                "",
                "test",
                "negative_lookup",
                f"Aký názov práce má osoba {name}?",
                [name],
                oracle,
                [],
                no_answer=True,
                reason="Full name absent from all original profiles.",
                negative_constraints={"unknown_person": name},
            )
        )
    pool["negative_lookup"] = negatives
    return pool


def refresh_titles(case: dict, oracle: dict) -> dict:
    """Adjudicate title gold from raw sources before looking at TEST predictions."""
    if (
        case["category"] not in {"person_title_year", "person_keyword_title"}
        or not case["should_answer"]
    ):
        return copy.deepcopy(case)
    name = case["graph"]["start_entities"][0]["value"]
    years = [int(v) for v in re.findall(r"\b20\d{2}\b", case["question"])]
    tag = next(
        (b.get("keyword") for b in case["graph"]["expected_bindings"] if b.get("keyword")), None
    )
    text = norm(case["question"])
    kind = next(
        (
            kind
            for phrase, kind in [
                ("diplomoveho projektu", "diploma_project"),
                ("diplomovej prace", "diploma_thesis"),
                ("bakalarskej prace", "bachelor_thesis"),
            ]
            if phrase in text
        ),
        None,
    )
    records = [
        r
        for r in academics(oracle, name)
        if r.get("title")
        and (not years or r["year"] in years)
        and (kind is None or r["academic_type"] == kind)
    ]
    bindings = [academic_binding(name, r, title=True, keyword=tag) for r in records]
    return make_case(
        case["id"],
        case["split"],
        case["category"],
        case["question"],
        [name],
        oracle,
        bindings,
        no_answer=not bindings,
        reason="No explicit title for the requested academic record in original Markdown.",
        negative_constraints={
            "title_required": True,
            **({"year": years[-1]} if years else {}),
            **({"keyword": tag} if tag else {}),
            **({"academic_type": kind} if kind else {}),
        },
    )


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, default=BASE / "graphrag_questions_v4.json")
    parser.add_argument("--artifacts", type=Path, default=BASE / "results/v4_preparation")
    parser.add_argument("--seed", type=int, default=SEED)
    parser.add_argument("--extra-exposed", type=Path, action="append", default=[])
    args = parser.parse_args()
    if args.output.exists() or args.artifacts.exists():
        parser.error("Refusing to overwrite a declared benchmark or its preparation artifacts.")
    oracle = load_source_oracle()
    historic = json.loads((BASE / "graphrag_questions_v3.json").read_text())
    for path in args.extra_exposed:
        historic.extend(json.loads(path.read_text()))
    old_questions = {norm(q["question"]) for q in historic}
    exposed = {logical_group(q) for q in historic}
    adjudications = []
    for index, original in enumerate(historic):
        corrected = refresh_titles(original, oracle)
        if corrected["graph"]["expected_bindings"] != original["graph"]["expected_bindings"]:
            adjudications.append(
                {
                    "id": original["id"],
                    "split": "dev",
                    "old_bindings": original["graph"]["expected_bindings"],
                    "new_bindings": corrected["graph"]["expected_bindings"],
                    "source_evidence": corrected["source_evidence"],
                }
            )
        historic[index] = corrected
    exposed.update(logical_group(q) for q in historic)
    dev = []
    for index, original in enumerate(historic):
        case = copy.deepcopy(original)
        case["split"] = "dev"
        case["id"] = f"regression-{index + 1:04d}"
        case["graph"]["partition_group"] = logical_group(case)
        dev.append(case)
    rng = random.Random(args.seed)
    holdout = []
    capacities = {}
    for category, cases in build_pool(oracle).items():
        unique = {
            logical_group(c): c
            for c in cases
            if logical_group(c) not in exposed and norm(c["question"]) not in old_questions
        }
        candidates = sorted(unique.values(), key=lambda q: logical_group(q))
        rng.shuffle(candidates)
        quota = QUOTAS[category]
        capacities[category] = len(candidates)
        if len(candidates) < quota:
            raise RuntimeError(
                f"Insufficient unseen logical groups for {category}: {len(candidates)} < {quota}"
            )
        # Random selection is committed before running either split. Never
        # select or remove a TEST question based on a retrieval outcome.
        for index, case in enumerate(candidates[:quota]):
            case = refresh_titles(case, oracle)
            case["split"] = "test"
            case["id"] = f"wide-{len(holdout) + 1:03d}"
            case["question"] = paraphrase(case, index % 3)
            case["graph"]["partition_group"] = logical_group(case)
            holdout.append(case)
        for index, case in enumerate(
            candidates[quota : quota + min(12, max(0, len(candidates) - quota))]
        ):
            case = refresh_titles(case, oracle)
            case["split"] = "dev"
            case["id"] = f"wide-dev-{len(dev) + 1:04d}"
            case["question"] = paraphrase(case, index % 3)
            case["graph"]["partition_group"] = logical_group(case)
            dev.append(case)
    questions = dev + holdout
    problems = {q["id"]: validate_case(q, oracle) for q in questions if validate_case(q, oracle)}
    if problems:
        raise RuntimeError(f"Source oracle rejected cases: {problems}")
    groups = {
        split: {q["graph"]["partition_group"] for q in questions if q["split"] == split}
        for split in ["dev", "test"]
    }
    assert not groups["dev"] & groups["test"]
    assert not {norm(q["question"]) for q in dev} & {norm(q["question"]) for q in holdout}
    write_json(args.output, questions)
    audit = audit_dataset(questions, oracle, args.output)
    audit["partition_policy"] = {
        "method": "disjoint logical query groups",
        "shared_source_profiles": True,
        "group_function": "benchmark_v4.logical_group",
        "seed": args.seed,
        "previously_exposed_groups": sorted(exposed),
        "limitations": [
            "Known source corpus, persons, vocabulary and template families.",
            "Fresh questions/fact combinations; not an unseen-person or LLM-answer benchmark.",
        ],
    }
    write_json(args.artifacts / "source_audit.json", audit)
    write_json(
        args.artifacts / "source_adjudications.json",
        {
            "reason": "Raw-source review before any new TEST predictions; preserve historic datasets unchanged.",
            "questions": adjudications,
        },
    )
    write_json(
        args.artifacts / "partition.json",
        {
            **audit["partition_policy"],
            "dataset_sha256": audit["dataset_sha256"],
            "dev_count": len(dev),
            "test_count": len(holdout),
            "available_unseen_groups": capacities,
            "test_categories": dict(Counter(q["category"] for q in holdout)),
            "test_positive_count": sum(q["should_answer"] for q in holdout),
            "test_negative_count": sum(not q["should_answer"] for q in holdout),
            "dev_groups": sorted(groups["dev"]),
            "test_groups": sorted(groups["test"]),
            "test_profile_count": len({p for q in holdout for p in q["graph"]["source_profiles"]}),
            "source_corpus": oracle["source_corpus"],
        },
    )
    print(
        json.dumps(
            {
                "dev": len(dev),
                "test": len(holdout),
                "source_issues": 0,
                "available_unseen_groups": capacities,
            },
            ensure_ascii=False,
            indent=2,
        )
    )


if __name__ == "__main__":
    main()