#!/usr/bin/env python3
"""Local infinite-canvas UI backend for the complete MiniMax H3 workflow set."""

from __future__ import annotations

import argparse
import asyncio
import importlib.util
import json
import os
import re
import secrets
import subprocess
import time
import uuid
import webbrowser
from collections import deque
from datetime import datetime
from html import escape
from pathlib import Path
from typing import Any
from urllib.parse import quote, urlencode, urlparse

import av
from aiohttp import ClientSession, ClientTimeout, web

try:
    from .deployment import apply_plan, build_plan, current_deployment, probe_hardware
    from .repair import diagnose as diagnose_environment
    from .repair import run_repair
except ImportError:
    from deployment import apply_plan, build_plan, current_deployment, probe_hardware
    from repair import diagnose as diagnose_environment
    from repair import run_repair


ROOT = Path(__file__).resolve().parents[1]
COMFY_ROOT = ROOT / "ComfyUI"
COMFY_INPUT = COMFY_ROOT / "input"
COMFY_OUTPUT = COMFY_ROOT / "output"
UPLOAD_ROOT = COMFY_INPUT / "h3_studio"
JOB_ROOT = Path(__file__).with_name("jobs")
FL2VA_RUNNER = ROOT / "tools" / "run_h3_736p_smoke.ps1"
R2V_RUNNER = ROOT / "tools" / "run_h3_r2v.ps1"
REDRAW_RUNNER = ROOT / "tools" / "run_h3_redraw_rtx.ps1"
MAXINE_RUNNER = ROOT / "tools" / "maxine_vsr.py"
UNIVERSAL_RUNNER = ROOT / "tools" / "universal_video_finish.py"
INDEX_FILE = Path(__file__).with_name("index.html")
CLOUD_ACCESS_FILE = Path(__file__).with_name("cloud_access.json")

VIDEO_EXTENSIONS = {".mp4", ".mov", ".mkv", ".webm", ".avi"}
IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
AUDIO_EXTENSIONS = {".wav", ".mp3", ".flac", ".m4a", ".ogg", ".aac"}

PRESETS = {
    "clean": {"label": "清洁稳态", "denoise": 0.35, "steps": 8},
    "balanced": {"label": "平衡重绘", "denoise": 0.50, "steps": 8},
    "detail": {"label": "细节增强", "denoise": 0.60, "steps": 8},
}

RESOLUTIONS = {
    "draft": {"label": "草图预览", "width": 608, "height": 352},
    "standard": {"label": "标准电影", "width": 1280, "height": 736},
    "scope": {"label": "超宽银幕", "width": 1728, "height": 736},
    "direct2k": {"label": "Direct 2K 实验", "width": 1920, "height": 1088},
}

WORKFLOW_MODES = {
    "t2v": {"label": "文生视频", "model": "FL2VA + Turbo", "assets": []},
    "i2v": {"label": "首帧生成", "model": "FL2VA + Turbo", "assets": ["first_frame"]},
    "l2v": {"label": "尾帧生成", "model": "FL2VA + Turbo", "assets": ["last_frame"]},
    "fl2v": {"label": "首尾帧生成", "model": "FL2VA + Turbo", "assets": ["first_frame", "last_frame"]},
    "r2v": {"label": "多参考生成", "model": "Ref2VA 20 Steps", "assets": ["references"]},
    "redraw": {"label": "视频重绘", "model": "Concat AV + Turbo", "assets": ["input_video"]},
    "upscale": {"label": "降噪与超分", "model": "自适应本地终修", "assets": ["input_video"]},
    "direct2k": {"label": "Direct 2K 实验", "model": "FL2VA + Turbo", "assets": []},
}

DEFAULT_PROMPT = (
    "A cinematic shot with stable geometry, coherent motion, realistic lighting, atmospheric "
    "depth and native environmental audio. Smooth controlled camera movement, photorealistic, "
    "no flicker, no deformation, no text, no logo."
)

REDRAW_PROMPT = (
    "Preserve the source video's composition, architecture, subject identity, camera path, "
    "timing and physical motion. Refine surfaces, atmospheric depth, water, mist and lighting "
    "with cinematic realism. Stable geometry, coherent temporal detail, no flicker, no "
    "deformation, no new large structures, no text, no logo."
)


def normalize_cloud_origin(value: str) -> str | None:
    """Accept production HTTPS origins and loopback HTTP origins used for local preview."""
    try:
        parsed = urlparse(value.strip().rstrip("/"))
        hostname = (parsed.hostname or "").lower()
        loopback = hostname in {"127.0.0.1", "localhost", "::1"}
        if not parsed.netloc or parsed.username or parsed.password:
            return None
        if parsed.scheme != "https" and not (parsed.scheme == "http" and loopback):
            return None
        if parsed.path not in {"", "/"} or parsed.params or parsed.query or parsed.fragment:
            return None
        # Accessing .port also validates malformed port values.
        _ = parsed.port
        return f"{parsed.scheme}://{parsed.netloc}"
    except (TypeError, ValueError):
        return None


