import os
import time
import requests

def launch_comfy_server(version_id: str, server_type: str, duration: int, timeout: int = 300):
    """Launch RunComfy server and wait until ready"""
    user_id = os.getenv("RUNCOMFY_USER_ID")
    api_token = os.getenv("API_TOKEN")
    
    if not user_id or not api_token:
        raise ValueError("RUNCOMFY_USER_ID and API_TOKEN must be set")
    
    # Create server
    url = f"https://beta-api.runcomfy.net/prod/api/users/{user_id}/servers"
    headers = {"Authorization": f"Bearer {api_token}"}
    payload = {
        "workflow_version_id": version_id,
        "server_type": server_type,
        "estimated_duration": duration
    }
    
    resp = requests.post(url, json=payload, headers=headers)
    resp.raise_for_status()
    data = resp.json()
    server_id = data["server_id"]
    
    # Poll until ready
    check_url = f"{url}/{server_id}"
    start_time = time.time()
    
    while time.time() - start_time < timeout:
        resp = requests.get(check_url, headers=headers)
        resp.raise_for_status()
        server_data = resp.json()
        
        if server_data["current_status"] == "Ready":
            return {
                "server_id": server_id,
                "comfy_url": server_data["main_service_url"],
                "status": "Ready",
                "usage_instruction": (
                    f"请使用 run_comfy_workflow 工具，并传入此 server_id ({server_id}) "
                    "以及你的 workflow_api.json 来生成图片。"
                    "用完后请务必调用 stop_comfy_env 工具关闭机器。"
                )
            }
        
        time.sleep(5)
    
    raise TimeoutError(f"Server {server_id} not ready within {timeout}s")