#!/usr/bin/env python3
"""Validate committed SearchGen-Bench artifacts without access to ToolGen."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from types import SimpleNamespace

from build_leaderboard_data import build_aggregates


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--data-dir",
        type=Path,
        default=Path(__file__).resolve().parents[1] / "public" / "data",
    )
    return parser.parse_args()


def load_json(path: Path):
    return json.loads(path.read_text())


def main() -> None:
    data_dir = parse_args().data_dir.resolve()
    manifest = load_json(data_dir / "manifest.json")
    records = [
        json.loads(line)
        for line in (data_dir / "prompt_scores.jsonl").read_text().splitlines()
        if line.strip()
    ]
    errors = []
    if len(records) != manifest["dataset"]["n_prompts"]:
        errors.append("prompt count does not match manifest")
    if len({record["sample_id"] for record in records}) != len(records):
        errors.append("prompt IDs are not unique")

    expected_partition = manifest["partition"]
    actual_partition = {
        "NoSearch": sum(record["stratum"] == "NoSearch" for record in records),
        "SearchIntensive": sum(record["stratum"] == "SearchIntensive" for record in records),
        "VisualSearch": sum(record["search_type"] == "VisualSearch" for record in records),
        "TextualSearch": sum(record["search_type"] == "TextualSearch" for record in records),
    }
    if actual_partition != expected_partition:
        errors.append(f"partition mismatch: {actual_partition}")

    model_ids = list(manifest["models"])
    for record in records:
        if set(record["models"]) != set(model_ids):
            errors.append(f"model set mismatch in {record['sample_id']}")
        for model_id, result in record["models"].items():
            if result["status"] != "scored":
                continue
            for component, value in result["components_raw_0to3"].items():
                if value is not None and not 0 <= value <= 3:
                    errors.append(f"out-of-range value in {record['sample_id']} {model_id} {component}")

    canonical_stub = SimpleNamespace(
        TABLE1_GENS=model_ids,
        SKIP_MISSING_GENS=set(manifest["scoring"]["missing_policy_exceptions"]),
    )
    regenerated = build_aggregates(records, canonical_stub)
    expected_files = {
        "leaderboard_overall.json": regenerated["overall"],
        "leaderboard_by_stratum.json": regenerated["strata"],
        "leaderboard_by_domain.json": regenerated["domains"],
        "leaderboard_by_failure_mode.json": regenerated["failure_modes"],
    }
    for filename, expected in expected_files.items():
        if load_json(data_dir / filename) != expected:
            errors.append(f"{filename} does not match prompt-level recomputation")

    if errors:
        raise SystemExit("Artifact validation failed:\n- " + "\n- ".join(errors[:50]))
    print(f"Validated {len(records)} prompts, {len(model_ids)} models, and four aggregate artifacts")


if __name__ == "__main__":
    main()