"""Prepare the ComfyUI runtime used by the VoiceGate Space.

The script is intentionally idempotent. By default it prepares ComfyUI and
custom node repositories, but it does not download large model weights unless
explicitly requested.
"""

from __future__ import annotations

import argparse
import hashlib
import os
import shutil
import subprocess
import sys
from dataclasses import dataclass
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
COMFY_DIR = ROOT / "ComfyUI"
CUSTOM_NODES_DIR = COMFY_DIR / "custom_nodes"
DEFAULT_PERSISTENT_MODEL_ROOT = Path("/data/voicegate_models")
MELBAND_FILENAME = "MelBandRoformer_fp32.safetensors"
MELBAND_SIZE = 912_885_656
MELBAND_SHA256 = "450caec8e8e261ff79426f17ccf16d43490ba4b790ff84d573083cf94e111258"


@dataclass(frozen=True)
class GitRepo:
    name: str
    url: str
    commit: str
    target: Path
    requirements: str | None = "requirements.txt"


COMFYUI = GitRepo(
    name="ComfyUI",
    url="https://github.com/comfyanonymous/ComfyUI.git",
    commit="5aa71b9bc28809a16596bb9fa3d0a6300d8e3f0e",
    target=COMFY_DIR,
)

CUSTOM_NODE_REPOS = [
    GitRepo(
        name="comfyui_voicebridge",
        url="https://github.com/YanTianlong-01/comfyui_voicebridge.git",
        commit="3728962c0db7b9e05a1d0b341e3dbbd8adba4409",
        target=CUSTOM_NODES_DIR / "comfyui_voicebridge",
    ),
    GitRepo(
        name="ComfyUI_RH_VoxCPM",
        url="https://github.com/RH-RunningHub/ComfyUI_RH_VoxCPM.git",
        commit="8365fe0e1fa60d7547f83ae5db53453a3d9c627d",
        target=CUSTOM_NODES_DIR / "ComfyUI_RH_VoxCPM",
    ),
    GitRepo(
        name="ComfyUI-MelBandRoFormer",
        url="https://github.com/kijai/ComfyUI-MelBandRoFormer.git",
        commit="92c86854e6654f4aacc97484471af95c98ea16d4",
        target=CUSTOM_NODES_DIR / "ComfyUI-MelBandRoFormer",
    ),
    GitRepo(
        name="ComfyUI_RH_LLM_API",
        url="https://github.com/HM-RunningHub/ComfyUI_RH_LLM_API.git",
        commit="26e18d1a769bd08e115b59bfdf170f8a2166c0df",
        target=CUSTOM_NODES_DIR / "ComfyUI_RH_LLM_API",
        requirements=None,
    ),
    GitRepo(
        name="rgthree-comfy",
        url="https://github.com/rgthree/rgthree-comfy.git",
        commit="738105af5fb14e96fbecaf406dc356e284797e8c",
        target=CUSTOM_NODES_DIR / "rgthree-comfy",
    ),
    GitRepo(
        name="ComfyUI-Easy-Use",
        url="https://github.com/yolain/ComfyUI-Easy-Use.git",
        commit="625efbfa2fc20c31797dfffcbb41a26b6d91ab7b",
        target=CUSTOM_NODES_DIR / "ComfyUI-Easy-Use",
    ),
    GitRepo(
        name="ComfyUI_Comfyroll_CustomNodes",
        url="https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes.git",
        commit="d78b780ae43fcf8c6b7c6505e6ffb4584281ceca",
        target=CUSTOM_NODES_DIR / "ComfyUI_Comfyroll_CustomNodes",
        requirements=None,
    ),
    GitRepo(
        name="ComfyUI_AudioTools",
        url="https://github.com/billwuhao/ComfyUI_AudioTools.git",
        commit="41463715b476aa1d44de617119a68d8841aa04bd",
        target=CUSTOM_NODES_DIR / "ComfyUI_AudioTools",
    ),
]


MODEL_LINKS = {
    "voxcpm2": COMFY_DIR / "models" / "voxcpm" / "VoxCPM2",
    "melband": COMFY_DIR / "models" / "diffusion_models" / "MelBandRoFormer_comfy",
    "qwen3_asr": COMFY_DIR / "models" / "Qwen3-ASR",
}

MODEL_DIRS = {
    "voxcpm2": Path("voxcpm") / "VoxCPM2",
    "melband": Path("diffusion_models") / "MelBandRoFormer_comfy",
    "qwen3_asr": Path("Qwen3-ASR"),
}


