import io
import json
import uuid
import urllib.request
import urllib.parse
import websocket

from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import StreamingResponse

# --- Cấu hình --- 
# Địa chỉ máy chủ ComfyUI, mặc định là localhost
COMFYUI_SERVER_ADDRESS = "127.0.0.1:8188"
# ID của client, dùng để nhận diện trong WebSocket
CLIENT_ID = str(uuid.uuid4())

# Khởi tạo ứng dụng FastAPI
app = FastAPI(title="Clarity AI Upscaler API")

def get_image(filename, subfolder, folder_type):
    """Lấy file ảnh từ máy chủ ComfyUI."""
    data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
    url_values = urllib.parse.urlencode(data)
    with urllib.request.urlopen(f"http://{COMFYUI_SERVER_ADDRESS}/view?{url_values}") as response:
        return response.read()

def queue_prompt(prompt_workflow):
    """Gửi yêu cầu thực thi quy trình đến ComfyUI."""
    p = {"prompt": prompt_workflow, "client_id": CLIENT_ID}
    data = json.dumps(p).encode('utf-8')
    req = urllib.request.Request(f"http://{COMFYUI_SERVER_ADDRESS}/prompt", data=data)
    return json.loads(urllib.request.urlopen(req).read())

def upload_image(image_bytes: bytes, filename: str = "input_image.png"):
    """Tải ảnh lên thư mục input của ComfyUI."""
    import requests
    files = {"image": (filename, image_bytes, 'image/png'), "overwrite": (None, 'true')}
    response = requests.post(f"http://{COMFYUI_SERVER_ADDRESS}/upload/image", files=files)
    if response.status_code == 200:
        return response.json()
    else:
        raise HTTPException(status_code=500, detail=f"Failed to upload image: {response.text}")

def get_history(prompt_id):
    """Lấy lịch sử thực thi của một prompt."""
    with urllib.request.urlopen(f"http://{COMFYUI_SERVER_ADDRESS}/history/{prompt_id}") as response:
        return json.loads(response.read())

def track_execution_and_get_output(prompt_id):
    """Theo dõi tiến trình qua WebSocket và lấy ảnh kết quả."""
    ws_url = f"ws://{COMFYUI_SERVER_ADDRESS}/ws?clientId={CLIENT_ID}"
    ws = websocket.WebSocket()
    ws.connect(ws_url)
    
    while True:
        out = ws.recv()
        if isinstance(out, str):
            message = json.loads(out)
            if message['type'] == 'executing':
                data = message['data']
                if data['node'] is None and data['prompt_id'] == prompt_id:
                    break # Đã thực thi xong
        else:
            continue # Bỏ qua các tin nhắn không phải dạng text
    ws.close()

    history = get_history(prompt_id)[prompt_id]
    for node_id, node_output in history['outputs'].items():
        if 'images' in node_output:
            for image in node_output['images']:
                if image['type'] == 'output':
                    image_data = get_image(image['filename'], image['subfolder'], image['type'])
                    return image_data
    raise HTTPException(status_code=500, detail="Không tìm thấy ảnh kết quả.")

@app.post("/upscale/",
          summary="Nâng cấp và tinh chỉnh hình ảnh",
          response_description="Hình ảnh đã được xử lý ở định dạng PNG")
async def upscale_image(file: UploadFile = File(..., description="File ảnh cần xử lý.")):
    """
    Nhận một file ảnh, xử lý nó thông qua quy trình ComfyUI, và trả về kết quả.
    """
    try:
        # 1. Đọc file workflow từ định dạng API
        with open("workflow_api.json", "r", encoding="utf-8") as f:
            prompt_workflow = json.load(f)

        # 2. Đọc và tải ảnh đầu vào lên ComfyUI
        image_bytes = await file.read()
        upload_response = upload_image(image_bytes)
        input_filename = upload_response['name']

        # 3. Tìm node LoadImage (ID=17) và cập nhật tên file
        # LƯU Ý: ID '17' là ID của node LoadImage trong workflow của bạn.
        # Nếu bạn thay đổi workflow, ID này có thể thay đổi.
        prompt_workflow["17"]["inputs"]["image"] = input_filename

        # 4. Gửi quy trình để thực thi
        queue_response = queue_prompt(prompt_workflow)
        prompt_id = queue_response['prompt_id']

        # 5. Theo dõi tiến trình và lấy ảnh kết quả
        result_image_bytes = track_execution_and_get_output(prompt_id)

        # 6. Trả về ảnh kết quả
        return StreamingResponse(io.BytesIO(result_image_bytes), media_type="image/png")

    except FileNotFoundError:
        raise HTTPException(status_code=500, detail="Không tìm thấy file workflow_api.json.")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Đã xảy ra lỗi: {str(e)}")

# Để chạy máy chủ này, mở terminal và gõ lệnh:
# uvicorn api_server:app --reload