Blame for client.py

"""
Motor del Worker, conexión WebSocket y ejecución de tareas de compresión.
100% idéntico a la lógica y ciclo de vida de hf_worker/main.go.
"""

import os
import time
import random
import asyncio
import threading
import subprocess
from pathlib import Path
from typing import Optional, Dict, Any
from urllib.parse import urlparse

import psutil
import requests
import websockets

from config import (
    MANAGER_URL,
    WORKER_SECRET,
    WORKER_ID,
    WORK_DIR,
    MAX_JOBS,
    SPACE_URL,
    logger,
)
from state import STATE, CurrentJob
from protocol import (
    TYPE_HANDSHAKE,
    TYPE_HANDSHAKE_ACK,
    TYPE_PING,
    TYPE_PONG,
    TYPE_TASK_ASSIGN,
    TYPE_TASK_PROGRESS,
    TYPE_TASK_COMPLETE,
    TYPE_TASK_FAILED,
    TYPE_TASK_CANCEL,
    build_ws_packet,
    parse_ws_packet,
)
from ffmpeg import (
    get_duration,
    compress_with_cmd,
)


class UploadProgressReader:
    """Itera el archivo de salida y reporta velocidad/progreso al panel."""

    def __init__(self, file_obj, total_size: int, task_id: str, worker):
        self.file_obj = file_obj
        self.total_size = total_size
        self.task_id = task_id
        self.worker = worker

    def __len__(self):
        return self.total_size

    def __iter__(self):
        started = time.monotonic()
        last_report = started
        uploaded = 0
        while True:
            chunk = self.file_obj.read(1024 * 1024)
            if not chunk:
                break
            uploaded += len(chunk)
            now = time.monotonic()
            if now - last_report >= 1 or uploaded >= self.total_size:
                elapsed = max(now - started, 0.001)
                eta = ""
                if uploaded < self.total_size:
                    eta = f"{int((self.total_size - uploaded) / (uploaded / elapsed))}s"
                progress = {
                    "task_id": self.task_id,
                    "stage": "uploading",
                    "percent": min(100.0, uploaded * 100 / max(self.total_size, 1)),
                    "current": uploaded,
                    "total": self.total_size,
                    "speed": f"{uploaded / elapsed / (1024 * 1024):.1f} MB/s",
                    "eta": eta,
                }
                asyncio.run_coroutine_threadsafe(
                    self.worker._send_ws_safe(TYPE_TASK_PROGRESS, progress),
                    self.worker.loop,
                )
                last_report = now
            yield chunk