def load_cloud_access() -> dict[str, Any]:
    if not CLOUD_ACCESS_FILE.is_file():
        return {"enabled": False, "access_key": "", "allowed_origins": []}
    try:
        value = json.loads(CLOUD_ACCESS_FILE.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return {"enabled": False, "access_key": "", "allowed_origins": []}
    origins = value.get("allowed_origins") if isinstance(value, dict) else []
    normalized = [normalize_cloud_origin(str(origin)) for origin in origins] if isinstance(origins, list) else []
    return {
        "enabled": bool(value.get("enabled")) if isinstance(value, dict) else False,
        "access_key": str(value.get("access_key") or "") if isinstance(value, dict) else "",
        "allowed_origins": [origin for origin in normalized if origin],
    }


def pair_cloud_origin(origin: str) -> str:
    current = load_cloud_access()
    allowed = set(current["allowed_origins"])
    allowed.add(origin)
    access_key = current["access_key"] or secrets.token_urlsafe(32)
    value = {
        "version": 2,
        "enabled": True,
        "access_key": access_key,
        "allowed_origins": sorted(allowed),
    }
    CLOUD_ACCESS_FILE.write_text(json.dumps(value, ensure_ascii=False, indent=2), encoding="utf-8")
    return access_key


@web.middleware
async def cloud_access_middleware(request: web.Request, handler: Any) -> web.StreamResponse:
    origin = str(request.headers.get("Origin") or "").rstrip("/")
    host_port = request.host.rsplit(":", 1)[-1] if ":" in request.host else "80"
    local_origins = {
        f"{request.scheme}://{request.host}",
        f"http://127.0.0.1:{host_port}",
        f"http://localhost:{host_port}",
    }
    local_host = request.host.split(":", 1)[0].lower() in {"127.0.0.1", "localhost", "[::1]"}
    if origin in local_origins or (not origin and local_host):
        return await handler(request)
    if not origin:
        raise web.HTTPUnauthorized(text="远程访问必须携带网页来源和本地访问密钥")

    access = load_cloud_access()
    if not access["enabled"] or origin not in access["allowed_origins"]:
        raise web.HTTPForbidden(text="该网页来源未获准连接本地 H3 Studio")
    if request.method == "OPTIONS":
        response: web.StreamResponse = web.Response(status=204)
    else:
        if request.path.startswith("/api/"):
            supplied = str(request.headers.get("X-H3-Studio-Key") or "")
            if not access["access_key"] or not secrets.compare_digest(supplied, access["access_key"]):
                response = web.Response(status=401, text="本地 H3 Studio 访问密钥不正确")
            else:
                response = await handler(request)
        else:
            response = await handler(request)
    response.headers["Access-Control-Allow-Origin"] = origin
    response.headers["Vary"] = "Origin"
    response.headers["Access-Control-Allow-Methods"] = "GET, POST, OPTIONS"
    response.headers["Access-Control-Allow-Headers"] = "Content-Type, X-H3-Studio-Key"
    response.headers["Access-Control-Expose-Headers"] = "Content-Length, Content-Range, Accept-Ranges"
    response.headers["Access-Control-Max-Age"] = "600"
    if str(request.headers.get("Access-Control-Request-Private-Network") or "").lower() == "true":
        response.headers["Access-Control-Allow-Private-Network"] = "true"
    return response


def inside(path: Path, root: Path) -> bool:
    try:
        path.resolve().relative_to(root.resolve())
        return True
    except ValueError:
        return False


def safe_stem(value: str) -> str:
    cleaned = re.sub(r"[^0-9A-Za-z._\-\u4e00-\u9fff]+", "_", value).strip("._")
    return (cleaned or "asset")[:72]


def asset_kind(path: Path) -> str | None:
    suffix = path.suffix.lower()
    if suffix in VIDEO_EXTENSIONS:
        return "video"
    if suffix in IMAGE_EXTENSIONS:
        return "image"
    if suffix in AUDIO_EXTENSIONS:
        return "audio"
    return None


def media_url(path: Path) -> str:
    return f"/media?path={quote(str(path.resolve()))}"


def even_scaled(value: int, scale: float) -> int:
    return int(round(value * scale / 2.0) * 2)


def probe_video(path: Path) -> dict[str, Any]:
    with av.open(str(path)) as container:
        if not container.streams.video:
            raise ValueError("文件不包含视频轨")
        stream = container.streams.video[0]
        return {
            "width": int(stream.codec_context.width),
            "height": int(stream.codec_context.height),
            "frames": int(stream.frames or 0),
            "fps": float(stream.average_rate or 24),
            "duration": float(container.duration / av.time_base) if container.duration else 0,
            "audio": bool(container.streams.audio),
        }


class Studio:
    def __init__(self, comfy_port: int) -> None:
        self.comfy_port = comfy_port
        self.comfy_url = f"http://127.0.0.1:{comfy_port}"
        self.comfy_process: asyncio.subprocess.Process | None = None
        self.owns_comfy = False
        self.service = "starting"
        self.service_message = "正在启动 ComfyUI…"
        self.comfy_log: deque[str] = deque(maxlen=140)
        self.job_process: asyncio.subprocess.Process | None = None
        self.job_task: asyncio.Task[None] | None = None
        self.job: dict[str, Any] = self.empty_job()
        self.deployment = current_deployment()

    def python_command(self) -> list[str]:
        command = self.deployment.get("runtime", {}).get("command") or []
        return [str(part) for part in command] if command else [os.sys.executable]

    @staticmethod
    def empty_job() -> dict[str, Any]:
        return {
            "id": None,
            "state": "idle",
            "stage": "等待任务",
            "stage_key": "idle",
            "progress": 0,
            "message": "选择模板并连接素材后即可开始",
            "started_at": None,
            "elapsed": 0,
            "mode": None,
            "input": None,
            "source_output": None,
            "output": None,
            "output_url": None,
            "logs": [],
            "error": None,
        }

    async def comfy_ready(self) -> bool:
        try:
            async with ClientSession(timeout=ClientTimeout(total=2)) as session:
                async with session.get(f"{self.comfy_url}/system_stats") as response:
                    return response.status == 200
        except Exception:
            return False

    async def start_comfy(self) -> None:
        profile = self.deployment["profile"]
        if profile["h3_support"] == "unavailable":
            if self.deployment.get("capabilities", {}).get("upscale"):
                self.service = "ready"
                finish_label = "NVIDIA Maxine" if profile["finish_engine"] == "maxine" else "通用视频终修"
                self.service_message = f"{finish_label}已就绪；当前运行时未启用本地 H3"
            else:
                self.service = "setup"
                self.service_message = "未找到可用的视频终修运行时，请先完成部署"
            return
        required = [COMFY_ROOT / "main.py", FL2VA_RUNNER]
        missing = [path.name for path in required if not path.is_file()]
        if missing or not self.deployment["readiness"]["ready"]:
            self.service = "setup"
            self.service_message = "请先在部署向导中补全本地引擎、模型与插件"
            return
        if await self.comfy_ready():
            self.service = "ready"
            self.service_message = "已连接现有 ComfyUI"
            return

        cache_root = ROOT / ".runtime-cache"
        triton_root = cache_root / "triton"
        for directory in (
            triton_root / "cache",
            triton_root / "dump",
            triton_root / "override",
            cache_root / "torchinductor",
        ):
            directory.mkdir(parents=True, exist_ok=True)
        env = os.environ.copy()
        env.update(
            {
                "TRITON_HOME": str(triton_root),
                "TRITON_CACHE_DIR": str(triton_root / "cache"),
                "TRITON_DUMP_DIR": str(triton_root / "dump"),
                "TRITON_OVERRIDE_DIR": str(triton_root / "override"),
                "TORCHINDUCTOR_CACHE_DIR": str(cache_root / "torchinductor"),
                "PYTHONUTF8": "1",
            }
        )
        creationflags = subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0
        command = [
            *self.python_command(),
            "-u",
            "main.py",
            "--listen",
            "127.0.0.1",
            "--port",
            str(self.comfy_port),
            "--disable-api-nodes",
            "--reserve-vram",
            "1.5",
            "--cache-none",
        ]
        if profile["acceleration"] == "sageattention":
            command.append("--use-sage-attention")
        elif profile["acceleration"] == "pytorch":
            command.append("--use-pytorch-cross-attention")
        if profile["id"] in {"rtx_low_vram", "non_nvidia_experimental"}:
            command.append("--lowvram")
        if self.deployment.get("runtime", {}).get("backend") == "directml":
            command.append("--directml")
        self.comfy_process = await asyncio.create_subprocess_exec(
            *command,
            cwd=str(COMFY_ROOT),
            env=env,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.STDOUT,
            creationflags=creationflags,
        )
        self.owns_comfy = True
        asyncio.create_task(self.read_comfy_log())
        for _ in range(180):
            if await self.comfy_ready():
                self.service = "ready"
                finish_label = "Maxine RTX" if profile["finish_engine"] == "maxine" else "通用终修"
                self.service_message = f"H3 / {finish_label} 已就绪"
                return
            if self.comfy_process.returncode is not None:
                break
            await asyncio.sleep(1)
        self.service = "error"
        tail = "\n".join(list(self.comfy_log)[-8:])
        self.service_message = "ComfyUI 启动失败" + (f"：{tail}" if tail else "")

    async def read_comfy_log(self) -> None:
        if not self.comfy_process or not self.comfy_process.stdout:
            return
        while line := await self.comfy_process.stdout.readline():
            text = line.decode("utf-8", errors="replace").strip()
            if text:
                self.comfy_log.append(text)

    async def restart_comfy(self) -> None:
        if self.job["state"] in {"running", "cancelling"}:
            raise RuntimeError("当前工作流仍在运行，无法重启引擎")
        if self.comfy_process and self.comfy_process.returncode is None:
            if not self.owns_comfy:
                raise RuntimeError("当前 ComfyUI 由外部程序启动；请在原窗口中停止后再重启")
            self.comfy_process.terminate()
            try:
                await asyncio.wait_for(self.comfy_process.wait(), timeout=15)
            except TimeoutError as exc:
                raise RuntimeError("引擎没有正常退出，已停止自动操作；请手动关闭启动窗口") from exc
        self.comfy_process = None
        self.owns_comfy = False
        self.deployment = current_deployment(refresh_runtime=True)
        self.service = "starting"
        self.service_message = "故障中心正在重新启动本地引擎…"
        await self.start_comfy()

    def list_assets(self) -> list[dict[str, Any]]:
        assets: list[dict[str, Any]] = []
        for root, source in ((COMFY_OUTPUT, "output"), (COMFY_INPUT, "input")):
            if not root.exists():
                continue
            for path in root.rglob("*"):
                kind = asset_kind(path)
                if not path.is_file() or kind is None:
                    continue
                stat = path.stat()
                assets.append(
                    {
                        "name": path.name,
                        "path": str(path.resolve()),
                        "url": media_url(path),
                        "kind": kind,
                        "source": source,
                        "size": stat.st_size,
                        "modified": stat.st_mtime,
                    }
                )
        assets.sort(key=lambda item: item["modified"], reverse=True)
        return assets[:160]

    def snapshot(self) -> dict[str, Any]:
        job = dict(self.job)
        if job["state"] in {"running", "cancelling"} and job["started_at"]:
            job["elapsed"] = int(time.time() - job["started_at"])
        return {
            "service": self.service,
            "service_message": self.service_message,
            "job": job,
            "presets": PRESETS,
            "resolutions": RESOLUTIONS,
            "modes": WORKFLOW_MODES,
            "deployment": self.deployment,
        }

    def set_stage(self, key: str, title: str, progress: int, message: str) -> None:
        self.job.update(
            {
                "stage_key": key,
                "stage": title,
                "progress": max(self.job["progress"], min(progress, 99)),
                "message": message,
            }
        )

    def add_log(self, line: str) -> None:
        self.job["logs"].append(line)
        if len(self.job["logs"]) > 140:
            del self.job["logs"][:-140]

    def parse_runner_line(self, line: str) -> None:
        self.add_log(line)
        if line.startswith(("SAMPLE_PROMPT_ID=", "R2V_SAMPLE_PROMPT_ID=")):
            label = "Ref2VA 多参考采样" if line.startswith("R2V") else "H3 Turbo 采样"
            self.set_stage("generate", label, 7, "加载文本编码器、H3 主模型与条件素材")
        elif line.startswith(("SAMPLE_WAITING_SECONDS=", "R2V_SAMPLE_WAITING_SECONDS=")):
            seconds = int(line.rsplit("=", 1)[1])
            self.set_stage("generate", "H3 视频生成", 10 + min(seconds // 8, 53), "SageAttention 正在生成联合音视频潜空间")
        elif line.startswith(("SAMPLE_ELAPSED_SECONDS=", "R2V_SAMPLE_ELAPSED_SECONDS=")):
            self.set_stage("generate", "H3 采样完成", 65, "已保存可续跑的联合音视频潜空间")
        elif line.startswith("DECODE_PROMPT_ID="):
            self.set_stage("decode", "音画解码", 68, "已释放 H3，正在分别解码视频和原生音频")
        elif line.startswith("DECODE_WAITING_SECONDS="):
            seconds = int(line.rsplit("=", 1)[1])
            self.set_stage("decode", "音画解码", 69 + min(seconds // 5, 7), "Video VAE 与 Audio VAE 正在解码")
        elif line.startswith("DECODE_ELAPSED_SECONDS="):
            self.set_stage("decode", "源片完成", 77, "H.264 与 AAC 源片已经封装")
        elif line.startswith("ENCODE_SOURCE_AV_PROMPT_ID="):
            self.set_stage("encode", "编码源视频", 5, "分块编码视频与音频潜空间")
        elif line.startswith("ENCODE_SOURCE_AV_WAITING_SECONDS="):
            seconds = int(line.rsplit("=", 1)[1])
            self.set_stage("encode", "编码源视频", 6 + min(seconds // 6, 18), "构建 Concat AV 重绘起点")
        elif line.startswith("ENCODE_SOURCE_AV_ELAPSED_SECONDS="):
            self.set_stage("encode", "源潜空间完成", 25, "已保存源 AV 潜空间")
        elif line.startswith("H3_REDRAW_PROMPT_ID="):
            self.set_stage("redraw", "H3 稳态重绘", 28, "加载双锚、Turbo LoRA 与源潜空间")
        elif line.startswith("H3_REDRAW_WAITING_SECONDS="):
            seconds = int(line.rsplit("=", 1)[1])
            self.set_stage("redraw", "H3 稳态重绘", 30 + min(seconds // 10, 43), "SageAttention 正在执行 8 步重绘")
        elif line.startswith("H3_REDRAW_ELAPSED_SECONDS="):
            self.set_stage("redraw", "H3 重绘完成", 75, "重绘潜空间已保存")
        elif line.startswith(("DECODE_RTX_MUX_PROMPT_ID=", "DECODE_MUX_PROMPT_ID=")):
            is_rtx = line.startswith("DECODE_RTX")
            self.set_stage("upscale" if is_rtx else "decode", "解码与 RTX 超分" if is_rtx else "解码重绘成片", 78, "解码音画并执行 RTX VSR 1.5×" if is_rtx else "保留源分辨率，随后进入通用终修")
        elif line.startswith(("DECODE_RTX_MUX_WAITING_SECONDS=", "DECODE_MUX_WAITING_SECONDS=")):
            seconds = int(line.rsplit("=", 1)[1])
            self.set_stage("decode", "解码重绘成片", 80 + min(seconds // 4, 17), "正在生成高码率中间片")
        elif line.startswith(("DECODE_RTX_MUX_ELAPSED_SECONDS=", "DECODE_MUX_ELAPSED_SECONDS=")):
            self.set_stage("mux", "封装成片", 98, "写入 H.264 视频与 AAC 音轨")

    async def run_command(self, command: list[str], cwd: Path = ROOT, post_total_frames: int = 0) -> int:
        creationflags = subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0
        self.job_process = await asyncio.create_subprocess_exec(
            *command,
            cwd=str(cwd),
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.STDOUT,
            creationflags=creationflags,
        )
        assert self.job_process.stdout is not None
        while line := await self.job_process.stdout.readline():
            decoded = line.decode("utf-8", errors="replace").strip()
            if not decoded:
                continue
            self.parse_runner_line(decoded)
            if post_total_frames and decoded.startswith("frames="):
                match = re.match(r"frames=(\d+)", decoded)
                if match:
                    frame = int(match.group(1))
                    progress = 80 + min(int(18 * frame / max(post_total_frames, 1)), 18)
                    engine = self.deployment["profile"]["finish_engine"]
                    title = "NVIDIA 终修" if engine == "maxine" else "通用终修"
                    self.set_stage("post", title, progress, f"已处理 {frame}/{post_total_frames} 帧")
        return await self.job_process.wait()

    @staticmethod
    def latest_output(prefix: str) -> Path:
        name = Path(prefix).name
        candidates = sorted(
            (COMFY_OUTPUT / "video").glob(f"{name}_*.mp4"),
            key=lambda path: path.stat().st_mtime,
            reverse=True,
        )
        if not candidates:
            raise RuntimeError(f"没有找到输出视频：{name}")
        return candidates[0].resolve()

    @staticmethod
    def comfy_name(path_value: str) -> str:
        path = Path(path_value).resolve()
        if inside(path, COMFY_INPUT):
            return path.relative_to(COMFY_INPUT.resolve()).as_posix()
        if inside(path, COMFY_OUTPUT):
            relative = path.relative_to(COMFY_OUTPUT.resolve()).as_posix()
            return f"{relative} [output]"
        raise ValueError("素材必须位于 ComfyUI/input 或 ComfyUI/output")

    async def run_fl2va(self, config: dict[str, Any], base: str) -> Path:
        output_prefix = f"video/{base}_Source"
        latent_prefix = f"latents/{base}_Latent"
        command = [
            "powershell.exe", "-NoProfile", "-ExecutionPolicy", "Bypass", "-File", str(FL2VA_RUNNER),
            "-Server", self.comfy_url,
            "-Steps", str(config["steps"]),
            "-Duration", str(config["duration"]),
            "-Width", str(config["width"]),
            "-Height", str(config["height"]),
            "-Seed", str(config["seed"]),
            "-PromptText", config["prompt"],
            "-OutputPrefix", output_prefix,
            "-LatentPrefix", latent_prefix,
            "-TurboMode", "true" if config["turbo"] else "false",
        ]
        if config.get("first_frame"):
            command.extend(["-FirstFrame", self.comfy_name(config["first_frame"])])
        if config.get("last_frame"):
            command.extend(["-LastFrame", self.comfy_name(config["last_frame"])])
        if await self.run_command(command) != 0:
            raise RuntimeError("H3 FL2VA 工作流执行失败")
        return self.latest_output(output_prefix)

    async def run_r2v(self, config: dict[str, Any], base: str) -> Path:
        JOB_ROOT.mkdir(parents=True, exist_ok=True)
        output_prefix = f"video/{base}_Source"
        latent_prefix = f"latents/{base}_Latent"
        references = config["references"]
        r2v_config = {
            "width": config["width"],
            "height": config["height"],
            "steps": config["r2v_steps"],
            "duration": config["duration"],
            "seed": config["seed"],
            "prompt": config["prompt"],
            "ref_image_size": config["ref_image_size"],
            "output_prefix": output_prefix,
            "latent_prefix": latent_prefix,
            "ref_images": [self.comfy_name(ref["path"]) for ref in references if ref["kind"] == "image"],
            "ref_videos": [
                {"file": self.comfy_name(ref["path"]), "use_audio": bool(ref.get("use_audio", True))}
                for ref in references if ref["kind"] == "video"
            ],
            "ref_audios": [self.comfy_name(ref["path"]) for ref in references if ref["kind"] == "audio"],
        }
        config_path = JOB_ROOT / f"{self.job['id']}_r2v.json"
        config_path.write_text(json.dumps(r2v_config, ensure_ascii=False, indent=2), encoding="utf-8")
        command = [
            "powershell.exe", "-NoProfile", "-ExecutionPolicy", "Bypass", "-File", str(R2V_RUNNER),
            "-Server", self.comfy_url,
            "-ConfigFile", str(config_path),
        ]
        if await self.run_command(command) != 0:
            raise RuntimeError("H3 Ref2VA 工作流执行失败")
        return self.latest_output(output_prefix)

    async def run_redraw(self, config: dict[str, Any], base: str) -> Path:
        use_maxine = config["finish_engine"] == "maxine"
        output_prefix = f"video/{base}_Redraw_{'RTX' if use_maxine else 'Source'}"
        command = [
            "powershell.exe", "-NoProfile", "-ExecutionPolicy", "Bypass", "-File", str(REDRAW_RUNNER),
            "-Server", self.comfy_url,
            "-InputFile", config["input_video"],
            "-Steps", str(config["steps"]),
            "-Denoise", str(config["redraw_denoise"]),
            "-Seed", str(config["seed"]),
            "-RtxScale", str(config["rtx_scale"]),
            "-PromptText", config["prompt"],
            "-OutputPrefix", output_prefix,
            "-SourceLatentPrefix", f"latents/{base}_SourceAV",
            "-SampledLatentPrefix", f"latents/{base}_RedrawLatent",
        ]
        if not use_maxine:
            command.append("-SkipRtx")
        if await self.run_command(command) != 0:
            raise RuntimeError("H3 视频重绘工作流执行失败")
        return self.latest_output(output_prefix)

    async def run_postprocess(self, source: Path, config: dict[str, Any], base: str, force_upscale: bool | None = None) -> Path:
        denoise = config["finish_denoise"]
        use_upscale = config["rtx_enabled"] if force_upscale is None else force_upscale
        if not use_upscale and denoise == "none":
            return source
        info = probe_video(source)
        engine = config["finish_engine"]
        suffix = "RTX" if engine == "maxine" and use_upscale else ("Universal" if engine == "universal" else "Denoise")
        output = (COMFY_OUTPUT / "video" / f"{base}_{suffix}_{denoise}.mp4").resolve()
        runner = MAXINE_RUNNER if engine == "maxine" else UNIVERSAL_RUNNER
        command = [*self.python_command(), str(runner), str(source), str(output), "--denoise", denoise, "--force"]
        if use_upscale:
            scale = config["rtx_scale"]
            sr_width = even_scaled(info["width"], scale)
            sr_height = even_scaled(info["height"], scale)
            if info["width"] == 1280 and info["height"] == 736 and abs(scale - 1.5) < 0.001:
                target_width, target_height = 1920, 1088
            else:
                target_width, target_height = sr_width, sr_height
            command.extend(["--scale", str(scale), "--width", str(target_width), "--height", str(target_height)])
            if engine == "maxine":
                command.extend(["--quality", config["rtx_quality"]])
        else:
            command.append("--denoise-only")
        title = "NVIDIA 终修" if engine == "maxine" else "通用终修"
        detail = "时序降噪与 RTX 视频超分" if engine == "maxine" else "保结构时序降噪与 Lanczos 高质量缩放"
        self.set_stage("post", title, 80, f"准备{detail}")
        if await self.run_command(command, post_total_frames=info["frames"]) != 0:
            raise RuntimeError(f"{title}工作流执行失败")
        if not output.is_file():
            raise RuntimeError(f"{title}未生成输出文件")
        return output

    async def run_job(self, config: dict[str, Any]) -> None:
        stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        base = f"H3Studio_{config['mode']}_{stamp}_{self.job['id'][-6:]}"
        try:
            if config["mode"] in {"t2v", "i2v", "l2v", "fl2v", "direct2k"}:
                source = await self.run_fl2va(config, base)
                final = await self.run_postprocess(source, config, base)
            elif config["mode"] == "r2v":
                source = await self.run_r2v(config, base)
                final = await self.run_postprocess(source, config, base)
            elif config["mode"] == "redraw":
                source = await self.run_redraw(config, base)
                final = await self.run_postprocess(
                    source,
                    config,
                    base,
                    force_upscale=config["finish_engine"] != "maxine",
                )
            elif config["mode"] == "upscale":
                source = Path(config["input_video"]).resolve()
                final = await self.run_postprocess(source, config, base)
            else:
                raise RuntimeError("未知工作流模式")

            if self.job["state"] == "cancelling":
                self.job.update(
                    {
                        "state": "cancelled",
                        "stage": "已取消",
                        "stage_key": "cancelled",
                        "message": "任务已安全中断；已经保存的潜空间仍可复用",
                    }
                )
                return
            self.job.update(
                {
                    "state": "complete",
                    "stage": "工作流完成",
                    "stage_key": "complete",
                    "progress": 100,
                    "message": "音画生成、优化与封装已经完成",
                    "source_output": str(source),
                    "output": str(final),
                    "output_url": media_url(final),
                    "elapsed": int(time.time() - self.job["started_at"]),
                }
            )
        except asyncio.CancelledError:
            raise
        except Exception as error:
            if self.job["state"] == "cancelling":
                self.job.update(
                    {
                        "state": "cancelled",
                        "stage": "已取消",
                        "stage_key": "cancelled",
                        "message": "任务已安全中断",
                    }
                )
            else:
                tail = "\n".join(self.job["logs"][-16:])
                self.job.update(
                    {
                        "state": "error",
                        "stage": "工作流失败",
                        "stage_key": "error",
                        "message": str(error),
                        "error": f"{error}\n{tail}".strip(),
                    }
                )
        finally:
            self.job_process = None

    async def cancel_job(self) -> None:
        if self.job["state"] != "running":
            return
        self.job["state"] = "cancelling"
        self.job["message"] = "正在安全中断当前阶段…"
        timeout = ClientTimeout(total=5)
        try:
            async with ClientSession(timeout=timeout) as session:
                await session.post(f"{self.comfy_url}/interrupt")
        except Exception:
            pass
        if self.job_process and self.job_process.returncode is None:
            try:
                await asyncio.wait_for(self.job_process.wait(), timeout=15)
            except TimeoutError:
                self.job_process.terminate()
        try:
            async with ClientSession(timeout=timeout) as session:
                await session.post(f"{self.comfy_url}/free", json={"unload_models": True, "free_memory": True})
        except Exception:
            pass

    async def cleanup(self) -> None:
        await self.cancel_job()
        if self.owns_comfy and self.comfy_process and self.comfy_process.returncode is None:
            self.comfy_process.terminate()
            try:
                await asyncio.wait_for(self.comfy_process.wait(), timeout=10)
            except TimeoutError:
                self.comfy_process.kill()
                await self.comfy_process.wait()


def resolve_asset(value: Any, expected: str | None = None) -> Path:
    path = Path(str(value or "")).resolve()
    if not path.is_file() or not (inside(path, COMFY_INPUT) or inside(path, COMFY_OUTPUT)):
        raise web.HTTPBadRequest(text="素材必须位于 ComfyUI/input 或 ComfyUI/output")
    kind = asset_kind(path)
    if kind is None or (expected and kind != expected):
        raise web.HTTPBadRequest(text=f"素材类型不正确，需要 {expected or '媒体'}")
    return path


async def index(_: web.Request) -> web.FileResponse:
    return web.FileResponse(INDEX_FILE)


async def cloud_connect(request: web.Request) -> web.Response:
    return_value = str(request.query.get("return") or "")
    origin = normalize_cloud_origin(return_value)
    if not origin:
        raise web.HTTPBadRequest(
            text="回跳地址无效：Cloudflare Pages 必须使用 HTTPS；本地预览仅允许 localhost/127.0.0.1。"
        )
    safe_origin = escape(origin, quote=True)
    body = f"""<!doctype html>
<html lang="zh-CN"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1">
<title>H3 Studio 本机配对</title><style>
*{{box-sizing:border-box}}body{{margin:0;min-height:100vh;display:grid;place-items:center;background:#0b0c0f;color:#f1f0eb;font-family:Inter,'Microsoft YaHei',sans-serif;padding:24px}}
main{{width:min(520px,100%);padding:28px;border:1px solid #303239;border-radius:20px;background:#17191e;box-shadow:0 28px 90px #0008}}
.tag{{font-size:11px;letter-spacing:.14em;color:#d8ff45}}h1{{font-size:23px;margin:12px 0 10px}}p,li{{font-size:13px;line-height:1.75;color:#a4a7af}}code{{color:#c7d9ff;word-break:break-all}}.actions{{display:flex;gap:10px;margin-top:22px}}button,a{{height:42px;display:inline-flex;align-items:center;justify-content:center;border-radius:10px;padding:0 16px;font-size:13px;text-decoration:none}}button{{border:0;background:#d8ff45;color:#101116;font-weight:800;cursor:pointer}}a{{border:1px solid #3a3c43;color:#c5c7cc}}
</style></head><body><main><div class="tag">LOCAL CONSENT · H3 STUDIO</div><h1>允许此网页连接本机引擎？</h1>
<p>请求来源：<br><code>{safe_origin}</code></p>
<ul><li>模型、提示词和媒体仍只在你的电脑上处理。</li><li>该来源将获得 H3 Studio API 权限；访问密钥只回传到浏览器地址片段，不会发送给 Pages 服务器。</li><li>请只授权你本人部署或信任的网页域名。</li></ul>
<form method="post" action="/cloud/connect"><input type="hidden" name="return" value="{safe_origin}"><div class="actions"><button type="submit">允许并返回创作系统</button><a href="{safe_origin}">取消</a></div></form>
</main></body></html>"""
    return web.Response(text=body, content_type="text/html")


async def cloud_connect_confirm(request: web.Request) -> web.StreamResponse:
    form = await request.post()
    origin = normalize_cloud_origin(str(form.get("return") or ""))
    if not origin:
        raise web.HTTPBadRequest(text="回跳地址无效")
    access_key = pair_cloud_origin(origin)
    port = int(request.app["studio_port"])
    endpoint = f"http://127.0.0.1:{port}"
    fragment = urlencode({"h3-engine": endpoint, "h3-pair": access_key})
    raise web.HTTPFound(location=f"{origin}/#{fragment}")


async def api_status(request: web.Request) -> web.Response:
    return web.json_response(request.app["studio"].snapshot())


async def api_setup_status(request: web.Request) -> web.Response:
    studio: Studio = request.app["studio"]
    studio.deployment = current_deployment(refresh_runtime=True)
    return web.json_response(studio.deployment)


async def api_setup_plan(request: web.Request) -> web.Response:
    studio: Studio = request.app["studio"]
    payload = await request.json()
    override = payload.get("hardware_override") if isinstance(payload.get("hardware_override"), dict) else None
    profile_id = str(payload.get("profile_id") or "") or None
    return web.json_response(
        build_plan(
            probe_hardware(),
            override=override,
            selected_profile=profile_id,
            runtime=studio.deployment.get("runtime"),
        )
    )


async def api_setup_apply(request: web.Request) -> web.Response:
    studio: Studio = request.app["studio"]
    if studio.job["state"] in {"running", "cancelling"}:
        raise web.HTTPConflict(text="请先等待当前任务完成，再切换部署规格")
    payload = await request.json()
    previous = studio.deployment["profile"]["id"]
    studio.deployment = apply_plan(payload)
    profile = studio.deployment["profile"]
    if profile["h3_support"] == "unavailable":
        if studio.deployment.get("capabilities", {}).get("upscale"):
            studio.service = "ready"
            finish_label = "NVIDIA Maxine" if profile["finish_engine"] == "maxine" else "通用视频终修"
            studio.service_message = f"{finish_label}已就绪；当前运行时未启用本地 H3"
        else:
            studio.service = "setup"
            studio.service_message = "未找到可用的视频终修运行时，请先完成部署"
    elif not studio.deployment["readiness"]["ready"]:
        studio.service = "setup"
        studio.service_message = "部署规格已保存；请按清单补全必需组件"
    elif studio.service == "setup":
        studio.service = "starting"
        studio.service_message = "配置已保存，正在启动本地 H3 引擎…"
        asyncio.create_task(studio.start_comfy())
    return web.json_response(
        {
            **studio.deployment,
            "restart_required": previous != profile["id"] and studio.comfy_process is not None,
        }
    )


async def api_repair_status(request: web.Request) -> web.Response:
    studio: Studio = request.app["studio"]
    reachable = await studio.comfy_ready()
    report = diagnose_environment(
        studio.deployment,
        service=studio.service,
        service_message=studio.service_message,
        comfy_reachable=reachable,
        owns_comfy=studio.owns_comfy,
        comfy_log=list(studio.comfy_log),
    )
    return web.json_response(report)


async def api_repair_run(request: web.Request) -> web.Response:
    studio: Studio = request.app["studio"]
    if studio.job["state"] in {"running", "cancelling"}:
        raise web.HTTPConflict(text="请先等待当前工作流结束，再执行环境修复")
    payload = await request.json()
    action = str(payload.get("action") or "")
    try:
        if action == "restart_engine":
            await studio.restart_comfy()
            result = {"message": studio.service_message}
        else:
            if action == "reset_runtime_caches" and await studio.comfy_ready():
                raise RuntimeError("缓存正在被引擎使用。请先点击左下角停止 Studio，再重新启动并在引擎报错状态下执行修复。")
            result = await asyncio.to_thread(run_repair, action, studio.deployment)
            studio.deployment = current_deployment(refresh_runtime=True)
        reachable = await studio.comfy_ready()
        report = diagnose_environment(
            studio.deployment,
            service=studio.service,
            service_message=studio.service_message,
            comfy_reachable=reachable,
            owns_comfy=studio.owns_comfy,
            comfy_log=list(studio.comfy_log),
        )
        return web.json_response({"ok": True, "result": result, "report": report})
    except (RuntimeError, ValueError, subprocess.SubprocessError) as exc:
        raise web.HTTPBadRequest(text=str(exc)) from exc


async def api_assets(request: web.Request) -> web.Response:
    return web.json_response({"assets": request.app["studio"].list_assets()})


async def api_asset_info(request: web.Request) -> web.Response:
    path = resolve_asset(request.query.get("path"))
    info = {"path": str(path), "kind": asset_kind(path), "size": path.stat().st_size}
    if info["kind"] == "video":
        info.update(probe_video(path))
    return web.json_response(info)


async def api_upload(request: web.Request) -> web.Response:
    reader = await request.multipart()
    field = await reader.next()
    if field is None or not field.filename:
        raise web.HTTPBadRequest(text="缺少素材文件")
    suffix = Path(field.filename).suffix.lower()
    kind = asset_kind(Path(field.filename))
    if kind is None:
        raise web.HTTPBadRequest(text="不支持该素材格式")
    destination_dir = UPLOAD_ROOT / kind
    destination_dir.mkdir(parents=True, exist_ok=True)
    destination = destination_dir / f"{datetime.now():%Y%m%d_%H%M%S}_{safe_stem(Path(field.filename).stem)}{suffix}"
    with destination.open("wb") as output:
        while chunk := await field.read_chunk(size=1024 * 1024):
            output.write(chunk)
    return web.json_response(
        {
            "name": destination.name,
            "path": str(destination.resolve()),
            "url": media_url(destination),
            "kind": kind,
            "size": destination.stat().st_size,
        }
    )


async def api_start(request: web.Request) -> web.Response:
    studio: Studio = request.app["studio"]
    if studio.service != "ready":
        raise web.HTTPConflict(text="ComfyUI 尚未就绪")
    if studio.job["state"] in {"running", "cancelling"}:
        raise web.HTTPConflict(text="已有任务正在运行")
    payload = await request.json()
    mode = str(payload.get("mode", ""))
    if mode not in WORKFLOW_MODES:
        raise web.HTTPBadRequest(text="未知工作流模板")
    profile = studio.deployment["profile"]
    capabilities = studio.deployment.get("capabilities") or {}
    if not capabilities.get(mode, mode == "upscale"):
        if mode == "upscale" and capabilities.get("upscale"):
            pass
        else:
            raise web.HTTPBadRequest(text=f"当前运行时未启用{WORKFLOW_MODES[mode]['label']}，请检查部署规格和组件")
    if profile["h3_support"] == "unavailable" and mode != "upscale":
        raise web.HTTPBadRequest(text="当前部署规格仅支持视频终修；请打开部署向导启用 H3 生成")
    capability = {"r2v": "r2v", "redraw": "redraw", "direct2k": "direct2k"}.get(mode)
    if capability and not profile[capability]:
        raise web.HTTPBadRequest(text=f"{profile['label']}未启用{WORKFLOW_MODES[mode]['label']}，请更换部署规格")

    resolution_key = "direct2k" if mode == "direct2k" else str(payload.get("resolution", profile["default_resolution"]))
    if resolution_key not in RESOLUTIONS:
        raise web.HTTPBadRequest(text="未知分辨率预设")
    resolution = RESOLUTIONS[resolution_key]
    duration = float(payload.get("duration", 5))
    if not 0 < duration <= 15:
        raise web.HTTPBadRequest(text="时长必须在 0–15 秒之间")
    if mode == "direct2k" and duration > 0.5:
        raise web.HTTPBadRequest(text="Direct 2K 在 24GB 显存上最多允许 0.5 秒")
    steps = int(payload.get("steps", 8))
    r2v_steps = int(payload.get("r2v_steps", 20))
    if not 4 <= steps <= 8 or not 8 <= r2v_steps <= 30:
        raise web.HTTPBadRequest(text="采样步数超出允许范围")

    config: dict[str, Any] = {
        "mode": mode,
        "width": resolution["width"],
        "height": resolution["height"],
        "duration": duration,
        "steps": steps,
        "r2v_steps": r2v_steps,
        "seed": int(payload.get("seed", 556589502035082)),
        "prompt": str(payload.get("prompt", REDRAW_PROMPT if mode == "redraw" else DEFAULT_PROMPT)).strip()[:6000],
        "turbo": bool(payload.get("turbo", True)),
        "redraw_denoise": float(payload.get("redraw_denoise", 0.35)),
        "finish_denoise": str(payload.get("finish_denoise", "medium")),
        "rtx_enabled": bool(payload.get("rtx_enabled", True)),
        "rtx_scale": float(payload.get("rtx_scale", 1.5)),
        "rtx_quality": str(payload.get("rtx_quality", "HIGHBITRATE_HIGH")),
        "ref_image_size": str(payload.get("ref_image_size", "match")),
        "finish_engine": profile["finish_engine"],
    }
    if not config["prompt"] and mode != "upscale":
        raise web.HTTPBadRequest(text="提示词不能为空")
    if config["finish_denoise"] not in {"none", "mild", "medium", "finish"}:
        raise web.HTTPBadRequest(text="未知终修降噪预设")
    if config["rtx_quality"] not in {"HIGHBITRATE_HIGH", "HIGHBITRATE_ULTRA"}:
        raise web.HTTPBadRequest(text="未知 RTX 质量预设")
    if not 0 < config["redraw_denoise"] <= 1 or not 1 < config["rtx_scale"] <= 4:
        raise web.HTTPBadRequest(text="重绘或 RTX 参数超出范围")

    if mode == "i2v":
        config["first_frame"] = str(resolve_asset(payload.get("first_frame"), "image"))
    elif mode == "l2v":
        config["last_frame"] = str(resolve_asset(payload.get("last_frame"), "image"))
    elif mode == "fl2v":
        config["first_frame"] = str(resolve_asset(payload.get("first_frame"), "image"))
        config["last_frame"] = str(resolve_asset(payload.get("last_frame"), "image"))
    elif mode == "r2v":
        references = payload.get("references", [])
        if not isinstance(references, list) or not references:
            raise web.HTTPBadRequest(text="多参考生成至少需要一个参考素材")
        counts = {"image": 0, "video": 0, "audio": 0}
        normalized = []
        for reference in references:
            path = resolve_asset(reference.get("path"))
            kind = asset_kind(path)
            counts[kind] += 1
            normalized.append({"path": str(path), "kind": kind, "use_audio": bool(reference.get("use_audio", True))})
        if counts["image"] > 9 or counts["video"] > 3 or counts["audio"] > 3:
            raise web.HTTPBadRequest(text="参考素材数量超过 H3 限制（9 图 / 3 视频 / 3 音频）")
        config["references"] = normalized
    elif mode in {"redraw", "upscale"}:
        input_video = resolve_asset(payload.get("input_video"), "video")
        if mode == "redraw" and not probe_video(input_video)["audio"]:
            raise web.HTTPBadRequest(text="H3 Concat AV 重绘需要带音轨的视频")
        config["input_video"] = str(input_video)

    job_id = uuid.uuid4().hex
    studio.job = studio.empty_job()
    studio.job.update(
        {
            "id": job_id,
            "state": "running",
            "stage": "准备工作流",
            "stage_key": "prepare",
            "progress": 2,
            "message": f"载入 {WORKFLOW_MODES[mode]['label']} 模板并清理显存",
            "started_at": time.time(),
            "mode": mode,
            "input": config.get("input_video") or config.get("first_frame"),
        }
    )
    studio.job_task = asyncio.create_task(studio.run_job(config))
    return web.json_response({"ok": True, "job": studio.snapshot()["job"]})


async def api_cancel(request: web.Request) -> web.Response:
    await request.app["studio"].cancel_job()
    return web.json_response({"ok": True})


async def api_shutdown(request: web.Request) -> web.Response:
    request.app["shutdown_event"].set()
    return web.json_response({"ok": True})


async def serve_media(request: web.Request) -> web.FileResponse:
    path = Path(request.query.get("path", "")).resolve()
    if not path.is_file() or not (inside(path, COMFY_INPUT) or inside(path, COMFY_OUTPUT)):
        raise web.HTTPNotFound()
    return web.FileResponse(path)


def build_installer_archive(output: Path) -> Path:
    builder = ROOT / "tools" / "build_h3_cn_installer.py"
    if not builder.is_file():
        raise RuntimeError("缺少一键安装包构建器 tools\\build_h3_cn_installer.py")
    spec = importlib.util.spec_from_file_location("h3_installer_builder", builder)
    if spec is None or spec.loader is None:
        raise RuntimeError("无法加载一键安装包构建器")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module.build_installer(output)


async def download_installer(request: web.Request) -> web.FileResponse:
    installer = ROOT / "H3-Studio-OneClick-CN.zip"
    if not installer.is_file():
        async with request.app["installer_lock"]:
            if not installer.is_file():
                try:
                    await asyncio.to_thread(build_installer_archive, installer)
                except Exception as exc:
                    raise web.HTTPServiceUnavailable(text=f"一键安装包生成失败：{exc}") from exc
    return web.FileResponse(installer, headers={"Content-Disposition": 'attachment; filename="H3-Studio-OneClick-CN.zip"'})


async def run_app(host: str, port: int, comfy_port: int, open_browser: bool) -> None:
    required = [INDEX_FILE, UNIVERSAL_RUNNER]
    if not all(path.is_file() for path in required):
        raise RuntimeError("MiniMax H3 Studio 工作区不完整")
    studio = Studio(comfy_port)
    app = web.Application(client_max_size=20 * 1024**3, middlewares=[cloud_access_middleware])
    app["studio"] = studio
    app["studio_port"] = port
    app["shutdown_event"] = asyncio.Event()
    app["installer_lock"] = asyncio.Lock()
    app.add_routes(
        [
            web.get("/", index),
            web.get("/cloud/connect", cloud_connect),
            web.post("/cloud/connect", cloud_connect_confirm),
            web.get("/api/status", api_status),
            web.get("/api/setup/status", api_setup_status),
            web.post("/api/setup/plan", api_setup_plan),
            web.post("/api/setup/apply", api_setup_apply),
            web.get("/api/repair/status", api_repair_status),
            web.post("/api/repair/run", api_repair_run),
            web.get("/api/assets", api_assets),
            web.get("/api/asset-info", api_asset_info),
            web.post("/api/upload", api_upload),
            web.post("/api/start", api_start),
            web.post("/api/cancel", api_cancel),
            web.post("/api/shutdown", api_shutdown),
            web.get("/media", serve_media),
            web.get("/downloads/H3-Studio-OneClick-CN.zip", download_installer),
        ]
    )
    runner = web.AppRunner(app)
    await runner.setup()
    await web.TCPSite(runner, host, port).start()
    url = f"http://{host}:{port}"
    print(f"H3 Studio: {url}", flush=True)
    if open_browser:
        asyncio.get_running_loop().call_later(1.0, webbrowser.open, url)
    comfy_task = asyncio.create_task(studio.start_comfy())
    try:
        await app["shutdown_event"].wait()
    finally:
        if not comfy_task.done():
            comfy_task.cancel()
        await studio.cleanup()
        await runner.cleanup()


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8191)
    parser.add_argument("--comfy-port", type=int, default=8189)
    parser.add_argument("--no-browser", action="store_true")
    args = parser.parse_args()
    try:
        asyncio.run(run_app(args.host, args.port, args.comfy_port, not args.no_browser))
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()
