#!/usr/bin/env python3
"""Download the pinned Supertonic 3 CPU model. No packages or credentials required."""
import hashlib
import json
from pathlib import Path
import sys
import urllib.request

ROOT = Path(__file__).resolve().parent.parent
MANIFEST = json.loads((ROOT / "scripts/supertonic-3-model.json").read_text())


def valid(path, entry):
    if not path.is_file() or path.stat().st_size != entry["size"]:
        return False
    digest = hashlib.sha256() if "sha256" in entry else hashlib.sha1()
    if "git_blob" in entry:
        digest.update(f"blob {entry['size']}\0".encode())
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest() == entry.get("sha256", entry.get("git_blob"))


def download(manifest, destination):
    total = sum(entry["size"] for entry in manifest["files"])
    print(f"{manifest['repository']}: {total / 1e6:.1f} MB; destination: {destination}", flush=True)
    base = f"https://huggingface.co/{manifest['repository']}/resolve/{manifest['revision']}"
    for entry in manifest["files"]:
        target = destination / entry["path"]
        if valid(target, entry):
            print(f"Verified {entry['path']}", flush=True)
            continue
        target.parent.mkdir(parents=True, exist_ok=True)
        temporary = target.with_suffix(target.suffix + ".part")
        print(f"Downloading {entry['path']} ({entry['size'] / 1e6:.1f} MB)", flush=True)
        try:
            with urllib.request.urlopen(f"{base}/{entry['path']}", timeout=60) as source, temporary.open("wb") as output:
                while chunk := source.read(1024 * 1024):
                    output.write(chunk)
            if not valid(temporary, entry):
                raise RuntimeError(f"Checksum mismatch: {entry['path']}")
            temporary.replace(target)
        finally:
            temporary.unlink(missing_ok=True)


def main():
    destination = Path(sys.argv[1]) if len(sys.argv) > 1 else ROOT / ".local/tts/supertonic-3"
    print("Supertonic model license: OpenRAIL-M (saved as LICENSE alongside the model).", flush=True)
    download(MANIFEST, destination)
    print("Ready. Run ./gradlew :composeApp:run and use Read aloud in a book.", flush=True)


if __name__ == "__main__":
    main()