"""v9.0 D6 — tests for the v3 few-shot bundle mined from LoCoMo.

The v3 bundle is generated by `scripts/mine_locomo_fewshot.py --version 3
--seed 42` and consumed by `benchmarks/locomo_bench_llm._augment_category_prompts`.
These tests guard schema, determinism, dataset traceability and consumer
compatibility.
"""

from __future__ import annotations

import hashlib
import json
import sys
from pathlib import Path

import pytest

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "scripts"))
sys.path.insert(0, str(ROOT / "benchmarks"))

import mine_locomo_fewshot as mine  # noqa: E402

V3_PATH = ROOT / "benchmarks" / "data" / "locomo_few_shot_v3.json"
LOCOMO_PATH = ROOT / "benchmarks" / "data" / "locomo" / "data" / "locomo10.json"
EXPECTED_MD5 = "29258307fdfb6a33c2668298a515fe65"
EXPECTED_TRAIN_IDS = [0, 1, 2, 3, 4, 5, 6]


@pytest.fixture(scope="module")
def v3_bundle() -> dict:
    assert V3_PATH.exists(), (
        f"missing {V3_PATH}; run scripts/mine_locomo_fewshot.py --version 3 first."
    )
    return json.loads(V3_PATH.read_text())


@pytest.fixture(scope="module")
def locomo_dataset() -> list:
    return json.loads(LOCOMO_PATH.read_text())


def test_v3_bundle_has_5_categories(v3_bundle: dict) -> None:
    assert {"1", "2", "3", "4", "5"}.issubset(v3_bundle.keys())
    assert "_meta" in v3_bundle
    meta = v3_bundle["_meta"]
    assert meta["version"] == 3
    assert meta["seed"] == 42
    assert meta["n_per_category"] == 15
    assert meta["train_conv_ids"] == EXPECTED_TRAIN_IDS


def test_v3_each_category_has_15_pairs(v3_bundle: dict) -> None:
    for cat in ("1", "2", "3", "4", "5"):
        pairs = v3_bundle[cat]
        assert isinstance(pairs, list)
        assert len(pairs) == 15, f"cat {cat} has {len(pairs)} pairs, want 15"
        for p in pairs:
            assert set(p.keys()) == {"q", "a"}, f"unexpected keys in {p}"
            assert p["q"].strip(), "empty question"
            assert p["a"].strip(), "empty answer"


def test_v3_no_duplicate_questions_within_category(v3_bundle: dict) -> None:
    for cat in ("1", "2", "3", "4", "5"):
        questions = [p["q"] for p in v3_bundle[cat]]
        assert len(questions) == len(set(questions)), (
            f"cat {cat} has duplicate questions"
        )


def test_v3_pairs_sorted_lexicographically(v3_bundle: dict) -> None:
    for cat in ("1", "2", "3", "4", "5"):
        questions = [p["q"] for p in v3_bundle[cat]]
        assert questions == sorted(questions), (
            f"cat {cat} not lex-sorted by question"
        )


def test_v3_md5_deterministic(tmp_path: Path) -> None:
    on_disk_md5 = hashlib.md5(V3_PATH.read_bytes()).hexdigest()
    assert on_disk_md5 == EXPECTED_MD5, (
        f"v3 bundle md5 drift: got {on_disk_md5}, want {EXPECTED_MD5}. "
        "Re-run scripts/mine_locomo_fewshot.py --version 3 --seed 42."
    )

    # Regenerate fresh and compare bytes — proves the mining function itself
    # is deterministic, not just the on-disk artefact.
    bundle = mine.mine_v3(
        LOCOMO_PATH,
        train_conv_ids=EXPECTED_TRAIN_IDS,
        n_per_category=15,
        sample_pool=mine.V3_SAMPLE_POOL,
        seed=42,
    )
    out = tmp_path / "regen.json"
    out.write_text(json.dumps(bundle, indent=2, ensure_ascii=False))
    regen_md5 = hashlib.md5(out.read_bytes()).hexdigest()
    assert regen_md5 == EXPECTED_MD5


def test_v3_distinct_question_stems_per_category(v3_bundle: dict) -> None:
    """Diversity criterion — every category must have ≥4 distinct question stems."""
    for cat in ("1", "2", "3", "4", "5"):
        stems = {mine._question_stem(p["q"]) for p in v3_bundle[cat]}
        assert len(stems) >= 4, (
            f"cat {cat} only has {len(stems)} distinct stems: {stems}"
        )


def test_v3_distinct_answer_surfaces_per_category(v3_bundle: dict) -> None:
    """Diversity criterion — every category must have ≥3 surface-form classes."""
    for cat in ("1", "2", "3", "4", "5"):
        surfaces = {mine._answer_surface_form(p["a"]) for p in v3_bundle[cat]}
        assert len(surfaces) >= 3, (
            f"cat {cat} only has surfaces {surfaces}"
        )


