#<--BẮT ĐẦU TOÀN BỘ CODE CHO FILE load_balancer.py (WORKER ĐẦU MỐI) - PHIÊN BẢN CUỐI CÙNG-->
import os
import asyncio
import aiohttp
import random
from fastapi import FastAPI, HTTPException, Request, Header
from pydantic import BaseModel, Field
from typing import List, Dict, Any

# --- Cấu hình ---
# Đọc key xác thực bí mật từ Secrets của Hugging Face. 
# Bạn cần vào Settings -> Secrets của Space này để thêm một Secret mới tên là AUTH_KEY
API_AUTH_KEY = os.getenv("AUTH_KEY", "default-secret-key-change-me")

# !!! THAY THẾ BẰNG DANH SÁCH URL WORKER THỢ THẬT CỦA BẠN !!!
WORKER_URLS = [
    # Dán danh sách URL của các worker thợ vào đây. Ví dụ:
    "https://whf-manager-api-key-sheet-worker-checker-01.hf.space",
    "https://whf-manager-api-key-sheet-elevenlabs-worker-02.hf.space",
    "https://whf-manager-api-key-sheet-elevenlabs-worker-03.hf.space",
    "https://whf-manager-api-key-sheet-elevenlabs-worker-04.hf.space",
    "https://whf-manager-api-key-sheet-elevenlabs-worker-05.hf.space",
    "https://whf-manager-api-key-sheet-elevenlabs-worker-06.hf.space",
    "https://whf-manager-api-key-sheet-elevenlabs-worker-07.hf.space",
    "https://whf-manager-api-key-sheet-elevenlabs-worker-08.hf.space",
    "https://whf-manager-api-key-sheet-elevenlabs-worker-09.hf.space",
    "https://whf-manager-api-key-sheet-elevenlabs-worker-10.hf.space",
    "https://whf-manager-api-key-sheet-elevenlabs-worker-11.hf.space",
    "https://whf-manager-api-key-sheet-elevenlabs-worker-12.hf.space",
    "https://whf-manager-api-key-sheet-elevenlabs-worker-13.hf.space",
    "https://whf-manager-api-key-sheet-elevenlabs-worker-14.hf.space",
    "https://whf-manager-api-key-sheet-elevenlabs-worker-15.hf.space",
    "https://whf-manager-api-key-sheet-elevenlabs-worker-16.hf.space",
    "https://whf-manager-api-key-sheet-elevenlabs-worker-17.hf.space",
    "https://whf-manager-api-key-sheet-elevenlabs-worker-18.hf.space",
    "https://whf-manager-api-key-sheet-elevenlabs-worker-19.hf.space"
]

# --- Pydantic Models ---
class APIKeyPayload(BaseModel):
    api_keys: List[str] = Field(..., min_items=1)

# --- Khởi tạo ứng dụng FastAPI ---
app = FastAPI(
    title="Maintenance Load Balancer (Đầu mối)",
    description="Xác thực, nhận yêu cầu từ client và phân phối đến các worker thợ.",
    version="2.0.0",
)

# --- Logic của Load Balancer ---
worker_load = {url: 0 for url in WORKER_URLS}
load_lock = asyncio.Lock()

async def select_worker() -> str:
    if not WORKER_URLS: raise ValueError("Danh sách WORKER_URLS chưa được cấu hình!")
    async with load_lock:
        min_load = min(worker_load.values())
        least_busy_workers = [url for url, load in worker_load.items() if load == min_load]
    return random.choice(least_busy_workers)

async def forward_request_to_worker(worker_url: str, payload: Dict[str, Any], auth_key: str):
    endpoint = "/check-credits" # Endpoint trên worker thợ
    target_url = f"{worker_url.rstrip('/')}{endpoint}"
    
    # Chuyển tiếp cả header xác thực đến worker thợ
    headers = {"Content-Type": "application/json", "x-api-key": auth_key}
    
    async with load_lock: worker_load[worker_url] += 1
    print(f"Forwarding to {worker_url}. Current loads: {worker_load}")
    
    try:
        async with aiohttp.ClientSession() as session:
            async with session.post(target_url, json=payload, headers=headers, timeout=125) as response:
                response.raise_for_status()
                return await response.json()
    finally:
        async with load_lock: worker_load[worker_url] -= 1
        print(f"Request to {worker_url} finished. Current loads: {worker_load}")

# --- API Endpoint của Load Balancer ---
@app.post("/dispatch-maintenance")
async def dispatch_maintenance_endpoint(payload: APIKeyPayload, x_api_key: str = Header(...)):
    """Endpoint DUY NHẤT mà Client Tool sẽ gọi."""
    # 1. Xác thực request từ client
    if x_api_key != API_AUTH_KEY:
        raise HTTPException(status_code=401, detail="Invalid API Key for Load Balancer")
    
    # 2. Chọn worker và chuyển tiếp request
    try:
        worker_to_use = await select_worker()
        result = await forward_request_to_worker(worker_to_use, payload.dict(), API_AUTH_KEY)
        return result
    except ValueError as e:
        raise HTTPException(status_code=503, detail=str(e))
    except aiohttp.ClientResponseError as e:
        raise HTTPException(status_code=e.status, detail=f"Worker thợ báo lỗi: {e.message}")
    except Exception as e:
        raise HTTPException(status_code=504, detail=f"Lỗi kết nối đến worker thợ: {str(e)}")

@app.get("/")
def read_root():
    return {"message": "Load Balancer is running."}

#<--KẾT THÚC TOÀN BỘ CODE CHO FILE load_balancer.py -->