def model_root() -> Path:
    configured = os.environ.get("VOICEGATE_MODEL_ROOT")
    if configured:
        return Path(configured)
    if DEFAULT_PERSISTENT_MODEL_ROOT.parent.exists():
        return DEFAULT_PERSISTENT_MODEL_ROOT
    return COMFY_DIR / "models"


def model_target(name: str) -> Path:
    return model_root() / MODEL_DIRS[name]


def redacted(cmd: list[str]) -> list[str]:
    result = cmd[:]
    for index, value in enumerate(result[:-1]):
        if value == "--token":
            result[index + 1] = "***"
    return result


def run(cmd: list[str], cwd: Path | None = None, dry_run: bool = False) -> None:
    printable = " ".join(redacted(cmd))
    prefix = f"cd {cwd} && " if cwd else ""
    print(f"+ {prefix}{printable}", flush=True)
    if not dry_run:
        subprocess.run(cmd, cwd=cwd, check=True)


def ensure_git_repo(repo: GitRepo, dry_run: bool = False) -> None:
    if not dry_run:
        repo.target.parent.mkdir(parents=True, exist_ok=True)
    if repo.target.exists() and not (repo.target / ".git").exists():
        raise RuntimeError(f"{repo.target} exists but is not a git repository")

    if not repo.target.exists():
        run(["git", "clone", repo.url, str(repo.target)], dry_run=dry_run)

    run(["git", "fetch", "--all", "--tags"], cwd=repo.target, dry_run=dry_run)
    run(["git", "checkout", repo.commit], cwd=repo.target, dry_run=dry_run)


def install_requirements(repo: GitRepo, dry_run: bool = False) -> None:
    if not repo.requirements:
        return
    requirements = repo.target / repo.requirements
    if not requirements.exists():
        return
    run(
        [sys.executable, "-m", "pip", "install", "-r", str(requirements)],
        dry_run=dry_run,
    )


def ensure_model_link(name: str, dry_run: bool = False) -> None:
    target = model_target(name)
    link = MODEL_LINKS[name]
    print(f"+ mkdir -p {target}")
    print(f"+ link {link} -> {target}")
    if dry_run:
        return

    target.mkdir(parents=True, exist_ok=True)
    link.parent.mkdir(parents=True, exist_ok=True)
    if link.exists() or link.is_symlink():
        if link.resolve() == target.resolve():
            return
        if link.is_dir() and not link.is_symlink() and not any(link.iterdir()):
            link.rmdir()
        else:
            raise RuntimeError(f"{link} exists and does not point to {target}")
    link.symlink_to(target, target_is_directory=True)


def prepare_model_dirs(dry_run: bool = False) -> None:
    for name in MODEL_LINKS:
        ensure_model_link(name, dry_run=dry_run)


