"""Download only pinned diagnostic assets into owned ignored data."""
import argparse
import hashlib
import json
from pathlib import Path
import time
import urllib.request
import uuid

ROOT = Path(__file__).resolve().parents[2]
DESTINATION = ROOT / ".searchi" / "rerank-eval" / "model"


def digest(path):
    checksum = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(65536), b""):
            checksum.update(chunk)
    return checksum.hexdigest()


def main():
    parser=argparse.ArgumentParser()
    parser.add_argument("--reference",action="store_true")
    args=parser.parse_args()
    name="reference-manifest.json" if args.reference else "model-manifest.json"
    manifest = json.loads(Path(__file__).with_name(name).read_text(encoding="utf-8"))
    for item in manifest["files"]:
        destination = DESTINATION / item["path"]
        if destination.is_file() and destination.stat().st_size == item["bytes"] and digest(destination) == item["sha256"]:
            print("verified cache", item["path"], flush=True)
            continue
        destination.parent.mkdir(parents=True, exist_ok=True)
        temporary = destination.with_name(destination.name + "." + uuid.uuid4().hex + ".partial")
        total = 0
        start = time.monotonic()
        try:
            with urllib.request.urlopen(item["url"], timeout=20) as response, temporary.open("xb") as output:
                while chunk := response.read(65536):
                    total += len(chunk)
                    if total > item["bytes"] or time.monotonic() - start > 180:
                        raise ValueError("Download exceeded size/deadline")
                    output.write(chunk)
            if total != item["bytes"] or digest(temporary) != item["sha256"]:
                raise ValueError("Pinned asset hash/size mismatch")
            temporary.replace(destination)
            print("verified download", item["path"], total, flush=True)
        finally:
            temporary.unlink(missing_ok=True)


if __name__ == "__main__":
    main()