def test_v3_pairs_match_locomo_dataset(v3_bundle: dict, locomo_dataset: list) -> None:
    """Every (q, a) in v3 must be traceable to a real QA in locomo10.json."""
    legal: set[tuple[str, str]] = set()
    for sample_idx, sample in enumerate(locomo_dataset):
        if sample_idx not in EXPECTED_TRAIN_IDS:
            continue
        for qa in sample.get("qa", []):
            q = str(qa.get("question", "")).strip()
            a = str(qa.get("answer", "")).strip()
            adv = str(qa.get("adversarial_answer", "")).strip()
            if q and a:
                legal.add((q, a))
            if q and adv:
                legal.add((q, adv))

    for cat in ("1", "2", "3", "4", "5"):
        for p in v3_bundle[cat]:
            assert (p["q"], p["a"]) in legal, (
                f"cat {cat} pair not in train slice: {p}"
            )


def test_v3_only_uses_train_conv_ids(locomo_dataset: list) -> None:
    """Eval-leakage guard: v3 must not pull pairs from held-out (7, 8, 9)."""
    held_out = {7, 8, 9}
    held_pairs: set[tuple[str, str]] = set()
    for sample_idx, sample in enumerate(locomo_dataset):
        if sample_idx not in held_out:
            continue
        for qa in sample.get("qa", []):
            q = str(qa.get("question", "")).strip()
            a = str(qa.get("answer", "")).strip()
            if q and a:
                held_pairs.add((q, a))

    bundle = json.loads(V3_PATH.read_text())
    for cat in ("1", "2", "3", "4", "5"):
        for p in bundle[cat]:
            assert (p["q"], p["a"]) not in held_pairs, (
                f"leakage: cat {cat} pair from held-out conv: {p}"
            )


def test_v3_canonical_answers_are_unique_within_category(v3_bundle: dict) -> None:
    """Dedupe check — no two pairs in a category share canonicalised answer."""
    for cat in ("1", "2", "3", "4", "5"):
        canon = [mine._canonical_answer(p["a"]) for p in v3_bundle[cat]]
        # Allow at most 1 duplicate in case the dataset is too small for
        # strict uniqueness; current bundle has zero, but the assertion guards
        # against silent corruption.
        assert len(set(canon)) >= len(canon) - 1, (
            f"cat {cat} canonical answers: {canon}"
        )


def test_augment_category_prompts_accepts_v3() -> None:
    """The bench's consumer must augment all 5 categories from the v3 bundle."""
    import importlib

    bench = importlib.import_module("locomo_bench_llm")
    augmented = bench._augment_category_prompts(V3_PATH)
    assert isinstance(augmented, dict)
    marker = "ADDITIONAL FEW-SHOT (mined from LoCoMo train, not eval):"
    for cat in (1, 2, 3, 4, 5):
        assert cat in augmented, f"missing cat {cat} in augmented prompts"
        prompt = augmented[cat]
        assert marker in prompt, (
            f"cat {cat} prompt was not augmented from the bundle"
        )
        # Count Q: lines only inside the appended few-shot block.
        block = prompt.split(marker, 1)[1]
        assert block.count('Q: "') == 15, (
            f"cat {cat} expected 15 Q lines in few-shot block, "
            f"got {block.count('Q: ')}"
        )


def test_augment_falls_back_when_file_missing(tmp_path: Path) -> None:
    """Sanity check: consumer must not raise when v3 file is missing."""
    import importlib

    bench = importlib.import_module("locomo_bench_llm")
    out = bench._augment_category_prompts(tmp_path / "does-not-exist.json")
    assert out == bench.CATEGORY_PROMPTS


def test_v3_meta_source_is_locomo10(v3_bundle: dict) -> None:
    src = v3_bundle["_meta"]["source"]
    assert src.endswith("locomo10.json")


def test_question_stem_helper_basic() -> None:
    assert mine._question_stem("Where did Alice go?") == "where"
    assert mine._question_stem("How many dogs does Andrew have?") == "how many"
    assert mine._question_stem("Did James meet Bob?") == "did"
    assert mine._question_stem("") == ""


def test_answer_surface_form_helper_basic() -> None:
    assert mine._answer_surface_form("Yes") == "yesno"
    assert mine._answer_surface_form("3") == "number"
    assert mine._answer_surface_form("two") == "number"
    assert mine._answer_surface_form("2023") == "year"
    assert mine._answer_surface_form("July 4") == "date"
    assert mine._answer_surface_form("Paris") == "entity"
    assert mine._answer_surface_form("a, b, c") == "list"
