#!/usr/bin/env python3
"""Build a deterministic 500-example CV-Bench validation subset.

The source parquet files are fetched without any proxy, but only 250 examples
from each of the 2D and 3D configurations are materialized. Sampling is
stratified by the dataset's task column and is reproducible from ``--seed``.
"""

from __future__ import annotations

import argparse
import hashlib
import io
import json
import os
import random
import shutil
import subprocess
from collections import Counter, defaultdict
from datetime import datetime, timezone
from pathlib import Path

import pyarrow as pa
import pyarrow.parquet as pq
from PIL import Image


SOURCE_REPO = "nyu-visionx/CV-Bench"
SOURCE_REVISION = "bc284db50d036958861cb60cdd7b77612052ce0d"
SOURCE_FILES = {
    "2D": "test_2d.parquet",
    "3D": "test_3d.parquet",
}
EXPECTED_ROWS = {"2D": 1438, "3D": 1200}
PROXY_ENV_VARS = (
    "http_proxy",
    "https_proxy",
    "HTTP_PROXY",
    "HTTPS_PROXY",
    "all_proxy",
    "ALL_PROXY",
)
IMAGE_SUFFIXES = {
    "JPEG": ".jpg",
    "PNG": ".png",
    "WEBP": ".webp",
    "BMP": ".bmp",
    "GIF": ".gif",
    "TIFF": ".tiff",
}


def direct_environment() -> dict[str, str]:
    environment = os.environ.copy()
    for name in PROXY_ENV_VARS:
        environment.pop(name, None)
    return environment


def download_direct(url: str, destination: Path, retries: int = 5) -> None:
    if destination.is_file() and destination.stat().st_size > 0:
        print(f"[cvbench] reuse {destination} ({destination.stat().st_size} bytes)")
        return
    destination.parent.mkdir(parents=True, exist_ok=True)
    part = destination.with_suffix(destination.suffix + ".part")
    print(f"[cvbench] direct resumable download {url}", flush=True)
    subprocess.run(
        [
            "curl",
            "--noproxy",
            "*",
            "--fail",
            "--location",
            "--retry",
            str(retries),
            "--retry-all-errors",
            "--retry-delay",
            "2",
            "--connect-timeout",
            "30",
            "--speed-limit",
            "1024",
            "--speed-time",
            "30",
            "--continue-at",
            "-",
            "--output",
            str(part),
            url,
        ],
        env=direct_environment(),
        check=True,
    )
    if not part.is_file() or part.stat().st_size == 0:
        raise IOError(f"download did not produce data: {part}")
    part.replace(destination)


def proportional_quotas(groups: dict[str, list[int]], total: int) -> dict[str, int]:
    population = sum(len(rows) for rows in groups.values())
    exact = {name: total * len(rows) / population for name, rows in groups.items()}
    quotas = {name: int(value) for name, value in exact.items()}
    remainder = total - sum(quotas.values())
    order = sorted(groups, key=lambda name: (-(exact[name] - quotas[name]), name))
    for name in order[:remainder]:
        quotas[name] += 1
    return quotas


def stable_rng(seed: int, config: str, task: str) -> random.Random:
    digest = hashlib.sha256(f"{seed}:{config}:{task}".encode()).digest()
    return random.Random(int.from_bytes(digest[:8], "big"))


def select_rows(
    parquet_path: Path,
    config: str,
    count: int,
    seed: int,
    seen_image_hashes: set[bytes],
) -> list[int]:
    metadata = pq.read_table(parquet_path, columns=["type", "task", "image"])
    if metadata.num_rows != EXPECTED_ROWS[config]:
        raise ValueError(
            f"unexpected {config} row count: {metadata.num_rows} != {EXPECTED_ROWS[config]}"
        )
    groups: dict[str, list[int]] = defaultdict(list)
    for row_index, task in enumerate(metadata.column("task").to_pylist()):
        groups[str(task)].append(row_index)
    quotas = proportional_quotas(groups, count)
    selected: list[int] = []
    duplicate_images_skipped = 0
    for task in sorted(groups):
        candidates = groups[task].copy()
        stable_rng(seed, config, task).shuffle(candidates)
        task_selected = 0
        for row_index in candidates:
            image_value = metadata.column("image")[row_index].as_py()
            raw = image_value.get("bytes")
            if raw is None:
                raise ValueError(f"parquet image has no embedded bytes at {config} row {row_index}")
            image_hash = hashlib.sha256(raw).digest()
            if image_hash in seen_image_hashes:
                duplicate_images_skipped += 1
                continue
            seen_image_hashes.add(image_hash)
            selected.append(row_index)
            task_selected += 1
            if task_selected == quotas[task]:
                break
        if task_selected != quotas[task]:
            raise ValueError(
                f"not enough unique {config}/{task} images: {task_selected} != {quotas[task]}"
            )
    selected.sort()
    if len(selected) != count or len(set(selected)) != count:
        raise AssertionError(f"invalid {config} selection size")
    print(f"[cvbench] {config} task quotas: {dict(sorted(quotas.items()))}")
    print(f"[cvbench] {config} duplicate candidate images skipped: {duplicate_images_skipped}")
    return selected


