#!/usr/bin/env python3
"""Launch ComfyUI with arguments selected by the active deployment profile."""

from __future__ import annotations

import argparse
import os
import subprocess
import sys
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
COMFY_ROOT = ROOT / "ComfyUI"
sys.path.insert(0, str(ROOT))

from studio.deployment import current_deployment  # noqa: E402


def runtime_environment() -> dict[str, str]:
    cache_root = ROOT / ".runtime-cache"
    triton_root = cache_root / "triton"
    values = {
        "TRITON_HOME": triton_root,
        "TRITON_CACHE_DIR": triton_root / "cache",
        "TRITON_DUMP_DIR": triton_root / "dump",
        "TRITON_OVERRIDE_DIR": triton_root / "override",
        "TORCHINDUCTOR_CACHE_DIR": cache_root / "torchinductor",
    }
    for directory in values.values():
        directory.mkdir(parents=True, exist_ok=True)
    env = os.environ.copy()
    env.update({name: str(path) for name, path in values.items()})
    env["PYTHONUTF8"] = "1"
    return env


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", type=int, default=8189)
    parser.add_argument("--listen", default="127.0.0.1")
    parser.add_argument("--open-browser", action="store_true")
    args = parser.parse_args()

    plan = current_deployment(refresh_runtime=True)
    profile = plan["profile"]
    runtime = plan["runtime"]
    if not plan["capabilities"]["h3"]:
        print("H3 generation is not ready for the selected runtime.", file=sys.stderr)
        for warning in plan["warnings"]:
            print(f"- {warning}", file=sys.stderr)
        return 2

    command = [
        *runtime["command"],
        "-u",
        "main.py",
        "--listen",
        args.listen,
        "--port",
        str(args.port),
        "--disable-api-nodes",
        "--reserve-vram",
        "1.5",
        "--cache-none",
    ]
    if args.open_browser:
        command.append("--auto-launch")
    if profile["acceleration"] == "sageattention":
        command.append("--use-sage-attention")
    else:
        command.append("--use-pytorch-cross-attention")
    if profile["id"] in {"rtx_low_vram", "non_nvidia_experimental"}:
        command.append("--lowvram")
    if runtime.get("backend") == "directml":
        command.append("--directml")

    print(f"Runtime: {runtime['executable']} ({runtime['backend']})", flush=True)
    print(f"Profile: {profile['label']} / {profile['acceleration']}", flush=True)
    try:
        return subprocess.call(command, cwd=COMFY_ROOT, env=runtime_environment())
    except KeyboardInterrupt:
        return 130


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