"""Download and verify the public tau-bench inputs declared in study_config.json."""

from __future__ import annotations

import argparse
import hashlib
import json
import shutil
import urllib.request
from pathlib import Path

HERE = Path(__file__).resolve().parent


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", required=True, type=Path)
    args = parser.parse_args()
    config = json.loads((HERE / "study_config.json").read_text(encoding="utf-8"))
    args.output.mkdir(parents=True, exist_ok=True)
    for source in config["sources"]:
        destination = args.output / source["local_filename"]
        if destination.exists() and sha256_file(destination) == source["sha256"]:
            print(f"verified existing {destination.name}")
            continue
        if destination.exists():
            raise ValueError(f"refusing to overwrite hash-mismatched file: {destination}")
        temporary = destination.with_suffix(destination.suffix + ".part")
        try:
            with urllib.request.urlopen(source["url"], timeout=120) as response:
                with temporary.open("wb") as handle:
                    shutil.copyfileobj(response, handle, length=1024 * 1024)
            actual = sha256_file(temporary)
            if actual != source["sha256"]:
                raise ValueError(f"download hash mismatch for {source['source_id']}: {actual}")
            temporary.replace(destination)
            print(f"downloaded and verified {destination.name}")
        finally:
            if temporary.exists():
                temporary.unlink()


if __name__ == "__main__":
    main()
