"""Explicitly provision the pinned Google model; search never calls this script."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import urllib.request


def verified(path, spec):
    if not path.is_file() or path.stat().st_size != spec["bytes"]:
        return False
    with path.open("rb") as source:
        return hashlib.file_digest(source, "sha256").hexdigest() == spec["sha256"]


def main():
    parser = argparse.ArgumentParser(description="Download the verified Google EmbeddingGemma 2 Q8 model (865 MB)")
    parser.add_argument("--model-dir", type=Path, default=Path(".searchi/models/embeddinggemma2"))
    parser.add_argument("--text-only", action="store_true", help="Download only the 310 MB text backbone")
    args = parser.parse_args()
    manifest = json.loads(Path(__file__).with_name("model-manifest.json").read_text())
    for spec in manifest["files"][:1] if args.text_only else manifest["files"]:
        destination = args.model_dir / spec["path"]
        if verified(destination, spec):
            print(f"Verified {spec['path']}", flush=True)
            continue
        destination.parent.mkdir(parents=True, exist_ok=True)
        temporary = destination.with_suffix(".part")
        url = f"https://huggingface.co/{manifest['export_repository']}/resolve/{manifest['revision']}/{spec['path']}"
        try:
            request = urllib.request.Request(url, headers={"User-Agent": "Searchi/0.1"})
            print(f"Downloading {spec['path']} ({spec['bytes']} bytes)", flush=True)
            with urllib.request.urlopen(request, timeout=60) as response, temporary.open("wb") as output:
                total = 0
                last_report = 0
                while chunk := response.read(1024 * 1024):
                    total += len(chunk)
                    if total > spec["bytes"]:
                        raise ValueError("Model download exceeds pinned size")
                    output.write(chunk)
                    if total - last_report >= 50 * 1024 * 1024:
                        print(f"Downloaded {total}/{spec['bytes']} bytes", flush=True)
                        last_report = total
            if not verified(temporary, spec):
                raise ValueError("Model size or SHA256 mismatch")
            os.replace(temporary, destination)
            print(f"Verified SHA256: {spec['sha256']}", flush=True)
        finally:
            temporary.unlink(missing_ok=True)


if __name__ == "__main__":
    main()