from __future__ import annotations

import argparse
import hashlib
import json
from dataclasses import asdict
from pathlib import Path
from typing import Any

from training.recipes import all_recipes, get_recipe
from training.recipes.base import ROOT, SPLITS, Case, Recipe, Split

DEFAULT_RECIPE = "capability-placement"


def validate_cases(recipe: Recipe, cases: dict[Split, tuple[Case, ...]]) -> None:
    if set(cases) != set(SPLITS):
        raise ValueError(f"{recipe.id} does not contain every split")
    seen_ids: set[str] = set()
    seen_slugs: set[str] = set()
    for split in SPLITS:
        rows = cases[split]
        expected_count = len(recipe.kinds) * recipe.split_counts[split]
        if len(rows) != expected_count:
            raise ValueError(f"{recipe.id}/{split} has {len(rows)} rows, expected {expected_count}")
        kinds = {row.kind for row in rows}
        if kinds != set(recipe.kinds):
            raise ValueError(f"{recipe.id}/{split} does not cover every kind")
        for kind in recipe.kinds:
            count = sum(row.kind == kind for row in rows)
            if count != recipe.split_counts[split]:
                raise ValueError(f"{recipe.id}/{split}/{kind} has {count} rows")
        for row in rows:
            if row.split != split:
                raise ValueError(f"split drift in {row.id}")
            if row.id in seen_ids or row.slug in seen_slugs:
                raise ValueError(f"duplicate id or slug in {recipe.id}: {row.id} / {row.slug}")
            seen_ids.add(row.id)
            seen_slugs.add(row.slug)
            if not row.prompt.strip() or not row.target:
                raise ValueError(f"empty prompt or target in {row.id}")
            if not recipe.validate_target(row, row.target):
                raise ValueError(f"target schema drift in {recipe.id}/{row.id}")


def _jsonl_bytes(rows: tuple[Case, ...]) -> bytes:
    return ("".join(json.dumps(asdict(row), separators=(",", ":")) + "\n" for row in rows)).encode()


def rendered_files(recipe: Recipe) -> dict[Path, bytes]:
    cases = recipe.generate_cases()
    validate_cases(recipe, cases)
    files = {recipe.data_dir / f"{split}.jsonl": _jsonl_bytes(cases[split]) for split in SPLITS}
    test_fixture = [
        {"id": row.id, "prompt": row.prompt, "expected": row.target} for row in cases["test"]
    ]
    files[recipe.eve_eval_path] = (
        json.dumps(test_fixture, indent=2, ensure_ascii=False) + "\n"
    ).encode()
    return files


def dataset_sha256(recipe: Recipe) -> str:
    digest = hashlib.sha256()
    for path, content in sorted(rendered_files(recipe).items(), key=lambda item: str(item[0])):
        if path == recipe.eve_eval_path:
            continue
        digest.update(path.name.encode())
        digest.update(b"\0")
        digest.update(content)
    return digest.hexdigest()


def write_files(recipes: tuple[Recipe, ...], check: bool) -> dict[str, str]:
    mismatches: list[str] = []
    hashes: dict[str, str] = {}
    for recipe in recipes:
        for path, content in rendered_files(recipe).items():
            if check:
                if not path.exists() or path.read_bytes() != content:
                    mismatches.append(str(path.relative_to(ROOT)))
                continue
            path.parent.mkdir(parents=True, exist_ok=True)
            path.write_bytes(content)
        hashes[recipe.id] = dataset_sha256(recipe)
    if mismatches:
        raise SystemExit("generated data is stale: " + ", ".join(mismatches))
    return hashes


def load_split(recipe: Recipe, split: Split) -> tuple[Case, ...]:
    rows: list[Case] = []
    path = recipe.data_dir / f"{split}.jsonl"
    for line in path.read_text().splitlines():
        payload: dict[str, Any] = json.loads(line)
        rows.append(Case(**payload))
    return tuple(rows)


def parser() -> argparse.ArgumentParser:
    result = argparse.ArgumentParser(description="Generate deterministic Eve recipe corpora")
    selection = result.add_mutually_exclusive_group()
    selection.add_argument(
        "--recipe",
        choices=tuple(recipe.id for recipe in all_recipes()),
        help=f"one recipe to generate; defaults to {DEFAULT_RECIPE}",
    )
    selection.add_argument("--all", action="store_true", help="generate every recipe")
    action = result.add_mutually_exclusive_group(required=True)
    action.add_argument("--write", action="store_true", help="write generated files")
    action.add_argument("--check", action="store_true", help="fail when generated files are stale")
    return result


def main() -> None:
    args = parser().parse_args()
    recipes = all_recipes() if args.all else (get_recipe(args.recipe or DEFAULT_RECIPE),)
    hashes = write_files(recipes, check=args.check)
    for recipe_id, digest in hashes.items():
        print(f"{recipe_id} dataset_sha256={digest}")


if __name__ == "__main__":
    main()