#!/usr/bin/env python3
"""Build the sub-25 MiB China-friendly H3 Studio bootstrap package."""

from __future__ import annotations

import shutil
import zipfile
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_OUTPUT = ROOT / "H3-Studio-OneClick-CN.zip"
MAX_PAGES_ASSET = 25 * 1024 * 1024

PROJECT_FILES = (
    "setup_h3_cn.cmd",
    "start_h3_studio.cmd",
    "start_h3.cmd",
    "installer/setup_h3_cn.cmd",
    "installer/install_h3_cn.ps1",
    "installer/requirements-cn.txt",
    "installer/requirements-studio.txt",
    "installer/README_FIRST.txt",
    "studio/h3_studio.py",
    "studio/deployment.py",
    "studio/repair.py",
    "studio/runtime.py",
    "studio/index.html",
    "tools/check_sageattention.py",
    "tools/build_h3_cn_installer.py",
    "tools/configure_cloud_access.py",
    "tools/launch_comfy.py",
    "tools/maxine_vsr.py",
    "tools/resolve_h3_runtime.cmd",
    "tools/prepare_736p_workflow.ps1",
    "tools/run_h3_736p_smoke.ps1",
    "tools/run_h3_r2v.ps1",
    "tools/run_h3_redraw_rtx.ps1",
    "tools/universal_video_finish.py",
)

COMFY_EXCLUDED_PREFIXES = (
    ".git/",
    "models/",
    "input/",
    "output/",
    "temp/",
    "user/",
    "custom_nodes/",
    "tests/",
    "tests-unit/",
    "__pycache__/",
)

CUSTOM_NODES = (
    "ComfyUI-MiniMax-H3-Turbo",
    "ComfyUI-PT_H3ConcatAVLatent",
    "Nvidia_RTX_Nodes_ComfyUI",
)


def normalized(path: Path, root: Path) -> str:
    return path.relative_to(root).as_posix()


def allowed_file(relative: str) -> bool:
    lowered = relative.lower()
    parts = lowered.split("/")
    return ".git" not in parts and "__pycache__" not in parts and not lowered.endswith((".pyc", ".pyo"))


def build_installer(output: Path = DEFAULT_OUTPUT) -> Path:
    output = output.resolve()
    output.parent.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(output, "w", compression=zipfile.ZIP_DEFLATED, compresslevel=9, allowZip64=True) as archive:
        for relative in PROJECT_FILES:
            path = ROOT / relative
            if not path.is_file():
                raise RuntimeError(f"一键安装器缺少文件：{relative}")
            archive.write(path, relative)

        comfy = ROOT / "ComfyUI"
        for path in comfy.rglob("*"):
            if not path.is_file():
                continue
            relative = normalized(path, comfy)
            if any(relative.startswith(prefix) for prefix in COMFY_EXCLUDED_PREFIXES):
                continue
            if allowed_file(relative):
                archive.write(path, f"ComfyUI/{relative}")

        custom_root = comfy / "custom_nodes"
        for node_name in CUSTOM_NODES:
            node_root = custom_root / node_name
            if not node_root.is_dir():
                raise RuntimeError(f"一键安装器缺少节点：{node_name}")
            for path in node_root.rglob("*"):
                if path.is_file():
                    relative = normalized(path, node_root)
                    if allowed_file(relative):
                        archive.write(path, f"ComfyUI/custom_nodes/{node_name}/{relative}")

        for directory in (
            "ComfyUI/input/",
            "ComfyUI/output/",
            "ComfyUI/models/diffusion_models/",
            "ComfyUI/models/text_encoders/",
            "ComfyUI/models/vae/",
            "ComfyUI/models/loras/",
        ):
            archive.writestr(directory, b"")

    size = output.stat().st_size
    if size >= MAX_PAGES_ASSET:
        output.unlink(missing_ok=True)
        raise RuntimeError(f"一键安装器 {size / 1024**2:.1f} MiB，超过 Cloudflare Pages 单文件 25 MiB 限制")
    return output


def main() -> int:
    result = build_installer()
    pages_copy = ROOT / "cloudflare-pages-package" / "downloads" / result.name
    pages_copy.parent.mkdir(parents=True, exist_ok=True)
    shutil.copy2(result, pages_copy)
    print(f"Installer: {result} ({result.stat().st_size / 1024**2:.1f} MiB)")
    print(f"Pages copy: {pages_copy}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
