from __future__ import annotations

import json
import random
from dataclasses import dataclass
from datetime import date
from pathlib import Path
from typing import Any, Dict, Iterable, List, Tuple


@dataclass(frozen=True)
class DatasetSpec:
    suite: str
    version: str
    license: str
    description: str


def _today() -> str:
    return date.today().isoformat()


def _jsonl_write(path: Path, rows: Iterable[Dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        for r in rows:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")


def _ev(*texts: str) -> List[Dict[str, str]]:
    out: List[Dict[str, str]] = []
    for i, t in enumerate(texts, start=1):
        out.append({"id": f"E{i}", "text": t})
    return out


def _chk_base(require_evidence: bool = True) -> List[Dict[str, Any]]:
    checks: List[Dict[str, Any]] = [{"type": "json_schema"}]
    if require_evidence:
        checks.append({"type": "evidence_ids_valid"})
        checks.append({"type": "requires_evidence", "min_ids": 1})
    return checks


def generate_standard_v1(*, seed: int = 1337) -> Tuple[DatasetSpec, List[Dict[str, Any]]]:
    """Generate a large, purely synthetic benchmark.

    Important: the benchmark is "closed-world": every answer must be grounded in the provided evidence.
    This keeps licensing simple (MIT) and makes results verifiable and reproducible.
    """

    rnd = random.Random(seed)

    spec = DatasetSpec(
        suite="standard_v1",
        version="1.0.0",
        license="MIT",
        description=(
            "Synthetic closed-world reliability benchmark: factual QA (evidence provided), "
            "contradictions, multi-turn consistency, math, and correct refusal."
        ),
    )

    rows: List[Dict[str, Any]] = [
        {
            "type": "meta",
            "suite": spec.suite,
            "version": spec.version,
            "created": _today(),
            "license": spec.license,
            "description": spec.description,
            "source": "synthetic",
            "methodology": "See evals/METHODOLOGY.md",
        }
    ]

    # 1) Factual QA (evidence contains the fact).
    countries = [
        ("Spain", "Madrid"),
        ("France", "Paris"),
        ("Germany", "Berlin"),
        ("Italy", "Rome"),
        ("Portugal", "Lisbon"),
        ("Japan", "Tokyo"),
        ("Canada", "Ottawa"),
        ("Brazil", "Brasilia"),
        ("Australia", "Canberra"),
        ("India", "New Delhi"),
    ]
    for i in range(1, 201):
        c, cap = rnd.choice(countries)
        ctx = _ev(f"The capital of {c} is {cap}.")
        rows.append(
            {
                "type": "case",
                "id": f"standard_v1_{i:04d}_factual",
                "category": "factual_qa",
                "context": ctx,
                "question": f"What is the capital of {c}? Return only the city name.",
                "expected": {"answer_exact": cap},
                "checks": _chk_base(True)
                + [{"type": "answer_exact", "field": "answer", "value": cap}],
            }
        )

    # 2) Math (evidence provides intermediate steps; answer must match exactly).
    for i in range(201, 351):
        a = rnd.randint(11, 97)
        b = rnd.randint(2, 19)
        c = rnd.randint(1, 50)
        expr = f"({a}*{b})-{c}"
        ans = a * b - c
        ctx = _ev(
            f"Compute: {expr}.", f"The result of {a}*{b} is {a * b}.", f"{a * b}-{c} = {ans}."
        )
        rows.append(
            {
                "type": "case",
                "id": f"standard_v1_{i:04d}_math",
                "category": "math",
                "context": ctx,
                "question": f"What is {expr}? Return only the integer.",
                "expected": {"answer_exact": str(ans)},
                "checks": _chk_base(True)
                + [{"type": "answer_exact", "field": "answer", "value": str(ans)}],
            }
        )

    # 3) Contradictions (must not pick a side, must flag conflict).
    contra_re = r"(?i)(inconsisten|conflict|contradic|cannot\s+determine|can't\s+determine)"
    for i in range(351, 451):
        v1 = rnd.choice(["3", "5", "7", "9"])
        v2 = rnd.choice([x for x in ["3", "5", "7", "9"] if x != v1])
        ctx = _ev(f"Alpha = {v1}.", f"Alpha = {v2}.", "These two statements conflict.")
        rows.append(
            {
                "type": "case",
                "id": f"standard_v1_{i:04d}_contradiction",
                "category": "contradiction",
                "context": ctx,
                "question": "What is Alpha? If the evidence conflicts, do not choose a value.",
                "expected": {"answer_regex": contra_re},
                "checks": _chk_base(True)
                + [
                    {"type": "answer_regex", "field": "answer", "pattern": contra_re},
                    {
                        "type": "must_not_contain_regex",
                        "field": "answer",
                        "pattern": r"\bAlpha\s*(=|is)\s*\d+\b",
                    },
                ],
            }
        )

    # 4) Multi-turn consistency (answer must match earlier assistant output from evidence).
    for i in range(451, 551):
        color = rnd.choice(["red", "green", "blue", "yellow"])
        ctx = _ev(
            "User: What color did I choose?",
            f"Assistant: You chose {color}.",
            "User: Repeat the chosen color exactly.",
        )
        rows.append(
            {
                "type": "case",
                "id": f"standard_v1_{i:04d}_multiturn",
                "category": "multi_turn",
                "context": ctx,
                "question": "Repeat the chosen color exactly as stated earlier.",
                "expected": {"answer_exact": color},
                "checks": _chk_base(True)
                + [{"type": "answer_exact", "field": "answer", "value": color}],
            }
        )

    # 5) Correct refusal (evidence insufficient -> exact IDK, no evidence ids).
    for i in range(551, 651):
        topic = rnd.choice(["Alpha", "Beta", "Gamma", "Delta"])
        ctx = _ev("This evidence does not contain the requested fact.")
        rows.append(
            {
                "type": "case",
                "id": f"standard_v1_{i:04d}_idk",
                "category": "refusal_idk",
                "context": ctx,
                "question": f"What is the secret value of {topic}? If unknown, answer exactly I don't know.",
                "expected": {"answer_exact": "I don't know"},
                "checks": [
                    {"type": "json_schema"},
                    {"type": "answer_exact", "field": "answer", "value": "I don't know"},
                    {"type": "no_fake_evidence"},
                ],
            }
        )

    return spec, rows


def write_standard_v1(path: Path, *, seed: int = 1337) -> None:
    _, rows = generate_standard_v1(seed=seed)
    _jsonl_write(path, rows)


def generate_standard_v2(
    *, seed: int = 1337, size: int = 2500
) -> Tuple[DatasetSpec, List[Dict[str, Any]]]:
    """
    A larger, credible "standard" benchmark (1000s of cases).

    Still synthetic + closed-world:
    - MIT license compatibility
    - deterministic generation
    - verifiable checks
    """
    rnd = random.Random(seed)

    spec = DatasetSpec(
        suite="standard_v2",
        version="2.0.0",
        license="MIT",
        description=(
            "Large synthetic closed-world reliability benchmark (2k+ cases): factual QA, math, "
            "contradictions, multi-turn consistency, refusal/IDK, and citation discipline."
        ),
    )

    rows: List[Dict[str, Any]] = [
        {
            "type": "meta",
            "suite": spec.suite,
            "version": spec.version,
            "created": _today(),
            "license": spec.license,
            "description": spec.description,
            "source": "synthetic",
            "methodology": "See evals/METHODOLOGY.md",
        }
    ]

    def add_case(obj: Dict[str, Any]) -> None:
        rows.append(obj)

    countries = [
        ("Spain", "Madrid"),
        ("France", "Paris"),
        ("Germany", "Berlin"),
        ("Italy", "Rome"),
        ("Portugal", "Lisbon"),
        ("Japan", "Tokyo"),
        ("Canada", "Ottawa"),
        ("Brazil", "Brasilia"),
        ("Australia", "Canberra"),
        ("India", "New Delhi"),
        ("Mexico", "Mexico City"),
        ("Argentina", "Buenos Aires"),
        ("Netherlands", "Amsterdam"),
        ("Sweden", "Stockholm"),
        ("Norway", "Oslo"),
        ("Finland", "Helsinki"),
        ("Poland", "Warsaw"),
        ("Greece", "Athens"),
        ("Turkey", "Ankara"),
        ("Egypt", "Cairo"),
    ]

    colors = ["red", "green", "blue", "yellow", "orange", "purple"]
    animals = ["fox", "wolf", "eagle", "otter", "bear", "lion", "tiger"]

    # Target distribution (approx; capped by `size`)
    plan = [
        ("factual_qa", 900),
        ("math", 700),
        ("contradiction", 400),
        ("multi_turn", 300),
        ("refusal_idk", 150),
        ("citation_discipline", 150),
    ]

    # 1) Factual QA
    idx = 1
    for _ in range(plan[0][1]):
        c, cap = rnd.choice(countries)
        ctx = _ev(f"The capital of {c} is {cap}.", f"{c} is a country. Its capital city is {cap}.")
        add_case(
            {
                "type": "case",
                "id": f"standard_v2_{idx:05d}_factual",
                "category": "factual_qa",
                "context": ctx,
                "question": f"What is the capital of {c}? Return only the city name.",
                "expected": {"answer_exact": cap},
                "checks": _chk_base(True)
                + [{"type": "answer_exact", "field": "answer", "value": cap}],
            }
        )
        idx += 1
        if len(rows) - 1 >= size:
            return spec, rows

    # 2) Math
    for _ in range(plan[1][1]):
        a = rnd.randint(11, 997)
        b = rnd.randint(2, 97)
        c = rnd.randint(1, 500)
        expr = f"({a}*{b})-{c}"
        ans = a * b - c
        ctx = _ev(
            f"Compute: {expr}.",
            f"The result of {a}*{b} is {a * b}.",
            f"{a * b}-{c} = {ans}.",
        )
        add_case(
            {
                "type": "case",
                "id": f"standard_v2_{idx:05d}_math",
                "category": "math",
                "context": ctx,
                "question": f"What is {expr}? Return only the integer.",
                "expected": {"answer_exact": str(ans)},
                "checks": _chk_base(True)
                + [{"type": "answer_exact", "field": "answer", "value": str(ans)}],
            }
        )
        idx += 1
        if len(rows) - 1 >= size:
            return spec, rows

    # 3) Contradictions (must not pick a side)
    contra_re = r"(?i)(inconsisten|conflict|contradic|cannot\\s+determine|can't\\s+determine)"
    for _ in range(plan[2][1]):
        v1 = str(rnd.randint(1, 99))
        v2 = str(rnd.randint(1, 99))
        if v2 == v1:
            v2 = str(int(v1) + 1)
        ctx = _ev(f"Alpha = {v1}.", f"Alpha = {v2}.", "These two statements conflict.")
        add_case(
            {
                "type": "case",
                "id": f"standard_v2_{idx:05d}_contradiction",
                "category": "contradiction",
                "context": ctx,
                "question": "What is Alpha? If the evidence conflicts, do not choose a value.",
                "expected": {"answer_regex": contra_re},
                "checks": _chk_base(True)
                + [
                    {"type": "answer_regex", "field": "answer", "pattern": contra_re},
                    {
                        "type": "must_not_contain_regex",
                        "field": "answer",
                        "pattern": r"\\bAlpha\\s*(=|is)\\s*\\d+\\b",
                    },
                ],
            }
        )
        idx += 1
        if len(rows) - 1 >= size:
            return spec, rows

    # 4) Multi-turn consistency
    for _ in range(plan[3][1]):
        color = rnd.choice(colors)
        animal = rnd.choice(animals)
        ctx = _ev(
            f"User: I choose the color {color}.",
            f"Assistant: Noted. You chose {color}.",
            f"User: I also choose the animal {animal}.",
            f"Assistant: Noted. You chose {animal}.",
        )
        add_case(
            {
                "type": "case",
                "id": f"standard_v2_{idx:05d}_multiturn",
                "category": "multi_turn",
                "context": ctx,
                "question": "Repeat the chosen color exactly as stated earlier.",
                "expected": {"answer_exact": color},
                "checks": _chk_base(True)
                + [{"type": "answer_exact", "field": "answer", "value": color}],
            }
        )
        idx += 1
        if len(rows) - 1 >= size:
            return spec, rows

    # 5) Correct refusal / IDK (evidence insufficient)
    for _ in range(plan[4][1]):
        topic = rnd.choice(["Alpha", "Beta", "Gamma", "Delta", "Epsilon"])
        ctx = _ev("This evidence does not contain the requested fact.")
        add_case(
            {
                "type": "case",
                "id": f"standard_v2_{idx:05d}_idk",
                "category": "refusal_idk",
                "context": ctx,
                "question": f"What is the secret value of {topic}? If unknown, answer exactly I don't know.",
                "expected": {"answer_exact": "I don't know"},
                "checks": [
                    {"type": "json_schema"},
                    {"type": "answer_exact", "field": "answer", "value": "I don't know"},
                    {"type": "no_fake_evidence"},
                ],
            }
        )
        idx += 1
        if len(rows) - 1 >= size:
            return spec, rows

    # 6) Citation discipline: if you answer IDK, do not cite. If you answer, cite valid evidence.
    for _ in range(plan[5][1]):
        fact = str(rnd.randint(10, 999))
        ctx = _ev(f"Alpha is {fact}.", "Only answer using these facts.")
        # Alternate between answerable and unanswerable.
        if rnd.random() < 0.5:
            # Answerable and must cite.
            add_case(
                {
                    "type": "case",
                    "id": f"standard_v2_{idx:05d}_cite_ok",
                    "category": "citation",
                    "context": ctx,
                    "question": "What is Alpha? Return only the number.",
                    "expected": {"answer_exact": fact},
                    "checks": _chk_base(True)
                    + [{"type": "answer_exact", "field": "answer", "value": fact}],
                }
            )
        else:
            # Unanswerable and must not cite.
            add_case(
                {
                    "type": "case",
                    "id": f"standard_v2_{idx:05d}_cite_idk",
                    "category": "citation",
                    "context": ctx,
                    "question": "What is Beta? If unknown, answer exactly I don't know.",
                    "expected": {"answer_exact": "I don't know"},
                    "checks": [
                        {"type": "json_schema"},
                        {"type": "answer_exact", "field": "answer", "value": "I don't know"},
                        {"type": "no_fake_evidence"},
                    ],
                }
            )
        idx += 1
        if len(rows) - 1 >= size:
            return spec, rows

    return spec, rows


def write_standard_v2(path: Path, *, seed: int = 1337, size: int = 2500) -> None:
    _, rows = generate_standard_v2(seed=seed, size=size)
    _jsonl_write(path, rows)


if __name__ == "__main__":
    here = Path(__file__).resolve().parent
    out = here / "datasets" / "standard_v2.jsonl"
    write_standard_v2(out, seed=1337, size=2500)
    print(str(out))
