#!/usr/bin/env python3
"""Independent integrity and metric audit for a CV-Bench-500 evaluation."""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import re
from collections import Counter, defaultdict
from pathlib import Path

from PIL import Image


EXPECTED_SELECTION_HASH = "da5ca19b4a65b4df7f45bea069db43c76e48d9a33adc5a224d1abc64dfce1b92"
EXPECTED_SOURCE_REVISION = "bc284db50d036958861cb60cdd7b77612052ce0d"
TASK_COUNTS = {
    ("2D", "Count"): 137,
    ("2D", "Relation"): 113,
    ("3D", "Depth"): 125,
    ("3D", "Distance"): 125,
}


def _jsonl(path: Path) -> list[dict]:
    return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line]


def _close(actual: float, expected: float) -> bool:
    return math.isclose(float(actual), float(expected), rel_tol=0.0, abs_tol=1e-12)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset-dir", type=Path, required=True)
    parser.add_argument("--eval-dir", type=Path, required=True)
    args = parser.parse_args()
    dataset = args.dataset_dir.resolve()
    evaluation = args.eval_dir.resolve()

    manifest = _jsonl(dataset / "validation.jsonl")
    summary = json.loads((dataset / "summary.json").read_text(encoding="utf-8"))
    results = _jsonl(evaluation / "results.jsonl")
    metrics_artifact = json.loads((evaluation / "metrics.json").read_text(encoding="utf-8"))
    reported = metrics_artifact["metrics"]

    if len(manifest) != 500 or len(results) != 500:
        raise AssertionError(f"Expected 500 manifest/results rows, got {len(manifest)}/{len(results)}")
    manifest_by_id = {row["sample_id"]: row for row in manifest}
    result_by_id = {row["sample_id"]: row for row in results}
    if len(manifest_by_id) != 500 or len(result_by_id) != 500:
        raise AssertionError("Duplicate sample_id in manifest or results")
    if set(manifest_by_id) != set(result_by_id):
        raise AssertionError("Manifest/result sample_id coverage differs")

    selection_items = [f"{row['config']}:{row['row_index']}" for row in manifest]
    selection_hash = hashlib.sha256("\n".join(selection_items).encode()).hexdigest()
    if selection_hash != EXPECTED_SELECTION_HASH or summary["selection_sha256"] != selection_hash:
        raise AssertionError(f"Selection fingerprint mismatch: {selection_hash}")
    if summary["source_revision"] != EXPECTED_SOURCE_REVISION:
        raise AssertionError(f"Unexpected CV-Bench source revision: {summary['source_revision']}")

    image_hashes: set[str] = set()
    manifest_counts = Counter()
    for row in manifest:
        key = (str(row["config"]), str(row["task"]))
        manifest_counts[key] += 1
        image_path = (dataset / row["image"]).resolve()
        if dataset not in image_path.parents or not image_path.is_file():
            raise FileNotFoundError(f"Invalid image path: {image_path}")
        raw = image_path.read_bytes()
        image_hashes.add(hashlib.sha256(raw).hexdigest())
        with Image.open(image_path) as image:
            if image.size != (int(row["image_width"]), int(row["image_height"])):
                raise AssertionError(f"Image dimension mismatch: {row['sample_id']}")
            image.verify()
        expected_match = re.fullmatch(r"\(([A-Z])\)", str(row["answer"]))
        if not expected_match:
            raise AssertionError(f"Malformed answer: {row['sample_id']}={row['answer']!r}")
        expected_index = ord(expected_match.group(1)) - ord("A")
        if expected_index >= len(row["choices"]):
            raise AssertionError(f"Answer outside choices: {row['sample_id']}")
    if manifest_counts != Counter(TASK_COUNTS) or len(image_hashes) != 500:
        raise AssertionError(
            f"Dataset coverage mismatch: tasks={dict(manifest_counts)}, unique_images={len(image_hashes)}"
        )

    shard_rows: list[dict] = []
    for shard_path in sorted(evaluation.glob("rank_*.jsonl")):
        shard_rows.extend(_jsonl(shard_path))
    if len(shard_rows) != 500 or {row["sample_id"] for row in shard_rows} != set(result_by_id):
        raise AssertionError("Per-rank shards do not cover the final result set exactly")

    grouped: dict[tuple[str, str], list[dict]] = defaultdict(list)
    correct_total = 0
    parsed_total = 0
    error_total = 0
    for sample_id, result in result_by_id.items():
        source = manifest_by_id[sample_id]
        expected = re.fullmatch(r"\(([A-Z])\)", source["answer"]).group(1)
        if result["config"] != source["config"] or result["task"] != source["task"]:
            raise AssertionError(f"Result grouping mismatch: {sample_id}")
        if result["expected"] != expected:
            raise AssertionError(f"Result label mismatch: {sample_id}")
        error_total += int(result["error"] is not None)
        parsed_total += int(result["predicted"] is not None)
        if result["predicted"] is not None and result["predicted"] not in [
            chr(ord("A") + i) for i in range(len(source["choices"]))
        ]:
            raise AssertionError(f"Predicted option outside choices: {sample_id}")
        independently_correct = result["predicted"] == expected
        if bool(result["correct"]) != independently_correct:
            raise AssertionError(f"Incorrect correctness flag: {sample_id}")
        correct_total += int(independently_correct)
        grouped[(source["config"], source["task"])].append(result)

    task_numerators = {
        key: sum(bool(row["correct"]) for row in rows) for key, rows in grouped.items()
    }
    task_accuracy = {key: task_numerators[key] / TASK_COUNTS[key] for key in TASK_COUNTS}
    independently_computed = {
        "eval/cvbench_samples": 500.0,
        "eval/cvbench_errors": float(error_total),
        "eval/cvbench_parse_rate": parsed_total / 500,
        "eval/cvbench_accuracy": correct_total / 500,
        "eval/cvbench_2d_accuracy": (
            task_numerators[("2D", "Count")] + task_numerators[("2D", "Relation")]
        ) / 250,
        "eval/cvbench_3d_accuracy": (
            task_numerators[("3D", "Depth")] + task_numerators[("3D", "Distance")]
        ) / 250,
        "eval/cvbench_2d_count_accuracy": task_accuracy[("2D", "Count")],
        "eval/cvbench_2d_relation_accuracy": task_accuracy[("2D", "Relation")],
        "eval/cvbench_3d_depth_accuracy": task_accuracy[("3D", "Depth")],
        "eval/cvbench_3d_distance_accuracy": task_accuracy[("3D", "Distance")],
    }
    independently_computed["eval/cvbench_2d_task_macro_accuracy"] = (
        task_accuracy[("2D", "Count")] + task_accuracy[("2D", "Relation")]
    ) / 2
    independently_computed["eval/cvbench_3d_task_macro_accuracy"] = (
        task_accuracy[("3D", "Depth")] + task_accuracy[("3D", "Distance")]
    ) / 2
    independently_computed["eval/cvbench_task_macro_accuracy"] = sum(task_accuracy.values()) / 4

    for key, expected in independently_computed.items():
        if key not in reported or not _close(reported[key], expected):
            raise AssertionError(f"Metric mismatch {key}: reported={reported.get(key)}, expected={expected}")

    swanlab = json.loads((evaluation / "swanlab.json").read_text(encoding="utf-8"))
    if not swanlab.get("enabled") or not swanlab.get("run_id") or swanlab.get("error"):
        raise AssertionError(f"SwanLab logging is incomplete: {swanlab}")

    report = {
        "assessment": "ready_to_share",
        "dataset_rows": len(manifest),
        "unique_sample_ids": len(manifest_by_id),
        "unique_image_bytes": len(image_hashes),
        "selection_sha256": selection_hash,
        "source_revision": summary["source_revision"],
        "task_denominators": {f"{key[0]}/{key[1]}": value for key, value in TASK_COUNTS.items()},
        "task_numerators": {f"{key[0]}/{key[1]}": task_numerators[key] for key in TASK_COUNTS},
        "correct_total": correct_total,
        "parsed_total": parsed_total,
        "error_total": error_total,
        "metrics_verified": independently_computed,
        "swanlab_run_id": swanlab["run_id"],
    }
    (evaluation / "validation_report.json").write_text(
        json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    print(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()