app/routers/workers.py

from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.orm import Session
from app.database import get_db
from app.models.process import Worker
from app.models.user import AdminUser, User
from app.routers.auth import get_current_user
from app.routers.admin import get_current_admin
from app.utils.notifications import create_admin_notification
from pydantic import BaseModel
from typing import List, Optional
from datetime import datetime
import uuid

from app.websocket import ws_manager

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

class WorkerRegister(BaseModel):
    hostname: str
    name: Optional[str] = None
    status: Optional[str] = "online"

class WorkerResponse(BaseModel):
    id: int
    hostname: str
    name: Optional[str]
    status: str
    last_seen: datetime

    class Config:
        orm_mode = True

class WorkerUpdate(BaseModel):
    name: Optional[str] = None
    status: Optional[str] = None

@router.post("/register", response_model=WorkerResponse)
def register_worker(data: WorkerRegister, db: Session = Depends(get_db)):
    worker = db.query(Worker).filter(Worker.hostname == data.hostname).first()
    is_new = worker is None
    previous_status = None if is_new else worker.status
    if is_new:
        worker = Worker(
            hostname=data.hostname,
            name=data.name or data.hostname,
            status=data.status,
            last_seen=datetime.utcnow()
        )
        db.add(worker)
    else:
        worker.status = data.status
        worker.last_seen = datetime.utcnow()
        if data.name:
            worker.name = data.name
            
    db.commit()
    db.refresh(worker)
    
    if is_new:
        create_admin_notification(
            db, 
            "Worker Registrado", 
            f"El equipo worker '{worker.hostname}' se ha registrado en la red."
        )
    elif previous_status == "offline" and worker.status in ["online", "busy"]:
        create_admin_notification(
            db,
            "Worker Conectado",
            f"El equipo worker '{worker.name or worker.hostname}' se ha conectado (estado: {worker.status})."
        )
        
    ws_manager.broadcast_sync({"type": "refresh", "tab": "workers"})
    return worker

@router.get("", response_model=List[WorkerResponse])
def get_workers(db: Session = Depends(get_db), current_admin: AdminUser = Depends(get_current_admin)):
    if not current_admin.perm_workers:
        raise HTTPException(status_code=403, detail="No tiene permisos para gestionar workers")
    # Marcar automáticamente como 'offline' los que no han reportado en 2 minutos
    workers = db.query(Worker).order_by(Worker.id).all()
    now = datetime.utcnow()
    updated = False
    for w in workers:
        if w.status in ["online", "busy"] and (now - w.last_seen).total_seconds() > 120:
            w.status = "offline"
            updated = True
    if updated:
        db.commit()
        workers = db.query(Worker).order_by(Worker.id).all()
    return workers

@router.put("/{worker_id}", response_model=WorkerResponse)
def update_worker(worker_id: int, data: WorkerUpdate, db: Session = Depends(get_db), current_admin: AdminUser = Depends(get_current_admin)):
    if not current_admin.perm_workers:
        raise HTTPException(status_code=403, detail="No tiene permisos para gestionar workers")
    worker = db.query(Worker).filter(Worker.id == worker_id).first()
    if not worker:
        raise HTTPException(status_code=404, detail="Worker no encontrado")
    if data.name is not None:
        worker.name = data.name
    if data.status is not None:
        worker.status = data.status
        if data.status == "suspended":
            worker.is_active = False
            worker.command = "stop"
            try:
                from tasks import celery_app
                celery_app.control.cancel_consumer('celery', destination=[f'celery@{worker.hostname}'])
            except Exception as e:
                print(f"[Worker] Error cancelando consumidor: {e}")
        elif data.status == "online":
            worker.is_active = True
            worker.command = "start"
            try:
                from tasks import celery_app
                celery_app.control.add_consumer('celery', destination=[f'celery@{worker.hostname}'])
            except Exception as e:
                print(f"[Worker] Error agregando consumidor: {e}")
    db.commit()
    db.refresh(worker)
    
    # Registrar en notificaciones de admin
    action_desc = f"actualizado (estado: {worker.status})"
    if data.status == "suspended":
        action_desc = "suspendido"
    elif data.status == "online":
        action_desc = "activado"
    
    create_admin_notification(
        db,
        title=f"Worker {action_desc}",
        message=f"El Worker '{worker.name}' (ID: {worker.id}, Host: {worker.hostname}) ha sido {action_desc} por el administrador '{current_admin.username}'.",
        category="worker"
    )
    
    ws_manager.broadcast_sync({"type": "refresh", "tab": "workers"})
    return worker

@router.delete("/{worker_id}")
def delete_worker(worker_id: int, db: Session = Depends(get_db), current_admin: AdminUser = Depends(get_current_admin)):
    if not current_admin.perm_workers:
        raise HTTPException(status_code=403, detail="No tiene permisos para gestionar workers")
    worker = db.query(Worker).filter(Worker.id == worker_id).first()
    if not worker:
        raise HTTPException(status_code=404, detail="Worker no encontrado")
    
    db.delete(worker)
    db.commit()
    
    create_admin_notification(
        db,
        title="Worker eliminado",
        message=f"El Worker '{worker.name}' (ID: {worker_id}, Host: {worker.hostname}) ha sido eliminado permanentemente por el administrador '{current_admin.username}'.",
        category="worker"
    )
    
    ws_manager.broadcast_sync({"type": "refresh", "tab": "workers"})
    return {"message": "Worker eliminado exitosamente"}

