import os
import uvicorn
import requests
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional, List, Dict, Any
from dotenv import load_dotenv

env_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), ".env")
if os.path.exists(env_path):
    load_dotenv(env_path)
else:
    load_dotenv()

app = FastAPI(title="RunComfy Check Status Proxy")

class CheckStatusRequest(BaseModel):
    # 可以给一个可选参数，不填则查所有
    server_id: Optional[str] = None

@app.post("/")
def check_status(req: CheckStatusRequest):
    user_id = os.getenv("RUNCOMFY_USER_ID")
    api_token = os.getenv("API_TOKEN")
    
    if not user_id or not api_token:
        raise HTTPException(status_code=500, detail="RUNCOMFY_USER_ID and API_TOKEN must be set in .env")
    
    url = f"https://beta-api.runcomfy.net/prod/api/users/{user_id}/servers"
    headers = {"Authorization": f"Bearer {api_token}"}
    
    try:
        if req.server_id:
            # Check specific server
            check_url = f"{url}/{req.server_id}"
            resp = requests.get(check_url, headers=headers)
            resp.raise_for_status()
            server_data = resp.json()
            return {"servers": [server_data], "total": 1}
        else:
            # Check all servers
            resp = requests.get(url, headers=headers)
            resp.raise_for_status()
            servers = resp.json()
            return {"servers": servers, "total": len(servers)}
            
    except requests.exceptions.RequestException as e:
        error_msg = str(e)
        if hasattr(e, 'response') and e.response is not None:
            error_msg += f" - Response: {e.response.text}"
        raise HTTPException(status_code=500, detail=error_msg)

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", type=int, default=8102)
    args = parser.parse_args()
    uvicorn.run(app, host="0.0.0.0", port=args.port)