api/workers.py

import time
from datetime import datetime
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel
from typing import Optional, Dict, Any

from config import MASTER_SECRET, HEARTBEAT_TIMEOUT
from core.state import workers, add_worker, update_worker, get_worker_by_id
from core.task_manager import get_next_task_for_worker

router = APIRouter(prefix="/api/workers", tags=["workers"])

class WorkerRegisterRequest(BaseModel):
    worker_id: str
    worker_type: str
    capabilities: Dict[str, Any]
    cpu_cores: int
    memory_gb: float

class WorkerHeartbeatRequest(BaseModel):
    worker_id: str
    token: str
    current_load: float
    memory_used_gb: float
    gpu_used_gb: Optional[float] = None
    current_task: Optional[str] = None
    backup_task: Optional[str] = None
    progress: Optional[float] = 0.0
    status_detail: Optional[str] = None

def generate_token(worker_id: str) -> str:
    import hashlib
    salt = str(time.time())
    return hashlib.sha256(f"{worker_id}:{MASTER_SECRET}:{salt}".encode()).hexdigest()[:32]

def verify_token(worker_id: str, token: str) -> bool:
    if worker_id not in workers:
        return False
    return workers[worker_id].get("token") == token

@router.post("/register")
async def register_worker(data: WorkerRegisterRequest):
    try:
        worker_id = data.worker_id
        if worker_id in workers:
            workers[worker_id]["token"] = generate_token(worker_id)
            workers[worker_id]["last_heartbeat"] = time.time()
            workers[worker_id]["status"] = "idle"
            workers[worker_id]["status_detail"] = "idle"
            workers[worker_id]["info"] = data.dict()
            print(f"✅ Worker {worker_id} re-registered")
            return {"status": "updated", "worker_id": worker_id, "token": workers[worker_id]["token"]}
        token = generate_token(worker_id)
        workers[worker_id] = {
            "info": data.dict(),
            "token": token,
            "last_heartbeat": time.time(),
            "status": "idle",
            "status_detail": "idle",
            "current_task": None,
            "backup_task": None,
            "progress": 0.0,
            "completed_tasks": 0,
            "registered_at": datetime.now().isoformat(),
        }
        print(f"✅ Worker {worker_id} registered")
        return {"status": "registered", "worker_id": worker_id, "token": token}
    except Exception as e:
        print(f"❌ Register error: {e}")
        raise HTTPException(500, str(e))

@router.post("/heartbeat")
async def worker_heartbeat(data: WorkerHeartbeatRequest):
    try:
        if data.worker_id not in workers:
            raise HTTPException(404, "Worker not found")
        if not verify_token(data.worker_id, data.token):
            raise HTTPException(401, "Invalid token")
        w = workers[data.worker_id]
        w["last_heartbeat"] = time.time()
        w["status"] = "busy" if data.current_task else "idle"
        w["current_task"] = data.current_task
        w["backup_task"] = data.backup_task
        w["progress"] = data.progress or 0.0
        w["status_detail"] = data.status_detail or "idle"
        return {"status": "ok"}
    except HTTPException:
        raise
    except Exception as e:
        print(f"❌ Heartbeat error: {e}")
        raise HTTPException(500, str(e))

@router.get("/list")
async def list_workers():
    try:
        now = time.time()
        result = []
        for wid, w in workers.items():
            is_online = (now - w.get("last_heartbeat", 0)) < HEARTBEAT_TIMEOUT
            result.append({
                "worker_id": wid,
                "worker_type": w.get("info", {}).get("worker_type", "unknown"),
                "status": "online" if is_online else "offline",
                "status_detail": w.get("status_detail", "idle"),
                "current_task": w.get("current_task"),
                "backup_task": w.get("backup_task"),
                "progress": w.get("progress", 0),
                "completed_tasks": w.get("completed_tasks", 0),
                "last_heartbeat": datetime.fromtimestamp(w.get("last_heartbeat", 0)).isoformat() if w.get("last_heartbeat") else None,
            })
        return {"workers": result, "total": len(result)}
    except Exception as e:
        print(f"❌ List error: {e}")
        return {"workers": [], "total": 0}

@router.get("/next-task")
async def get_next_task(worker_id: str, token: str):
    try:
        if worker_id not in workers:
            raise HTTPException(404, "Worker not found")
        if not verify_token(worker_id, token):
            raise HTTPException(401, "Invalid token")
        task = get_next_task_for_worker(worker_id)
        if task:
            workers[worker_id]["current_task"] = task["task_id"]
            workers[worker_id]["status"] = "busy"
            return {"current_task": task, "backup_task": None}
        else:
            return {"current_task": None, "backup_task": None}
    except HTTPException:
        raise
    except Exception as e:
        print(f"❌ Get next task error: {e}")
        raise HTTPException(500, str(e))