# ==========================================
# BYOK / LOCAL WORKERS ENDPOINTS
# ==========================================

@router.get("/my-worker")
def get_my_worker(db: Session = Depends(get_db), current_user: User = Depends(get_current_user)):
    worker = db.query(Worker).filter(Worker.user_id == current_user.id).first()
    if not worker:
        # Generar token unico de worker y crearlo
        token = f"melodix_wk_{uuid.uuid4().hex}"
        worker = Worker(
            hostname="unregistered",
            name=f"Worker de {current_user.username}",
            status="offline",
            last_seen=datetime.utcnow(),
            user_id=current_user.id,
            token=token,
            command="none",
            is_active=False
        )
        db.add(worker)
        db.commit()
        db.refresh(worker)
    
    return {
        "id": worker.id,
        "name": worker.name,
        "status": worker.status,
        "last_seen": worker.last_seen,
        "token": worker.token,
        "is_active": worker.is_active,
        "command": worker.command,
        "hostname": worker.hostname
    }

@router.post("/my-worker/toggle")
def toggle_my_worker(db: Session = Depends(get_db), current_user: User = Depends(get_current_user)):
    worker = db.query(Worker).filter(Worker.user_id == current_user.id).first()
    if not worker:
        raise HTTPException(status_code=404, detail="Worker no encontrado")
    worker.is_active = not worker.is_active
    db.commit()
    db.refresh(worker)
    return {"is_active": worker.is_active, "status": worker.status}

@router.post("/my-worker/command")
def send_worker_command(command: str, db: Session = Depends(get_db), current_user: User = Depends(get_current_user)):
    if command not in ["start", "stop"]:
        raise HTTPException(status_code=400, detail="Comando invalido. Use 'start' o 'stop'")
    worker = db.query(Worker).filter(Worker.user_id == current_user.id).first()
    if not worker:
        raise HTTPException(status_code=404, detail="Worker no encontrado")
    worker.command = command
    db.commit()
    return {"message": f"Comando '{command}' enviado al worker"}

@router.get("/agent/poll")
def agent_poll(token: str, status: str, hostname: str, db: Session = Depends(get_db)):
    worker = db.query(Worker).filter(Worker.token == token).first()
    if not worker:
        raise HTTPException(status_code=404, detail="Token de Worker no registrado")
    
    worker.last_seen = datetime.utcnow()
    
    # Si el administrador suspendió el worker, mantener el estado "suspended" y detener el engine
    if worker.status == "suspended":
        current_command = "stop" if status == "busy" else worker.command
        if worker.command != "none":
            worker.command = "none"
        db.commit()
        return {
            "command": current_command,
            "is_active": False,
            "user_id": worker.user_id,
            "status": "suspended"
        }
    
    # De lo contrario, actualizar estado según el reporte del agente
    previous_status = worker.status
    if status == "offline":
        worker.status = "offline"
    elif status == "online":
        if worker.status != "busy":
            worker.status = "online"
    else:
        worker.status = status
    if hostname:
        worker.hostname = hostname
    
    current_command = worker.command
    if worker.command != "none":
        worker.command = "none"
    
    db.commit()
    
    if previous_status == "offline" and worker.status in ["online", "busy"]:
        create_admin_notification(
            db,
            "Worker Conectado",
            f"El worker local/BYOK '{worker.name or worker.hostname}' se ha conectado."
        )
        
    return {
        "command": current_command,
        "is_active": worker.is_active,
        "user_id": worker.user_id,
        "status": worker.status
    }

@router.get("/agent/poll-public")
def agent_poll_public(hostname: str, status: str, db: Session = Depends(get_db)):
    worker = db.query(Worker).filter(Worker.hostname == hostname, Worker.user_id == None).first()
    is_new = worker is None
    if not worker:
        # Registrar de forma automática si no existe en la base de datos
        worker = Worker(
            hostname=hostname,
            name=hostname,
            status="online",
            last_seen=datetime.utcnow(),
            command="none",
            is_active=True
        )
        db.add(worker)
        db.commit()
        db.refresh(worker)
        
        # Enviar notificación de nuevo worker
        create_admin_notification(
            db,
            "Worker Registrado",
            f"El equipo worker '{worker.hostname}' se ha autoregistrado vía poll público."
        )
    
    worker.last_seen = datetime.utcnow()
    
    # Si el administrador suspendió el worker, mantener el estado "suspended"
    if worker.status == "suspended":
        current_command = "stop" if status == "busy" else worker.command
        if worker.command != "none":
            worker.command = "none"
        db.commit()
        return {
            "command": current_command,
            "status": "suspended"
        }
    
    # De lo contrario, actualizar estado según el reporte del agente
    previous_status = worker.status
    if status == "offline":
        worker.status = "offline"
    elif status == "online":
        if worker.status != "busy":
            worker.status = "online"
    else:
        worker.status = status
    current_command = worker.command
    if worker.command != "none":
        worker.command = "none"
    
    db.commit()
    
    if not is_new and previous_status == "offline" and worker.status in ["online", "busy"]:
        create_admin_notification(
            db,
            "Worker Conectado",
            f"El equipo worker '{worker.name or worker.hostname}' se ha conectado."
        )
        
    return {
        "command": current_command,
        "status": worker.status
    }