"""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 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")


@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 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,
    )
    hf_hub_download(
        repo_id="Kijai/MelBandRoFormer_comfy",
        filename="MelBandRoformer_fp32.safetensors",
        local_dir=model_target("melband"),
        token=token,
    )
    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)

    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()