"""Read-only upstream pins; all downloaded data lives in this new app directory."""
import hashlib
import json
import shutil
import subprocess
from pathlib import Path

from config import FILES, ROOT, TOTAL_BYTES

COMFY_SHA = "5cc4b870094c89090db54bd8d12649dcaf89f2ba"
MOSS_SHA = "c1355e6b1986dfe8fba3c0455783a05899212342"
COMFY_ROOT = ROOT / "_runtime" / "ComfyUI"
MOSS_ROOT = ROOT / "_runtime" / "MOSS-Image"
MODELS_ROOT = ROOT / "model_files"


def checkout(path, url, revision):
    if path.exists():
        current = subprocess.check_output(
            ["git", "-C", str(path), "rev-parse", "HEAD"], text=True
        ).strip()
        if current != revision:
            raise RuntimeError(f"既存checkoutのrevisionが異なります: {path}")
        return
    path.parent.mkdir(parents=True, exist_ok=True)
    subprocess.run(["git", "init", str(path)], check=True)
    subprocess.run(["git", "-C", str(path), "remote", "add", "origin", url], check=True)
    subprocess.run(["git", "-C", str(path), "fetch", "--depth", "1", "origin", revision], check=True)
    subprocess.run(["git", "-C", str(path), "checkout", "--detach", "FETCH_HEAD"], check=True)


def verify(path, item):
    if path.stat().st_size != item["size"]:
        raise RuntimeError(f"ファイルサイズ不一致: {path.name}")
    stamp = {"size": path.stat().st_size, "mtime_ns": path.stat().st_mtime_ns,
             "sha256": item["sha256"]}
    receipt = path.with_suffix(path.suffix + ".verified.json")
    if receipt.exists():
        try:
            if json.loads(receipt.read_text()) == stamp:
                return
        except (OSError, ValueError):
            pass
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(8 * 1024 * 1024), b""):
            digest.update(block)
    if digest.hexdigest() != item["sha256"]:
        raise RuntimeError(f"SHA256不一致: {path.name}")
    receipt.write_text(json.dumps(stamp), encoding="utf-8")


def prepare():
    from huggingface_hub import hf_hub_download

    if TOTAL_BYTES > 48_000_000_000:
        raise RuntimeError("モデル容量が48GBを超えています。")
    MODELS_ROOT.mkdir(exist_ok=True)
    needed = sum(item["size"] for item in FILES
                 if not (MODELS_ROOT / item["path"]).is_file())
    if shutil.disk_usage(ROOT).free < needed + 3_000_000_000:
        raise RuntimeError("モデルと一時ファイル用の空きディスクが不足しています。")
    print(f"[bootstrap] selected files: {TOTAL_BYTES / 1e9:.2f} GB; missing: {needed / 1e9:.2f} GB", flush=True)
    checkout(COMFY_ROOT, "https://github.com/Comfy-Org/ComfyUI.git", COMFY_SHA)
    checkout(MOSS_ROOT, "https://github.com/OpenMOSS/MOSS-Image.git", MOSS_SHA)
    for item in FILES:
        print(f"[bootstrap] fetch/verify {item['id']} ({item['size'] / 1e9:.3f} GB)", flush=True)
        # Fetch only pinned files, with separate directories for the two templates.
        local_dir = MODELS_ROOT / Path(item["path"]).parent if "source_path" in item else MODELS_ROOT
        path = Path(hf_hub_download(item["repo"], item.get("source_path", item["path"]),
                                   revision=item["revision"], local_dir=local_dir))
        verify(path, item)
        print(f"[bootstrap] verified {item['id']}", flush=True)
    print("[bootstrap] all selected files verified", flush=True)


if __name__ == "__main__":
    prepare()