#!/usr/bin/env python3
"""RunComfy Workflow Executor (仅执行，不启动/关闭机器)"""

import argparse
import json
import os
import sys
import urllib.parse
import uuid
from pathlib import Path

import requests
import websocket
from dotenv import load_dotenv

load_dotenv()

BASE_URL = "https://beta-api.runcomfy.net/prod/api"
USER_ID = os.getenv("RUNCOMFY_USER_ID")
API_TOKEN = os.getenv("API_TOKEN")
HEADERS = {"Authorization": f"Bearer {API_TOKEN}", "Content-Type": "application/json"}

SUBDIR_UPLOAD_MAP = {
    "images": {"type": "input", "subfolder": ""},
    "loras": {"type": "input", "subfolder": "loras"},
    "checkpoints": {"type": "input", "subfolder": "checkpoints"},
    "vae": {"type": "input", "subfolder": "vae"},
    "controlnet": {"type": "input", "subfolder": "controlnet"},
    "upscale": {"type": "input", "subfolder": "upscale_models"},
}


def get_server_url(server_id: str) -> str:
    resp = requests.get(f"{BASE_URL}/users/{USER_ID}/servers/{server_id}", headers=HEADERS)
    resp.raise_for_status()
    data = resp.json()
    if data.get("current_status") != "Ready":
        raise Exception(f"机器未就绪: {data.get('current_status')}")
    return data["main_service_url"].rstrip("/")


def upload_file(comfy_url: str, file_path: Path, file_type: str = "input", subfolder: str = "") -> str:
    with open(file_path, "rb") as f:
        files = [("image", (file_path.name, f, "application/octet-stream"))]
        data = {"overwrite": "true", "type": file_type, "subfolder": subfolder}
        resp = requests.post(f"{comfy_url}/upload/image", data=data, files=files)
    resp.raise_for_status()
    server_name = resp.json()["name"]
    print(f"  上传: {file_path.name} → {subfolder}/{server_name}" if subfolder else f"  上传: {file_path.name}")
    return server_name


def upload_input_dir(comfy_url: str, input_dir: Path) -> dict[str, str]:
    if not input_dir.exists():
        print(f"  input 目录不存在: {input_dir}")
        return {}
    uploaded = {}
    ALL_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".mp4", ".avi", ".mov", ".webm", 
                ".safetensors", ".ckpt", ".pt", ".pth", ".gguf"}
    for f in input_dir.iterdir():
        if f.is_file() and f.suffix.lower() in ALL_EXTS:
            server_name = upload_file(comfy_url, f, "input", "")
            uploaded[f.name] = server_name
    for subdir in input_dir.iterdir():
        if not subdir.is_dir():
            continue
        mapping = SUBDIR_UPLOAD_MAP.get(subdir.name, {"type": "input", "subfolder": subdir.name})
        for f in subdir.iterdir():
            if f.is_file() and f.suffix.lower() in ALL_EXTS:
                server_name = upload_file(comfy_url, f, mapping["type"], mapping["subfolder"])
                uploaded[f.name] = server_name
    return uploaded


def submit_prompt(comfy_url: str, workflow_api: dict, client_id: str) -> str:
    payload = {"prompt": workflow_api, "client_id": client_id}
    resp = requests.post(f"{comfy_url}/prompt", json=payload)
    resp.raise_for_status()
    data = resp.json()
    if data.get("node_errors"):
        print(f"  节点错误: {data['node_errors']}")
    print(f"任务已提交: {data['prompt_id']}")
    return data["prompt_id"]


