"""Restore the two upstream HF cache entries used by the censor workers."""

import argparse
import hashlib
import json
from pathlib import Path
import shutil


def main():
    from huggingface_hub.constants import HF_HUB_CACHE

    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--repo-dir", type=Path, default=Path(__file__).resolve().parent)
    parser.add_argument("--cache-dir", type=Path, default=Path(HF_HUB_CACHE))
    args = parser.parse_args()
    root = args.repo_dir.resolve()
    manifest = json.loads((root / "model-manifest.json").read_text(encoding="utf-8"))
    models = {f["path"]: f for f in manifest["files"]}
    for entry in manifest["censor_hf_cache"]:
        source = root / entry["path"]
        expected = models[entry["path"]]
        if source.stat().st_size != expected["size_bytes"]:
            raise RuntimeError("File size mismatch: " + str(source))
        digest = hashlib.sha256(source.read_bytes()).hexdigest()
        if digest != expected["sha256"]:
            raise RuntimeError("SHA256 mismatch: " + str(source))
        cache = args.cache_dir / ("models--" + entry["repo_id"].replace("/", "--"))
        blob = cache / "blobs" / digest
        blob.parent.mkdir(parents=True, exist_ok=True)
        if blob.exists():
            if hashlib.sha256(blob.read_bytes()).hexdigest() != digest:
                raise RuntimeError("Existing cache blob is corrupt: " + str(blob))
        else:
            shutil.copyfile(source, blob)
        snapshot = cache / "snapshots" / entry["revision"] / entry["filename"]
        snapshot.parent.mkdir(parents=True, exist_ok=True)
        if snapshot.exists():
            if hashlib.sha256(snapshot.read_bytes()).hexdigest() != digest:
                raise RuntimeError("Existing cache snapshot differs: " + str(snapshot))
        else:
            snapshot.symlink_to(blob.resolve())
        ref = cache / "refs" / "main"
        ref.parent.mkdir(parents=True, exist_ok=True)
        ref.write_text(entry["revision"], encoding="utf-8")
        print("Restored " + entry["repo_id"] + " @ " + entry["revision"])
    print("Set CENSOR_MODELS_DIR=" + str(root / "auxiliary" / "censor") + " before starting ComfyUI.")
    print("The cache directory must be the one used by the ComfyUI service account.")


if __name__ == "__main__":
    main()