#!/usr/bin/env python3
"""停止 RunComfy 机器实例

用法:
    python stop_comfy_env.py --server-id <SERVER_ID>
"""

import argparse
import os
import sys
from pathlib import Path

import requests
from dotenv import load_dotenv

load_dotenv(Path(__file__).parent.parent.parent / ".env")

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",
}


def stop_machine(server_id: str) -> dict:
    """删除指定的机器实例
    
    Returns:
        dict: {
            "server_id": str,
            "status": str,  # "Deleted" 或错误信息
            "message": str
        }
    """
    try:
        resp = requests.delete(
            f"{BASE_URL}/users/{USER_ID}/servers/{server_id}",
            headers=HEADERS
        )
        
        if resp.status_code == 200:
            return {
                "server_id": server_id,
                "status": "Deleted",
                "message": f"机器 {server_id} 已成功删除"
            }
        elif resp.status_code == 404:
            return {
                "server_id": server_id,
                "status": "NotFound",
                "message": f"机器 {server_id} 不存在或已被删除"
            }
        else:
            return {
                "server_id": server_id,
                "status": "Error",
                "message": f"HTTP {resp.status_code}: {resp.text}"
            }
            
    except Exception as e:
        return {
            "server_id": server_id,
            "status": "Error",
            "message": f"请求失败: {str(e)}"
        }


def main():
    parser = argparse.ArgumentParser(description="停止 RunComfy 机器实例")
    parser.add_argument("--server-id", required=True, help="要关闭的机器 ID")
    args = parser.parse_args()

    if not USER_ID or not API_TOKEN:
        print("ERROR: 请设置 RUNCOMFY_USER_ID 和 API_TOKEN 环境变量")
        sys.exit(1)

    result = stop_machine(args.server_id)
    print(f"状态: {result['status']}")
    print(f"消息: {result['message']}")
    
    if result["status"] in ("Deleted", "NotFound"):
        sys.exit(0)
    else:
        sys.exit(1)


if __name__ == "__main__":
    main()