def wait_for_completion(comfy_url: str, client_id: str, prompt_id: str, timeout: int = 600):
    scheme = "wss" if comfy_url.startswith("https") else "ws"
    ws_url = f"{scheme}://{comfy_url.split('://', 1)[-1]}/ws?clientId={client_id}"
    print("WebSocket 监听中...")
    ws = websocket.WebSocket()
    ws.settimeout(timeout)
    ws.connect(ws_url)
    try:
        while True:
            out = ws.recv()
            if not out or isinstance(out, bytes):
                continue
            msg = json.loads(out)
            if msg.get("type") == "executing":
                data = msg.get("data", {})
                if data.get("prompt_id") == prompt_id and data.get("node") is None:
                    print("  执行完成")
                    break
                if data.get("node"):
                    print(f"  执行节点: {data['node']}")
            elif msg.get("type") == "progress":
                data = msg.get("data", {})
                print(f"  进度: {data.get('value', 0)}/{data.get('max', 1)}")
            elif msg.get("type") == "execution_error":
                if msg.get("data", {}).get("prompt_id") == prompt_id:
                    raise Exception(f"执行错误: {msg['data'].get('exception_message')}")
    finally:
        ws.close()


def download_outputs(comfy_url: str, prompt_id: str, output_dir: Path) -> list[str]:
    resp = requests.get(f"{comfy_url}/history/{prompt_id}")
    resp.raise_for_status()
    outputs = resp.json().get(prompt_id, {}).get("outputs", {})
    output_dir.mkdir(parents=True, exist_ok=True)
    saved = []
    for node_output in outputs.values():
        if "images" in node_output:
            for img in node_output["images"]:
                params = {"filename": img["filename"], "subfolder": img.get("subfolder", ""),
                         "type": img.get("temp") or img.get("type", "output")}
                resp = requests.get(f"{comfy_url}/view?{urllib.parse.urlencode(params)}")
                resp.raise_for_status()
                out_path = output_dir / img["filename"]
                out_path.write_bytes(resp.content)
                print(f"  图片: {out_path}")
                saved.append(str(out_path))
        if "gifs" in node_output:
            for video in node_output["gifs"]:
                params = {"filename": video["filename"], "subfolder": video.get("subfolder", ""),
                         "format": video.get("format", "mp4")}
                resp = requests.get(f"{comfy_url}/view?{urllib.parse.urlencode(params)}")
                resp.raise_for_status()
                out_path = output_dir / video["filename"]
                out_path.write_bytes(resp.content)
                print(f"  视频: {out_path}")
                saved.append(str(out_path))
    return saved


def main():
    parser = argparse.ArgumentParser(description="RunComfy workflow executor (不启动/关闭机器)")
    parser.add_argument("--server-id", required=True, help="已启动的机器 ID")
    parser.add_argument("--workflow", required=True, help="workflow_api.json 路径")
    parser.add_argument("--input-dir", default="input", help="输入文件目录，默认 input/")
    parser.add_argument("--output-dir", default="output", help="结果下载目录，默认 output/")
    parser.add_argument("--skip-upload", action="store_true", help="跳过文件上传")
    args = parser.parse_args()
    
    if not USER_ID or not API_TOKEN:
        print("ERROR: 请设置 RUNCOMFY_USER_ID 和 API_TOKEN 环境变量")
        sys.exit(1)
    
    workflow_path = Path(args.workflow)
    if not workflow_path.exists():
        print(f"ERROR: 文件不存在: {workflow_path}")
        sys.exit(1)
    
    with open(workflow_path, "r", encoding="utf-8") as f:
        workflow_api = json.load(f)
    
    client_id = str(uuid.uuid4())
    
    try:
        print(f"验证机器状态: {args.server_id}")
        comfy_url = get_server_url(args.server_id)
        print(f"  ComfyUI URL: {comfy_url}")
        
        if not args.skip_upload:
            print(f"\n上传 input 目录: {args.input_dir}")
            upload_input_dir(comfy_url, Path(args.input_dir))
        
        print(f"\n提交 workflow...")
        prompt_id = submit_prompt(comfy_url, workflow_api, client_id)
        wait_for_completion(comfy_url, client_id, prompt_id)
        
        print(f"\n下载结果...")
        saved = download_outputs(comfy_url, prompt_id, Path(args.output_dir))
        print(f"\n完成，共 {len(saved)} 个文件")
        print(f"注意: 机器 {args.server_id} 未自动关闭，请手动处理")
    except Exception as e:
        print(f"\n错误: {e}")
        print(f"机器 {args.server_id} 未自动关闭，请手动处理")
        sys.exit(1)


if __name__ == "__main__":
    main()