class WorkerEngine:
    """Cliente WebSocket y gestor de tareas del clúster."""

    def __init__(self):
        self.manager_url: str = MANAGER_URL
        self.secret: str = WORKER_SECRET
        self.worker_id: str = WORKER_ID
        self.work_dir: str = WORK_DIR

        self.ws_conn: Optional[websockets.WebSocketClientProtocol] = None
        self.loop: Optional[asyncio.AbstractEventLoop] = None
        self.session = requests.Session()
        self.session.headers.update({
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36",
        })

    def start(self):
        """Inicia el bucle de supervisión en un hilo secundario continuo."""
        thread = threading.Thread(target=self._run_event_loop, daemon=True, name="WorkerEngineThread")
        thread.start()

    def _run_event_loop(self):
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)
        self.loop.run_until_complete(self._supervise_connection())

    async def _supervise_connection(self):
        cores = psutil.cpu_count(logical=True) or 2
        STATE.log(f"🖥️ Nodo de Inferencia: {self.worker_id} (Hilos CPU: {cores})")
        if SPACE_URL:
            STATE.log(f"💓 URL Watchdog detectada: {SPACE_URL}")
        STATE.log(f"📁 Directorio temporal: {self.work_dir}")
        STATE.log(f"🌐 Upstream Manager: {self.manager_url}")

        # Bucle continuo de conexión con backoff exponencial + jitter (idéntico a main.go)
        backoff = 2.0
        while True:
            try:
                await self._run()
                if STATE.was_connected:
                    backoff = 2.0
                elif backoff < 60.0:
                    backoff = min(60.0, backoff * 2.0)
            except Exception as e:
                STATE.set_connection(False, str(e))
                STATE.log(f"⚠️ Desconectado del Manager: {e}. Reintentando en {int(round(backoff))}s...", level="WARNING")
                if STATE.was_connected:
                    backoff = 2.0
                elif backoff < 60.0:
                    backoff = min(60.0, backoff * 2.0)

            jitter = random.uniform(0.0, 1.0)
            await asyncio.sleep(backoff + jitter)

    def _resolve_url(self, target_url: str) -> str:
        """Resuelve URLs relativas contra la URL base del Manager."""
        if target_url.startswith("http://") or target_url.startswith("https://"):
            return target_url
        base = self.manager_url.rstrip("/")
        path = target_url.lstrip("/")
        return f"{base}/{path}"

    async def _send_ws_safe(self, msg_type: str, payload_obj: Any = None) -> bool:
        """Envía un paquete WS de forma segura."""
        if self.ws_conn and self.ws_conn.open:
            try:
                await self.ws_conn.send(build_ws_packet(msg_type, payload_obj))
                return True
            except Exception as e:
                logger.warning(f"Error enviando {msg_type}: {e}")
                return False
        return False

    async def _run(self):
        # Cada intento arranca marcado como no-conectado
        STATE.was_connected = False

        u = urlparse(self.manager_url)
        if not u.netloc and not u.path:
            raise ValueError(f"URL del manager inválida: {self.manager_url}")

        host = u.netloc or u.path
        ws_scheme = "wss" if u.scheme in ("https", "wss") else "ws"
        ws_url = f"{ws_scheme}://{host}/ws/worker"

        STATE.log(f"Intentando conectar a {ws_url}...")
        dial_headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36",
        }

        async with websockets.connect(
            ws_url,
            extra_headers=dial_headers,
            ping_interval=None,  # Manejado manualmente por protocolo
            max_size=10 * 1024 * 1024,
        ) as ws:
            self.ws_conn = ws

            # 1. Enviar Handshake (incluyendo cualquier tarea activa si reconecta tras reinicio del Manager)
            active_tasks = STATE.get_active_tasks_for_handshake()
            cores = psutil.cpu_count(logical=True) or 2
            handshake_payload = {
                "worker_id": self.worker_id,
                "secret": self.secret,
                "cpu_cores": cores,
                "max_jobs": MAX_JOBS,
                "space_url": SPACE_URL,
                "active_tasks": active_tasks,
            }
            await ws.send(build_ws_packet(TYPE_HANDSHAKE, handshake_payload))

            # 2. Esperar HandshakeAck
            ack_raw = await ws.recv()
            mtype, ack_payload = parse_ws_packet(ack_raw)
            if mtype != TYPE_HANDSHAKE_ACK or not isinstance(ack_payload, dict) or ack_payload.get("status") != "ok":
                err_msg = ack_payload.get("message", "rechazado por manager") if isinstance(ack_payload, dict) else "rechazado"
                STATE.set_connection(False, err_msg)
                raise ConnectionRefusedError(f"handshake rechazado por manager: {err_msg}")

            STATE.set_connection(True)
            STATE.log("✅ ¡Conectado y autenticado con éxito en el Manager! Esperando tareas...")

            # 3. Tarea asíncrona de Ping periódico (cada 20s)
            ping_stop = asyncio.Event()

            async def ping_loop():
                while not ping_stop.is_set():
                    try:
                        await asyncio.sleep(20)
                        memory = psutil.virtual_memory()
                        try:
                            disk = psutil.disk_usage('/')
                        except Exception:
                            disk = None
                        telemetry = {
                            "cpu_percent": psutil.cpu_percent(interval=None),
                            "memory_percent": memory.percent,
                            "memory_used_mb": int(memory.used / (1024 * 1024)),
                            "memory_total_mb": int(memory.total / (1024 * 1024)),
                            "disk_percent": disk.percent if disk else 0.0,
                            "disk_used_mb": int(disk.used / (1024 * 1024)) if disk else 0,
                            "disk_total_mb": int(disk.total / (1024 * 1024)) if disk else 0,
                        }
                        if not await self._send_ws_safe(TYPE_PING, telemetry):
                            break
                    except Exception:
                        break

            ping_task = asyncio.create_task(ping_loop())

            try:
                # 4. Bucle principal de lectura de mensajes
                async for raw_msg in ws:
                    msg_type, payload = parse_ws_packet(raw_msg)
                    if not msg_type:
                        continue

                    if msg_type == TYPE_PING:
                        await self._send_ws_safe(TYPE_PONG, None)

                    elif msg_type == TYPE_TASK_ASSIGN and isinstance(payload, dict):
                        task_id = payload.get("task_id", "")
                        file_name = payload.get("file_name", "")
                        upload_url = payload.get("upload_url", "")

                        job = CurrentJob(task_id, file_name, upload_url)
                        STATE.add_job(job)

                        # Ejecutar tarea en hilo independiente
                        asyncio.create_task(self._execute_task_wrapper(payload, job))

                    elif msg_type == TYPE_TASK_CANCEL and isinstance(payload, dict):
                        cancel_id = payload.get("task_id", "")
                        target_job = STATE.get_job(cancel_id)
                        if target_job:
                            STATE.log(f"⛔ [Worker] Cancelando tarea {cancel_id}...")
                            if target_job.cancel_func:
                                target_job.cancel_func()
                            if target_job.proc and target_job.proc.poll() is None:
                                try:
                                    STATE.log(f"💀 [Worker] Matando proceso FFmpeg PID {target_job.proc.pid} de tarea {cancel_id}...")
                                    target_job.proc.kill()
                                except Exception:
                                    pass
                            STATE.remove_job(cancel_id)
                        else:
                            STATE.log(f"⚠️ [Worker] Tarea {cancel_id} no encontrada en activeJobs para cancelar", level="WARNING")

            finally:
                ping_stop.set()
                ping_task.cancel()
                self.ws_conn = None

    async def _execute_task_wrapper(self, task_payload: Dict[str, Any], job: CurrentJob):
        """Envuelve la ejecución de la tarea en asyncio.to_thread."""
        await asyncio.to_thread(self._execute_task, task_payload, job)

    def _execute_task(self, task: Dict[str, Any], job: CurrentJob):
        task_id = task.get("task_id", "")
        file_name = task.get("file_name", "video.mp4")
        STATE.log(f"📥 Nueva tarea recibida: {task_id} (Archivo: {file_name})")

        cancel_event = threading.Event()
        job.cancel_func = cancel_event.set

        ext = Path(file_name).suffix or ".mp4"
        out_ext = ".avi" if task.get("settings", {}).get("format", "").lower() == "avi" else ".mp4"

        input_path = os.path.join(self.work_dir, f"{task_id}_in{ext}")
        output_path = os.path.join(self.work_dir, f"{task_id}_out{out_ext}")
        wm_image_path = ""

        try:
            # 1. Descargar archivo del Manager
            download_url = self._resolve_url(task.get("download_url", ""))
            req_headers = {
                "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36",
                "X-Worker-Secret": self.secret,
            }

            job.status = "Descargando..."
            # El tamaño autoritativo viaja en la asignación: un proxy delante
            # del manager puede eliminar el Content-Length (respuesta chunked).
            total_size = int(task.get("input_size") or 0)
            asyncio.run_coroutine_threadsafe(
                self._send_ws_safe(TYPE_TASK_PROGRESS, {
                    "task_id": task_id,
                    "stage": "downloading",
                    "percent": 0.0,
                    "current": 0,
                    "total": total_size,
                }),
                self.loop,
            )
            started = time.time()
            received = 0
            last_report = started
            last_report_bytes = 0
            smoothed_speed = 0.0
            max_attempts = 20

            # Descarga por HTTP reanudable con Range: si el stream del manager
            # muere a mitad de la transferencia, se reabre la conexión pidiendo
            # los bytes que faltan en lugar de fallar la tarea entera.
            with open(input_path, "wb") as f:
                attempt = 0
                while True:
                    if cancel_event.is_set():
                        self._report_failed(task_id, "descarga cancelada")
                        return
                    if total_size and received >= total_size:
                        break

                    attempt += 1
                    headers = dict(req_headers)
                    if received > 0:
                        headers["Range"] = f"bytes={received}-"

                    try:
                        with self.session.get(download_url, headers=headers, stream=True, timeout=120) as resp:
                            if resp.status_code == 416 and total_size and received >= total_size:
                                break  # archivo ya completo
                            if resp.status_code == 200 and received > 0:
                                # El servidor ignoró el Range: reescribir desde cero
                                f.seek(0)
                                f.truncate(0)
                                received = 0
                                last_report_bytes = 0
                            elif resp.status_code == 206 and received > 0:
                                f.seek(received)
                                f.truncate(received)
                            if total_size == 0:
                                if resp.status_code == 206:
                                    cr = resp.headers.get("Content-Range", "")
                                    if "/" in cr:
                                        try:
                                            total_size = int(cr.rsplit("/", 1)[1])
                                        except ValueError:
                                            pass
                                else:
                                    total_size = int(resp.headers.get("Content-Length", "0") or 0)
                            if resp.status_code not in (200, 206):
                                if resp.status_code in (408, 429) or resp.status_code >= 500 or (resp.status_code == 404 and attempt <= 3):
                                    raise IOError(f"falló descarga (status {resp.status_code})")
                                self._report_failed(task_id, f"falló descarga (status {resp.status_code})")
                                return

                            for chunk in resp.iter_content(chunk_size=512 * 1024):
                                if cancel_event.is_set():
                                    self._report_failed(task_id, "descarga cancelada")
                                    return
                                if chunk:
                                    # Con tamaño conocido no esperamos un EOF que puede no
                                    # llegar: recortamos el último chunk y cerramos.
                                    if total_size and received + len(chunk) > total_size:
                                        chunk = chunk[: total_size - received]
                                    if chunk:
                                        f.write(chunk)
                                        received += len(chunk)
                                    now = time.time()
                                    if now - last_report >= 1 or (total_size and received >= total_size):
                                        delta_sec = max(now - last_report, 0.001)
                                        delta_bytes = received - last_report_bytes
                                        inst_speed = delta_bytes / delta_sec / (1024 * 1024)
                                        if smoothed_speed <= 0:
                                            smoothed_speed = inst_speed
                                        else:
                                            smoothed_speed = smoothed_speed * 0.7 + inst_speed * 0.3
                                        pct = (received * 100 / total_size) if total_size else 0.0
                                        job.percent = pct
                                        job.speed = f"{smoothed_speed:.1f} MB/s"
                                        eta = ""
                                        if total_size and received < total_size and smoothed_speed > 0:
                                            eta = f"{int((total_size - received) / (smoothed_speed * 1024 * 1024))}s"
                                        progress = {
                                            "task_id": task_id,
                                            "stage": "downloading",
                                            "percent": pct,
                                            "current": received,
                                            "total": total_size,
                                            "speed": job.speed,
                                            "eta": eta,
                                        }
                                        asyncio.run_coroutine_threadsafe(
                                            self._send_ws_safe(TYPE_TASK_PROGRESS, progress), self.loop
                                        )
                                        last_report = now
                                        last_report_bytes = received
                                    if total_size and received >= total_size:
                                        break
                            if total_size and received < total_size:
                                # Conexión terminada sin excepción pero incompleta
                                raise IOError("conexión cerrada antes de completar la descarga")
                            f.flush()
                            break
                    except Exception as e:
                        try:
                            f.flush()
                        except Exception:
                            pass
                        if attempt >= max_attempts:
                            self._report_failed(task_id, f"falló descarga tras {max_attempts} reintentos ({received} bytes recibidos): {e}")
                            return
                        STATE.log(
                            f"[{task_id}] Descarga interrumpida en {received} bytes (intento {attempt}/{max_attempts}): {e}. Reanudando...",
                            level="WARNING",
                        )
                        time.sleep(min(3 * attempt, 15))

            original_size = os.path.getsize(input_path) if os.path.isfile(input_path) else 0

            # Notificar al Manager que la descarga finalizó para liberar el slot de descarga
            asyncio.run_coroutine_threadsafe(
                self._send_ws_safe(TYPE_TASK_PROGRESS, {"task_id": task_id, "stage": "compressing", "percent": 0.0}),
                self.loop,
            )

            # 1.5. Descargar el logo de marca de agua si la tarea incluye uno
            wm_url = task.get("watermark_image_url")
            if wm_url:
                candidate = os.path.join(self.work_dir, f"{task_id}_wm.png")
                try:
                    with self.session.get(self._resolve_url(wm_url), headers=req_headers, timeout=30) as wm_resp:
                        if wm_resp.status_code == 200:
                            with open(candidate, "wb") as f:
                                f.write(wm_resp.content)
                            wm_image_path = candidate
                        else:
                            STATE.log(f"[{task_id}] No se pudo descargar el logo (status {wm_resp.status_code}); se usará marca de texto")
                except Exception as ex:
                    STATE.log(f"[{task_id}] Error descargando logo ({ex}); se usará marca de texto")

            # 2. Obtener duración del video con ffprobe
            duration = get_duration(input_path)
            if duration <= 0:
                STATE.log(f"[{task_id}] Advertencia: no se pudo obtener duración exacta con ffprobe")

            # 3. Comprimir video con FFmpeg
            STATE.log(f"[{task_id}] Iniciando compresión con FFmpeg (Duración: {duration:.1f}s)...")
            start_t = time.time()
            job.status = "Comprimiendo..."

            def on_cmd_cb(proc: subprocess.Popen):
                job.proc = proc
                STATE.log(f"[{task_id}] Proceso FFmpeg iniciado con PID {proc.pid}")

            def on_progress_cb(p_data: Dict[str, Any]):
                job.percent = p_data.get("percent", 0.0)
                job.speed = p_data.get("speed", "--")
                job.eta = p_data.get("eta", "--")
                p_data["task_id"] = task_id
                p_data["stage"] = "compressing"
                asyncio.run_coroutine_threadsafe(
                    self._send_ws_safe(TYPE_TASK_PROGRESS, p_data),
                    self.loop,
                )

            try:
                compress_with_cmd(
                    input_path=input_path,
                    output_path=output_path,
                    settings=task.get("settings", {}),
                    wm_image_path=wm_image_path,
                    total_duration=duration,
                    on_cmd=on_cmd_cb,
                    on_progress=on_progress_cb,
                    is_cancelled=cancel_event.is_set,
                )
            except InterruptedError:
                STATE.log(f"[{task_id}] Compresión cancelada por solicitud del usuario/manager.")
                self._report_failed(task_id, "compresión cancelada")
                return
            except Exception as e:
                STATE.log(f"[{task_id}] Error en compresión: {e}", level="ERROR")
                self._report_failed(task_id, f"falló compresión: {e}")
                return

            job.done = True

            if not os.path.isfile(output_path) or os.path.getsize(output_path) == 0:
                self._report_failed(task_id, "archivo comprimido no encontrado o vacío")
                return

            compressed_size = os.path.getsize(output_path)
            elapsed = time.time() - start_t
            STATE.log(f"[{task_id}] Compresión exitosa! Original: {original_size} bytes -> Comprimido: {compressed_size} bytes ({elapsed:.1f}s)")

            # 4. Subir archivo comprimido al Manager en bloques (chunked resumable upload)
            # para evitar cortes por timeouts de reverse proxy (Caddy / Cloudflare / Nginx 60s)
            upload_base_url = self._resolve_url(task.get("upload_url", ""))
            job.status = "Subiendo..."
            upload_success = False

            chunk_size = 16 * 1024 * 1024  # 16 MB por bloque
            offset = 0
            started = time.monotonic()

            asyncio.run_coroutine_threadsafe(
                self._send_ws_safe(TYPE_TASK_PROGRESS, {
                    "task_id": task_id, "stage": "uploading", "percent": 0.0,
                    "current": 0, "total": compressed_size,
                }), self.loop
            )

            with open(output_path, "rb") as out_f:
                while offset < compressed_size:
                    if cancel_event.is_set():
                        self._report_failed(task_id, "subida cancelada")
                        return

                    out_f.seek(offset)
                    chunk_data = out_f.read(chunk_size)
                    if not chunk_data:
                        break

                    chunk_len = len(chunk_data)
                    end_offset = offset + chunk_len - 1

                    sep = "&" if "?" in upload_base_url else "?"
                    chunk_url = f"{upload_base_url}{sep}offset={offset}&total={compressed_size}"

                    chunk_headers = {
                        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36",
                        "Content-Type": "application/octet-stream",
                        "X-Worker-Secret": self.secret,
                        "Content-Length": str(chunk_len),
                        "Content-Range": f"bytes {offset}-{end_offset}/{compressed_size}",
                    }

                    chunk_uploaded = False
                    for chunk_attempt in range(1, 16):
                        if cancel_event.is_set():
                            self._report_failed(task_id, "subida cancelada")
                            return
                        try:
                            resp = self.session.post(chunk_url, data=chunk_data, headers=chunk_headers, timeout=60)
                            if resp.status_code in (200, 206):
                                chunk_uploaded = True
                                break
                            else:
                                STATE.log(f"[{task_id}] Manager respondió status {resp.status_code} en bloque @{offset}")
                                time.sleep(2)
                        except Exception as ex:
                            if chunk_attempt < 15:
                                STATE.log(f"[{task_id}] Reintentando bloque @{offset} ({chunk_attempt}/15): {ex}", level="WARNING")
                                time.sleep(3)

                    if not chunk_uploaded:
                        STATE.log(f"[{task_id}] Falló la subida del bloque @{offset} tras múltiples reintentos", level="ERROR")
                        break

                    offset += chunk_len

                    now = time.monotonic()
                    elapsed = max(now - started, 0.001)
                    percent = min(100.0, offset * 100.0 / max(compressed_size, 1))
                    speed_mb = (offset / elapsed) / (1024 * 1024)
                    eta_sec = int((compressed_size - offset) / (offset / elapsed)) if offset < compressed_size else 0
                    progress = {
                        "task_id": task_id,
                        "stage": "uploading",
                        "percent": percent,
                        "current": offset,
                        "total": compressed_size,
                        "speed": f"{speed_mb:.1f} MB/s",
                        "eta": f"{eta_sec}s" if eta_sec > 0 else "",
                    }
                    asyncio.run_coroutine_threadsafe(
                        self._send_ws_safe(TYPE_TASK_PROGRESS, progress),
                        self.loop,
                    )

                if offset >= compressed_size:
                    upload_success = True

            if not upload_success:
                self._report_failed(task_id, "falló subida al manager tras múltiples reintentos")
                return

            # 5. Notificar tarea completada al Manager (reintentando si WS aún está reconectando)
            complete_payload = {
                "task_id": task_id,
                "original_size": original_size,
                "compressed_size": compressed_size,
                "duration_sec": round(elapsed, 2),
            }
            for _ in range(10):
                fut = asyncio.run_coroutine_threadsafe(
                    self._send_ws_safe(TYPE_TASK_COMPLETE, complete_payload),
                    self.loop,
                )
                try:
                    if fut.result(timeout=3):
                        break
                except Exception:
                    pass
                time.sleep(2)

            STATE.record_success(original_size, compressed_size)
            STATE.log(f"[{task_id}] ✅ Tarea finalizada y confirmada.")

        except Exception as e:
            # Red/FS/errores inesperados (p.ej. timeout de lectura en la
            # descarga): notificar el fallo para que el manager no deje la
            # tarea colgada indefinidamente.
            STATE.log(f"[{task_id}] Error inesperado en la tarea: {e}", level="ERROR")
            self._report_failed(task_id, f"error inesperado: {e}")
        finally:
            STATE.remove_job(task_id)
            for p in (input_path, output_path, wm_image_path):
                if p and os.path.isfile(p):
                    try:
                        os.remove(p)
                    except Exception:
                        pass

    def _report_failed(self, task_id: str, err_msg: str):
        """Notifica fallo de tarea al Manager."""
        STATE.record_fail()
        asyncio.run_coroutine_threadsafe(
            self._send_ws_safe(TYPE_TASK_FAILED, {"task_id": task_id, "error": err_msg}),
            self.loop,
        )