def file_sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as file:
        for chunk in iter(lambda: file.read(8 * 1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def melband_model_path() -> Path:
    return model_target("melband") / MELBAND_FILENAME


def validate_melband_model(*, verify_hash: bool = True) -> tuple[bool, str]:
    path = melband_model_path()
    if not path.is_file():
        return False, "missing"
    size = path.stat().st_size
    if size != MELBAND_SIZE:
        return False, f"size_mismatch expected={MELBAND_SIZE} actual={size}"
    if verify_hash:
        try:
            digest = file_sha256(path)
        except OSError as exc:
            return False, f"read_error {type(exc).__name__}: {exc}"
        if digest != MELBAND_SHA256:
            return False, f"sha256_mismatch expected={MELBAND_SHA256} actual={digest}"
    return True, "ok"


def patch_melband_loader(dry_run: bool = False) -> None:
    """Avoid safetensors mmap on persistent Space storage.

    ComfyUI's generic loader uses safetensors.safe_open(), which memory maps the
    model file. A damaged file or an unstable mmap on /data can terminate the
    interpreter with SIGBUS before Python can report a normal exception.
    Loading from bytes uses regular reads and turns corruption into a catchable
    safetensors error instead.
    """

    nodes_path = CUSTOM_NODES_DIR / "ComfyUI-MelBandRoFormer" / "nodes.py"
    print(f"+ patch non-mmap MelBand loader: {nodes_path}", flush=True)
    if dry_run:
        return
    if not nodes_path.is_file():
        raise RuntimeError(f"MelBand node file is missing: {nodes_path}")

    text = nodes_path.read_text(encoding="utf-8")
    if "load_safetensors_bytes" not in text:
        text = text.replace(
            "import torchaudio.functional as TAF\n",
            "import torchaudio.functional as TAF\n"
            "from safetensors.torch import load as load_safetensors_bytes\n",
        )
        text = text.replace(
            "model.load_state_dict(load_torch_file(model_path), strict=True)",
            "with open(model_path, \"rb\") as model_file:\n"
            "            state_dict = load_safetensors_bytes(model_file.read())\n"
            "        model.load_state_dict(state_dict, strict=True)",
        )
    if "load_safetensors_bytes" not in text or "state_dict = load_safetensors_bytes" not in text:
        raise RuntimeError("Could not apply the non-mmap MelBand loader patch")
    nodes_path.write_text(text, encoding="utf-8")


def download_models(dry_run: bool = False) -> None:
    """Download large model assets.

    This is opt-in because the model set is large. Keep downloads separate from
    plain bootstrap so we can test repository setup without consuming Space
    storage and startup time.
    """

    token = os.environ.get("HF_TOKEN")
    if dry_run:
        print(f"+ snapshot_download openbmb/VoxCPM2 -> {model_target('voxcpm2')}")
        print(
            "+ hf_hub_download "
            f"Kijai/MelBandRoFormer_comfy/MelBandRoformer_fp32.safetensors -> {model_target('melband')}"
        )
        print(f"+ snapshot_download Qwen/Qwen3-ASR-1.7B -> {model_target('qwen3_asr') / 'Qwen3-ASR-1.7B'}")
        print(
            "+ snapshot_download "
            f"Qwen/Qwen3-ForcedAligner-0.6B -> {model_target('qwen3_asr') / 'Qwen3-ForcedAligner-0.6B'}"
        )
        return

    from huggingface_hub import hf_hub_download, snapshot_download

    snapshot_download(
        repo_id="openbmb/VoxCPM2",
        local_dir=model_target("voxcpm2"),
        token=token,
    )
    melband_valid, melband_reason = validate_melband_model(verify_hash=True)
    print(f"+ validate MelBand model: {melband_reason}", flush=True)
    if not melband_valid and melband_model_path().exists():
        print(f"+ remove invalid MelBand model: {melband_model_path()}", flush=True)
        melband_model_path().unlink()
    hf_hub_download(
        repo_id="Kijai/MelBandRoFormer_comfy",
        filename=MELBAND_FILENAME,
        local_dir=model_target("melband"),
        token=token,
        force_download=not melband_valid,
    )
    melband_valid, melband_reason = validate_melband_model(verify_hash=True)
    print(f"+ verify downloaded MelBand model: {melband_reason}", flush=True)
    if not melband_valid:
        raise RuntimeError(f"MelBand model validation failed: {melband_reason}")
    snapshot_download(
        repo_id="Qwen/Qwen3-ASR-1.7B",
        local_dir=model_target("qwen3_asr") / "Qwen3-ASR-1.7B",
        token=token,
    )
    snapshot_download(
        repo_id="Qwen/Qwen3-ForcedAligner-0.6B",
        local_dir=model_target("qwen3_asr") / "Qwen3-ForcedAligner-0.6B",
        token=token,
    )


def print_summary() -> None:
    print("\nRuntime layout:")
    print(f"  ComfyUI: {COMFY_DIR}")
    print(f"  custom_nodes: {CUSTOM_NODES_DIR}")
    print(f"  model_root: {model_root()}")
    for name, path in MODEL_LINKS.items():
        print(f"  model:{name}: {path} -> {model_target(name)}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--skip-pip", action="store_true")
    parser.add_argument("--with-models", action="store_true")
    parser.add_argument("--clean", action="store_true")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.clean:
        if COMFY_DIR.exists():
            print(f"Removing {COMFY_DIR}")
            if not args.dry_run:
                shutil.rmtree(COMFY_DIR)

    ensure_git_repo(COMFYUI, dry_run=args.dry_run)
    if not args.dry_run:
        CUSTOM_NODES_DIR.mkdir(parents=True, exist_ok=True)
    for repo in CUSTOM_NODE_REPOS:
        ensure_git_repo(repo, dry_run=args.dry_run)
    patch_melband_loader(dry_run=args.dry_run)

    if not args.skip_pip:
        install_requirements(COMFYUI, dry_run=args.dry_run)
        for repo in CUSTOM_NODE_REPOS:
            install_requirements(repo, dry_run=args.dry_run)

    prepare_model_dirs(dry_run=args.dry_run)
    if args.with_models:
        download_models(dry_run=args.dry_run)

    print_summary()


if __name__ == "__main__":
    main()