def image_bytes_and_suffix(image_value: dict) -> tuple[bytes, str, int, int]:
    raw = image_value.get("bytes")
    if raw is None:
        raise ValueError(f"parquet image has no embedded bytes: {image_value.get('path')!r}")
    with Image.open(io.BytesIO(raw)) as image:
        image.verify()
    with Image.open(io.BytesIO(raw)) as image:
        fmt = str(image.format or "").upper()
        width, height = image.size
    suffix = IMAGE_SUFFIXES.get(fmt, ".img")
    return raw, suffix, width, height


def materialize_config(
    parquet_path: Path,
    config: str,
    selected: list[int],
    staging: Path,
) -> list[dict]:
    columns = [
        "idx",
        "type",
        "task",
        "image",
        "question",
        "choices",
        "answer",
        "prompt",
        "filename",
        "source",
        "source_dataset",
        "source_filename",
        "target_class",
        "target_size",
        "bbox",
    ]
    table = pq.read_table(parquet_path, columns=columns)
    sampled = table.take(pa.array(selected, type=pa.int64())).to_pylist()
    images_dir = staging / "images"
    images_dir.mkdir(parents=True, exist_ok=True)
    records: list[dict] = []
    for ordinal, (row_index, row) in enumerate(zip(selected, sampled)):
        raw, suffix, width, height = image_bytes_and_suffix(row.pop("image"))
        image_name = f"{config}_{ordinal:03d}_row{row_index:04d}{suffix}"
        image_rel = Path("images") / image_name
        (staging / image_rel).write_bytes(raw)
        record = {
            "sample_id": f"{config}-{ordinal:03d}",
            "config": config,
            "split": "test",
            "row_index": row_index,
            "image": image_rel.as_posix(),
            "image_width": width,
            "image_height": height,
            **row,
        }
        records.append(record)
    return records


def validate_existing(output: Path, expected: int) -> bool:
    manifest = output / "validation.jsonl"
    summary = output / "summary.json"
    if not manifest.is_file() or not summary.is_file():
        return False
    records = [json.loads(line) for line in manifest.read_text().splitlines() if line.strip()]
    if len(records) != expected:
        return False
    return all((output / record["image"]).is_file() for record in records)


def main() -> None:
    root = Path(__file__).resolve().parents[1]
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, default=root / "data" / "cvbench_500")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--per-config", type=int, default=250)
    parser.add_argument("--keep-parquet", action="store_true")
    args = parser.parse_args()

    output = args.output.resolve()
    expected = 2 * args.per_config
    if validate_existing(output, expected):
        print(f"[cvbench] validated existing subset at {output}")
        return
    staging = output.with_name(output.name + ".tmp")
    cache = output.with_name(output.name + ".parquet-cache")
    if staging.exists():
        shutil.rmtree(staging)
    staging.mkdir(parents=True)
    cache.mkdir(parents=True, exist_ok=True)

    records: list[dict] = []
    selection_items: list[str] = []
    seen_image_hashes: set[bytes] = set()
    for config, filename in SOURCE_FILES.items():
        parquet_path = cache / filename
        url = (
            f"https://huggingface.co/datasets/{SOURCE_REPO}/resolve/"
            f"{SOURCE_REVISION}/{filename}"
        )
        download_direct(url, parquet_path)
        selected = select_rows(
            parquet_path,
            config,
            args.per_config,
            args.seed,
            seen_image_hashes,
        )
        records.extend(materialize_config(parquet_path, config, selected, staging))
        selection_items.extend(f"{config}:{row_index}" for row_index in selected)

    records.sort(key=lambda row: (row["config"], row["row_index"]))
    manifest = staging / "validation.jsonl"
    with manifest.open("w", encoding="utf-8") as handle:
        for record in records:
            handle.write(json.dumps(record, ensure_ascii=False) + "\n")

    by_type = Counter(record["config"] for record in records)
    by_task = Counter(f"{record['config']}/{record['task']}" for record in records)
    answer_counts = Counter(record["answer"] for record in records)
    summary = {
        "source_repo": SOURCE_REPO,
        "source_revision": SOURCE_REVISION,
        "source_split": "test",
        "seed": args.seed,
        "sampling": "task-stratified proportional without replacement; unique image bytes",
        "rows": len(records),
        "unique_images": len(seen_image_hashes),
        "counts_by_type": dict(sorted(by_type.items())),
        "counts_by_task": dict(sorted(by_task.items())),
        "counts_by_answer": dict(sorted(answer_counts.items())),
        "selection_sha256": hashlib.sha256("\n".join(selection_items).encode()).hexdigest(),
        "created_utc": datetime.now(timezone.utc).isoformat(),
    }
    (staging / "summary.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    if output.exists():
        raise FileExistsError(
            f"incomplete output already exists at {output}; move it aside before rerunning"
        )
    staging.replace(output)
    if not args.keep_parquet:
        shutil.rmtree(cache)
    if not validate_existing(output, expected):
        raise RuntimeError("materialized CV-Bench subset failed validation")
    print(json.dumps(summary, ensure_ascii=False, indent=2))
    print(f"[cvbench] ready at {output}")


if __name__ == "__main__":
    main()