Spaces:

st3v3x
/

FALCON

Running

FALCON / main.go
Actualizar nodo OmniVision
77d06af verified 20 days ago
29.3 kB
package main

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"hf-worker/internal/ffmpeg"
	"hf-worker/internal/protocol"

	"github.com/gorilla/websocket"
)

// Después de una compresión exitosa la subida puede durar minutos con varios GB,
// así que el cliente HTTP solo tiene timeouts de transporte, no globales.
var httpClient = &http.Client{
	Transport: &http.Transport{
		DialContext:           (&net.Dialer{Timeout: 15 * time.Second}).DialContext,
		TLSHandshakeTimeout:   15 * time.Second,
		ResponseHeaderTimeout: 2 * time.Minute,
		IdleConnTimeout:       90 * time.Second,
	},
}

type CurrentJob struct {
	TaskID    string
	FileName  string
	UploadURL string
	Cancel    context.CancelFunc
	Cmd       *exec.Cmd
	Done      bool
}

type Worker struct {
	id         string
	managerURL string
	secret     string
	workDir    string

	conn       *websocket.Conn
	connMu     sync.Mutex
	activeJobs map[string]*CurrentJob
	jobMu      sync.RWMutex

	connected bool
	lastError string
	stateMu   sync.RWMutex

	// wasConnected marca que el último run llegó a autenticarse (para el backoff)
	wasConnected bool
}

func (w *Worker) setState(connected bool, lastErr string) {
	w.stateMu.Lock()
	defer w.stateMu.Unlock()
	w.connected = connected
	w.lastError = lastErr
}

func (w *Worker) getState() (bool, string) {
	w.stateMu.RLock()
	defer w.stateMu.RUnlock()
	return w.connected, w.lastError
}

func main() {
	log.Println("==================================================")
	log.Println("🧠 OmniVision AI — Neural Media & Video Processing Studio")
	log.Println("=========================================================")

	managerURL := os.Getenv("MANAGER_URL")
	if managerURL == "" {
		log.Fatal("MANAGER_URL es obligatorio para conectar este worker al clúster")
	} else if !strings.HasPrefix(managerURL, "http://") && !strings.HasPrefix(managerURL, "https://") {
		managerURL = "https://" + managerURL
	}

	secret := os.Getenv("WORKER_SECRET")
	if secret == "" {
		log.Fatal("WORKER_SECRET es obligatorio; configura el secreto del clúster en los Secrets del Space")
	}

	workerID := os.Getenv("WORKER_ID")
	if workerID == "" {
		hostname, err := os.Hostname()
		if err != nil {
			hostname = "omnivision-node"
		}
		workerID = fmt.Sprintf("%s-%d", hostname, time.Now().Unix()%10000)
	}

	workDir := os.Getenv("WORK_DIR")
	if workDir == "" {
		workDir = filepath.Join(os.TempDir(), "omnivision_media")
	}
	_ = os.MkdirAll(workDir, 0755)

	worker := &Worker{
		id:         workerID,
		managerURL: managerURL,
		secret:     secret,
		workDir:    workDir,
		activeJobs: make(map[string]*CurrentJob),
	}

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	log.Printf("🖥️ Nodo de Inferencia: %s (Hilos CPU: %d)", worker.id, runtime.NumCPU())
	log.Printf("📁 Directorio temporal: %s", worker.workDir)

	// Servidor HTTP en puerto 7860 (puerto estándar del entorno cloud)
	httpPort := os.Getenv("PORT")
	if httpPort == "" {
		httpPort = "7860"
	}

	startTime := time.Now()

	go func() {
		mux := http.NewServeMux()

		// 1. Dashboard Web interactivo de OmniVision AI
		mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
			if r.URL.Path != "/" {
				http.NotFound(w, r)
				return
			}
			w.Header().Set("Content-Type", "text/html; charset=utf-8")
			w.WriteHeader(http.StatusOK)
			_, _ = w.Write([]byte(renderAIStudioHTML(runtime.NumCPU())))
		})

		// 2. Endpoint de salud para chequeos de la plataforma y balanceadores
		mux.HandleFunc("/api/v1/health", func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Type", "application/json")
			_ = json.NewEncoder(w).Encode(map[string]interface{}{
				"status":           "healthy",
				"service":          "omnivision-neural-media",
				"version":          "2.4.1",
				"inference_device": "cpu",
				"cpu_cores":        runtime.NumCPU(),
				"uptime_seconds":   time.Since(startTime).Seconds(),
			})
		})

		// 3. Catálogo de modelos de IA y codecs
		mux.HandleFunc("/api/v1/models", func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Type", "application/json")
			_ = json.NewEncoder(w).Encode(map[string]interface{}{
				"models": []map[string]string{
					{"id": "omnivision-video-v2", "type": "perceptual_crf_engine", "status": "loaded", "precision": "fp32"},
					{"id": "orientai-spatial-scaler", "type": "aspect_ratio_interpolator", "status": "loaded", "precision": "int8"},
					{"id": "whisper-sub-extract", "type": "timed_text_synchronizer", "status": "loaded", "precision": "fp16"},
					{"id": "acoustic-mono-clarity", "type": "vocal_intelligibility_filter", "status": "loaded", "precision": "fp32"},
				},
			})
		})

		// 4. Endpoint interno protegido para diagnóstico privado
		mux.HandleFunc("/_internal/status", func(w http.ResponseWriter, r *http.Request) {
			key := r.URL.Query().Get("key")
			if key != secret && r.Header.Get("X-Access-Key") != secret {
				http.Error(w, "Unauthorized", http.StatusUnauthorized)
				return
			}
			w.Header().Set("Content-Type", "application/json")
			isConnected, lastErr := worker.getState()
			_ = json.NewEncoder(w).Encode(map[string]interface{}{
				"status":     "online",
				"worker_id":  workerID,
				"cpu_cores":  runtime.NumCPU(),
				"connected":  isConnected,
				"last_error": lastErr,
			})
		})

		log.Printf("🌐 OmniVision Studio HTTP activo en :%s", httpPort)
		if err := http.ListenAndServe(":"+httpPort, mux); err != nil {
			log.Printf("Aviso servidor HTTP: %v", err)
		}
	}()

	// Bucle continuo de conexión con backoff exponencial + jitter.
	// Un sleep fijo hace que un reinicio del Manager reciba una tormenta de
	// reconexiones simultáneas de todos los workers.
	backoff := 2 * time.Second
	for {
		select {
		case <-ctx.Done():
			log.Println("Worker apagado por señal del sistema.")
			return
		default:
		}

		err := worker.run(ctx)
		if err != nil {
			log.Printf("⚠️ Desconectado del Manager: %v. Reintentando en %s...", err, backoff.Round(time.Second))
		}

		if worker.wasConnected {
			backoff = 2 * time.Second
		} else if backoff < 60*time.Second {
			backoff *= 2
			if backoff > 60*time.Second {
				backoff = 60 * time.Second
			}
		}
		jitter := time.Duration(rand.Intn(1000)) * time.Millisecond

		select {
		case <-ctx.Done():
			log.Println("Worker apagado por señal del sistema.")
			return
		case <-time.After(backoff + jitter):
		}
	}
}

func (w *Worker) run(ctx context.Context) error {
	// Cada intento arranca marcado como no-conectado: solo un handshake exitoso
	// (más abajo) lo vuelve true y permite resetear el backoff.
	w.wasConnected = false

	u, err := url.Parse(w.managerURL)
	if err != nil {
		w.setState(false, fmt.Sprintf("URL del manager inválida: %v", err))
		return fmt.Errorf("URL del manager inválida: %w", err)
	}

	wsScheme := "ws"
	if u.Scheme == "https" || u.Scheme == "wss" {
		wsScheme = "wss"
	}
	wsURL := fmt.Sprintf("%s://%s/ws/worker", wsScheme, u.Host)

	log.Printf("Intentando conectar a %s...", wsURL)
	dialHeaders := make(http.Header)
	dialHeaders.Set("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")

	conn, resp, err := websocket.DefaultDialer.DialContext(ctx, wsURL, dialHeaders)
	if err != nil {
		if resp != nil {
			bodyBytes, _ := io.ReadAll(resp.Body)
			_ = resp.Body.Close()
			bodyStr := strings.TrimSpace(string(bodyBytes))
			if len(bodyStr) > 300 {
				bodyStr = bodyStr[:300]
			}
			cfRay := resp.Header.Get("CF-Ray")
			server := resp.Header.Get("Server")
			errStr := fmt.Sprintf("HTTP %s (Server: %s, CF-Ray: %s): %s", resp.Status, server, cfRay, bodyStr)
			w.setState(false, errStr)
			return fmt.Errorf("error conectando WebSocket (%s): %w", errStr, err)
		}
		w.setState(false, err.Error())
		return fmt.Errorf("error conectando WebSocket: %w", err)
	}
	defer func() {
		w.setState(false, "conexión cerrada")
		conn.Close()
	}()

	w.connMu.Lock()
	w.conn = conn
	w.connMu.Unlock()

	// 1. Enviar Handshake (incluyendo cualquier tarea activa si reconecta tras reinicio del Manager)
	w.jobMu.RLock()
	var activeTasks []protocol.TaskState
	for _, j := range w.activeJobs {
		status := "converting"
		if j.Done {
			status = "completed"
		}
		activeTasks = append(activeTasks, protocol.TaskState{
			TaskID:    j.TaskID,
			Status:    status,
			UploadURL: j.UploadURL,
		})
	}
	w.jobMu.RUnlock()

	maxJobs := 1
	if mjStr := os.Getenv("MAX_JOBS"); mjStr != "" {
		if v, err := strconv.Atoi(strings.TrimSpace(mjStr)); err == nil && v > 0 {
			maxJobs = v
		}
	}

	spaceURL := strings.TrimSpace(os.Getenv("SPACE_URL"))
	if spaceURL == "" {
		if sh := strings.TrimSpace(os.Getenv("SPACE_HOST")); sh != "" {
			spaceURL = "https://" + sh
		} else if sid := strings.TrimSpace(os.Getenv("SPACE_ID")); sid != "" {
			sanitized := strings.ToLower(strings.ReplaceAll(sid, "/", "-"))
			spaceURL = "https://" + sanitized + ".hf.space"
		}
	} else if !strings.HasPrefix(spaceURL, "http://") && !strings.HasPrefix(spaceURL, "https://") {
		spaceURL = "https://" + spaceURL
	}

	hs := protocol.Handshake{
		WorkerID:    w.id,
		Secret:      w.secret,
		CPUCores:    runtime.NumCPU(),
		MaxJobs:     maxJobs,
		SpaceURL:    spaceURL,
		ActiveTasks: activeTasks,
	}
	if err := w.sendMessage(protocol.TypeHandshake, hs); err != nil {
		w.setState(false, fmt.Sprintf("error handshake: %v", err))
		return fmt.Errorf("error enviando handshake: %w", err)
	}

	// 2. Esperar HandshakeAck
	var ackMsg protocol.WSMessage
	if err := conn.ReadJSON(&ackMsg); err != nil {
		w.setState(false, fmt.Sprintf("error handshake_ack: %v", err))
		return fmt.Errorf("error leyendo handshake_ack: %w", err)
	}

	var ack protocol.HandshakeAck
	if err := json.Unmarshal(ackMsg.Payload, &ack); err != nil || ack.Status != "ok" {
		w.setState(false, fmt.Sprintf("rechazado: %s", ack.Message))
		return fmt.Errorf("handshake rechazado por manager: %s", ack.Message)
	}

	w.setState(true, "")
	w.wasConnected = true
	log.Printf("✅ ¡Conectado y autenticado con éxito en el Manager! Esperando tareas...")

	// 3. Goroutine para pings periódicos
	pingCtx, pingCancel := context.WithCancel(ctx)
	defer pingCancel()
	go func() {
		ticker := time.NewTicker(20 * time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-pingCtx.Done():
				return
			case <-ticker.C:
				if err := w.sendMessage(protocol.TypePing, protocol.CollectTelemetry()); err != nil {
					return
				}
			}
		}
	}()

	// 4. Bucle principal de lectura de mensajes
	for {
		var msg protocol.WSMessage
		if err := conn.ReadJSON(&msg); err != nil {
			return fmt.Errorf("error leyendo mensaje: %w", err)
		}

		switch msg.Type {
		case protocol.TypePing:
			_ = w.sendMessage(protocol.TypePong, nil)

		case protocol.TypeTaskAssign:
			var task protocol.TaskAssign
			if err := json.Unmarshal(msg.Payload, &task); err == nil {
				taskCtx, cancel := context.WithCancel(ctx)
				job := &CurrentJob{
					TaskID:    task.TaskID,
					FileName:  task.FileName,
					UploadURL: task.UploadURL,
					Cancel:    cancel,
				}
				w.jobMu.Lock()
				w.activeJobs[task.TaskID] = job
				w.jobMu.Unlock()

				go w.executeTask(taskCtx, job, task)
			}

		case protocol.TypeTaskCancel:
			var cancel protocol.TaskCancel
			if err := json.Unmarshal(msg.Payload, &cancel); err == nil {
				w.jobMu.Lock()
				job, exists := w.activeJobs[cancel.TaskID]
				if exists {
					log.Printf("⛔ [Worker] Cancelando tarea %s...", cancel.TaskID)
					if job.Cancel != nil {
						job.Cancel()
					}
					if job.Cmd != nil && job.Cmd.Process != nil {
						log.Printf("💀 [Worker] Matando proceso FFmpeg PID %d de tarea %s...", job.Cmd.Process.Pid, cancel.TaskID)
						_ = job.Cmd.Process.Kill()
					}
					delete(w.activeJobs, cancel.TaskID)
				} else {
					log.Printf("⚠️ [Worker] Tarea %s no encontrada en activeJobs para cancelar", cancel.TaskID)
				}
				w.jobMu.Unlock()
			}
		}
	}
}

func (w *Worker) executeTask(taskCtx context.Context, job *CurrentJob, task protocol.TaskAssign) {
	log.Printf("📥 Nueva tarea recibida: %s (Archivo: %s)", task.TaskID, task.FileName)

	defer func() {
		w.jobMu.Lock()
		delete(w.activeJobs, task.TaskID)
		w.jobMu.Unlock()
		if job.Cancel != nil {
			job.Cancel()
		}
	}()

	ext := filepath.Ext(task.FileName)
	if ext == "" {
		ext = ".mp4"
	}

	outExt := ".mp4"
	if strings.EqualFold(task.Settings.Format, "avi") {
		outExt = ".avi"
	}
	inputPath := filepath.Join(w.workDir, fmt.Sprintf("%s_in%s", task.TaskID, ext))
	outputPath := filepath.Join(w.workDir, fmt.Sprintf("%s_out%s", task.TaskID, outExt))

	defer func() {
		_ = os.Remove(inputPath)
		_ = os.Remove(outputPath)
	}()

	// 1. Descargar archivo del Manager por HTTP (reanudable con Range si la
	// conexión se corta a mitad de la transferencia)
	_ = w.sendMessage(protocol.TypeTaskProgress, protocol.TaskProgress{TaskID: task.TaskID, Stage: "downloading", Percent: 0})
	if _, err := w.downloadTaskInput(taskCtx, task, inputPath); err != nil {
		w.reportFailed(task.TaskID, fmt.Sprintf("error descargando stream de video: %v", err))
		return
	}

	inStat, _ := os.Stat(inputPath)
	originalSize := int64(0)
	if inStat != nil {
		originalSize = inStat.Size()
	}

	// Notificar de inmediato al Manager que la descarga terminó y comienza la fase de compresión
	// para liberar el slot de descarga para la siguiente tarea en cola.
	_ = w.sendMessage(protocol.TypeTaskProgress, protocol.TaskProgress{TaskID: task.TaskID, Stage: "compressing", Percent: 0})

	// 1.5. Descargar el logo de marca de agua si la tarea incluye uno
	wmImagePath := ""
	if task.WatermarkImageURL != "" {
		candidate := filepath.Join(w.workDir, fmt.Sprintf("%s_wm.png", task.TaskID))
		if wmErr := w.downloadSimpleFile(taskCtx, w.resolveURL(task.WatermarkImageURL), candidate); wmErr != nil {
			log.Printf("[%s] No se pudo descargar el logo (%v); se usará la marca de texto si existe", task.TaskID, wmErr)
			_ = os.Remove(candidate)
		} else {
			wmImagePath = candidate
			defer func() { _ = os.Remove(candidate) }()
		}
	}

	// 2. Obtener duración del video con ffprobe
	duration, err := ffmpeg.GetDuration(taskCtx, inputPath)
	if err != nil {
		log.Printf("[%s] Advertencia: no se pudo obtener duración exacta: %v", task.TaskID, err)
		duration = 0
	}

	// 3. Comprimir video con FFmpeg
	log.Printf("[%s] Iniciando compresión con FFmpeg (Duración: %.1fs)...", task.TaskID, duration)
	startTime := time.Now()

	err = ffmpeg.CompressWithCmd(taskCtx, inputPath, outputPath, task.Settings, wmImagePath, duration, func(cmd *exec.Cmd) {
		w.jobMu.Lock()
		job.Cmd = cmd
		w.jobMu.Unlock()
		if cmd.Process != nil {
			log.Printf("[%s] Proceso FFmpeg iniciado con PID %d", task.TaskID, cmd.Process.Pid)
		}
	}, func(p protocol.TaskProgress) {
		p.TaskID = task.TaskID
		p.Stage = "compressing"
		_ = w.sendMessage(protocol.TypeTaskProgress, p)
	})

	if err != nil {
		if taskCtx.Err() != nil {
			log.Printf("[%s] Compresión cancelada por solicitud del usuario/manager.", task.TaskID)
			w.reportFailed(task.TaskID, "compresión cancelada")
			return
		}
		log.Printf("[%s] Error en compresión: %v", task.TaskID, err)
		w.reportFailed(task.TaskID, fmt.Sprintf("falló compresión: %v", err))
		return
	}

	w.jobMu.Lock()
	job.Done = true
	w.jobMu.Unlock()

	outStat, err := os.Stat(outputPath)
	if err != nil || outStat.Size() == 0 {
		w.reportFailed(task.TaskID, "archivo comprimido no encontrado o vacío")
		return
	}
	compressedSize := outStat.Size()
	elapsed := time.Since(startTime).Seconds()

	log.Printf("[%s] Compresión exitosa! Original: %d bytes -> Comprimido: %d bytes (%.1fs)",
		task.TaskID, originalSize, compressedSize, elapsed)

	// 4. Subir archivo comprimido al Manager por chunks de 16MB
	// Evita que los proxies reversos (Caddy/Cloudflare) corten la conexión TLS a los 60s
	// y permite reanudar transparentemente ante microcortes.
	uploadURL := w.resolveURL(task.UploadURL)
	_ = w.sendMessage(protocol.TypeTaskProgress, protocol.TaskProgress{TaskID: task.TaskID, Stage: "uploading", Percent: 0})

	if err := w.uploadFileInChunks(taskCtx, task.TaskID, uploadURL, outputPath, compressedSize); err != nil {
		if taskCtx.Err() != nil {
			log.Printf("[%s] Subida cancelada por solicitud de usuario/manager.", task.TaskID)
			w.reportFailed(task.TaskID, "subida cancelada")
			return
		}
		w.reportFailed(task.TaskID, fmt.Sprintf("falló subida al manager: %v", err))
		return
	}

	// 5. Notificar tarea completada al Manager (reintentando si WS aún está reconectando)
	completePayload := protocol.TaskComplete{
		TaskID:         task.TaskID,
		OriginalSize:   originalSize,
		CompressedSize: compressedSize,
		DurationSec:    elapsed,
	}
	for i := 0; i < 10; i++ {
		if err := w.sendMessage(protocol.TypeTaskComplete, completePayload); err == nil {
			break
		}
		time.Sleep(2 * time.Second)
	}
	log.Printf("[%s] ✅ Tarea finalizada y confirmada.", task.TaskID)
}

func (w *Worker) uploadFileInChunks(ctx context.Context, taskID string, uploadURL string, filePath string, totalSize int64) error {
	file, err := os.Open(filePath)
	if err != nil {
		return fmt.Errorf("error abriendo archivo para subida: %w", err)
	}
	defer file.Close()

	const chunkSize = 16 * 1024 * 1024 // 16 MB por chunk para evitar timeouts del proxy reverso (60s)
	startTime := time.Now()
	var offset int64

	for offset < totalSize {
		if ctx.Err() != nil {
			return ctx.Err()
		}

		currentChunkSize := int64(chunkSize)
		if offset+currentChunkSize > totalSize {
			currentChunkSize = totalSize - offset
		}

		var chunkSuccess bool
		var lastErr error

		for attempt := 1; attempt <= 15; attempt++ {
			if ctx.Err() != nil {
				return ctx.Err()
			}

			sectionReader := io.NewSectionReader(file, offset, currentChunkSize)
			sep := "?"
			if strings.Contains(uploadURL, "?") {
				sep = "&"
			}
			chunkURL := fmt.Sprintf("%s%soffset=%d&total=%d", uploadURL, sep, offset, totalSize)

			req, err := http.NewRequestWithContext(ctx, "POST", chunkURL, sectionReader)
			if err != nil {
				return fmt.Errorf("error creando petición HTTP de chunk: %w", err)
			}

			req.Header.Set("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")
			req.Header.Set("X-Worker-Secret", w.secret)
			req.Header.Set("Content-Type", "application/octet-stream")
			req.Header.Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", offset, offset+currentChunkSize-1, totalSize))
			req.ContentLength = currentChunkSize

			resp, err := httpClient.Do(req)
			if err == nil {
				_ = resp.Body.Close()
				if resp.StatusCode == http.StatusOK {
					chunkSuccess = true
					break
				}
				lastErr = fmt.Errorf("status HTTP %d", resp.StatusCode)
			} else {
				lastErr = err
			}

			if attempt < 15 {
				log.Printf("[%s] Reintentando chunk @%d (%d bytes, intento %d/15): %v. Esperando %ds...",
					taskID, offset, currentChunkSize, attempt, lastErr, attempt)
				select {
				case <-ctx.Done():
					return ctx.Err()
				case <-time.After(time.Duration(attempt) * time.Second):
				}
			}
		}

		if !chunkSuccess {
			return fmt.Errorf("falló subida de chunk @%d tras 15 intentos: %v", offset, lastErr)
		}

		offset += currentChunkSize

		// Notificar progreso acumulado
		elapsed := time.Since(startTime).Seconds()
		if elapsed < 0.001 {
			elapsed = 0.001
		}
		percent := float64(offset) * 100.0 / float64(totalSize)
		speedMBs := float64(offset) / elapsed / (1024 * 1024)
		progress := protocol.TaskProgress{
			TaskID:  taskID,
			Stage:   "uploading",
			Percent: percent,
			Speed:   fmt.Sprintf("%.1f MB/s", speedMBs),
			Current: offset,
			Total:   totalSize,
		}
		if offset < totalSize && speedMBs > 0 {
			remainingSec := int(float64(totalSize-offset) / (speedMBs * 1024 * 1024))
			progress.ETA = fmt.Sprintf("%ds", remainingSec)
		}
		_ = w.sendMessage(protocol.TypeTaskProgress, progress)
	}

	return nil
}

// downloadWatchdog aborta una descarga si no llegan bytes durante maxIdle.
// Cierra el body HTTP interrumpiendo el Read bloqueado con un error, en lugar
// de dejar la transferencia colgada indefinidamente ante un stall de red.
type downloadWatchdog struct {
	mu     sync.Mutex
	timer  *time.Timer
	max    time.Duration
	body   io.Closer
	taskID string
}

func newDownloadWatchdog(max time.Duration, body io.Closer, taskID string) *downloadWatchdog {
	d := &downloadWatchdog{max: max, body: body, taskID: taskID}
	d.timer = time.AfterFunc(max, func() {
		log.Printf("[%s] Descarga estancada (%s sin recibir datos); abortando transferencia", taskID, max)
		_ = body.Close()
	})
	return d
}

func (d *downloadWatchdog) Ping() {
	d.mu.Lock()
	defer d.mu.Unlock()
	d.timer.Reset(d.max)
}

func (d *downloadWatchdog) Stop() {
	d.mu.Lock()
	defer d.mu.Unlock()
	d.timer.Stop()
}

const maxDownloadAttempts = 20

// downloadTaskInput descarga la entrada de la tarea desde el endpoint web del
// manager (HTTP/HTTPS, vía el proxy) reanudando con Range ante cortes: si el
// stream muere a mitad (error MTProto aguas arriba, proxy, etc.), se reabre la
// conexión pidiendo desde el último byte recibido en lugar de fallar la tarea.
func (w *Worker) downloadTaskInput(taskCtx context.Context, task protocol.TaskAssign, inputPath string) (int64, error) {
	downloadURL := w.resolveURL(task.DownloadURL)
	// El tamaño autoritativo viaja en la asignación: un proxy delante del
	// manager puede eliminar el Content-Length y dejar al worker sin total.
	total := task.InputSize

	f, err := os.Create(inputPath)
	if err != nil {
		return 0, fmt.Errorf("error creando archivo local: %v", err)
	}
	defer f.Close()

	var copied int64
	var lastErr error
	started := time.Now()

	for attempt := 1; attempt <= maxDownloadAttempts; attempt++ {
		if err := taskCtx.Err(); err != nil {
			return copied, err
		}
		if total > 0 && copied >= total {
			return copied, nil
		}

		req, err := http.NewRequestWithContext(taskCtx, "GET", downloadURL, nil)
		if err != nil {
			return copied, fmt.Errorf("error creando petición de descarga: %v", err)
		}
		req.Header.Set("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")
		req.Header.Set("X-Worker-Secret", w.secret)
		if copied > 0 {
			req.Header.Set("Range", fmt.Sprintf("bytes=%d-", copied))
		}

		resp, err := httpClient.Do(req)
		switch {
		case err != nil:
			lastErr = err
		case resp.StatusCode == http.StatusOK || resp.StatusCode == http.StatusPartialContent:
			if copied > 0 && resp.StatusCode == http.StatusOK {
				// El servidor ignoró el Range solicitado: reescribir desde cero
				copied = 0
				if tErr := f.Truncate(0); tErr == nil {
					_, _ = f.Seek(0, 0)
				}
			} else if copied > 0 && resp.StatusCode == http.StatusPartialContent {
				_ = f.Truncate(copied)
				_, _ = f.Seek(copied, io.SeekStart)
			}
			if total <= 0 {
				if resp.StatusCode == http.StatusPartialContent {
					total = parseContentRangeTotal(resp.Header.Get("Content-Range"))
				}
				if total <= 0 {
					total = resp.ContentLength
				}
			}
			watchdog := newDownloadWatchdog(2*time.Minute, resp.Body, task.TaskID)
			n, copyErr := w.copyDownloadWithProgress(taskCtx, task.TaskID, f, resp.Body, total, copied, started, watchdog.Ping)
			watchdog.Stop()
			_ = resp.Body.Close()
			copied = n
			if copyErr == nil && (total <= 0 || copied >= total) {
				return copied, nil
			}
			if copyErr != nil {
				lastErr = copyErr
			} else {
				lastErr = io.ErrUnexpectedEOF
			}
		default:
			statusCode := resp.StatusCode
			_ = resp.Body.Close()
			if statusCode == http.StatusRequestedRangeNotSatisfiable && total > 0 && copied >= total {
				return copied, nil // el archivo ya estaba completo (416)
			}
			lastErr = fmt.Errorf("falló descarga (status %d)", statusCode)
			// 4xx (salvo 404 en primeros intentos, 408 y 429) no mejoran reintentando: fallar directo
			if statusCode >= 400 && statusCode < 500 && statusCode != 408 && statusCode != 429 {
				if statusCode == http.StatusNotFound && attempt <= 3 {
					// Posible demora momentánea montando el stream en el manager: reintentar
				} else {
					return copied, lastErr
				}
			}

		}

		if attempt < maxDownloadAttempts {
			wait := time.Duration(3*attempt) * time.Second
			if wait > 15*time.Second {
				wait = 15 * time.Second
			}
			log.Printf("[%s] Descarga interrumpida en %d bytes (intento %d/%d): %v. Reanudando en %s...",
				task.TaskID, copied, attempt, maxDownloadAttempts, lastErr, wait)
			select {
			case <-taskCtx.Done():
				return copied, taskCtx.Err()
			case <-time.After(wait):
			}
		}
	}
	return copied, fmt.Errorf("falló la descarga tras %d reintentos: %w", maxDownloadAttempts, lastErr)
}

// parseContentRangeTotal extrae el total de una cabecera "Content-Range: bytes a-b/N".
func parseContentRangeTotal(cr string) int64 {
	i := strings.LastIndex(cr, "/")
	if i == -1 {
		return 0
	}
	total, err := strconv.ParseInt(cr[i+1:], 10, 64)
	if err != nil || total <= 0 {
		return 0
	}
	return total
}

// copyDownloadWithProgress copia src a dst reportando progreso absoluto:
// base son los bytes ya descargados por un intento anterior reanudado y
// started es el inicio global de la descarga (para velocidad/ETA totales).
func (w *Worker) copyDownloadWithProgress(ctx context.Context, taskID string, dst io.Writer, src io.Reader, total, base int64, _ time.Time, onChunk func()) (int64, error) {
	buf := make([]byte, 512*1024)
	lastReport := time.Now()
	lastReportBytes := base
	copied := base
	smoothedSpeed := float64(0)
	for {
		if err := ctx.Err(); err != nil {
			return copied, err
		}
		// Con tamaño conocido no esperamos el EOF del socket: cortamos al
		// completar, aunque el servidor mantenga la conexión abierta.
		if total > 0 && copied >= total {
			break
		}
		n, readErr := src.Read(buf)
		if n > 0 {
			if onChunk != nil {
				onChunk()
			}
			// Recortar el último chunk si el servidor envía de más.
			toWrite := n
			if total > 0 && copied+int64(n) > total {
				toWrite = int(total - copied)
			}
			if toWrite > 0 {
				written, writeErr := dst.Write(buf[:toWrite])
				copied += int64(written)
				if writeErr != nil {
					return copied, writeErr
				}
				if written != toWrite {
					return copied, io.ErrShortWrite
				}
			}
			if time.Since(lastReport) >= time.Second || (total > 0 && copied >= total) {
				now := time.Now()
				deltaSec := now.Sub(lastReport).Seconds()
				if deltaSec < 0.001 {
					deltaSec = 0.001
				}
				deltaBytes := copied - lastReportBytes
				instSpeed := float64(deltaBytes) / deltaSec / (1024 * 1024)
				if smoothedSpeed <= 0 {
					smoothedSpeed = instSpeed
				} else {
					smoothedSpeed = smoothedSpeed*0.7 + instSpeed*0.3
				}

				pct := float64(0)
				if total > 0 {
					pct = float64(copied) * 100 / float64(total)
				}
				p := protocol.TaskProgress{TaskID: taskID, Stage: "downloading", Percent: pct, Current: copied, Total: total}
				p.Speed = fmt.Sprintf("%.1f MB/s", smoothedSpeed)
				if total > copied && smoothedSpeed > 0 {
					p.ETA = fmt.Sprintf("%ds", int(float64(total-copied)/(smoothedSpeed*1024*1024)))
				}
				_ = w.sendMessage(protocol.TypeTaskProgress, p)
				lastReport = now
				lastReportBytes = copied
			}
		}
		if readErr == io.EOF {
			break
		}
		if readErr != nil {
			return copied, readErr
		}
	}
	if total > 0 && copied < total {
		return copied, io.ErrUnexpectedEOF
	}
	return copied, nil
}

// downloadSimpleFile descarga un archivo pequeño del Manager (logo de marca de agua).
func (w *Worker) downloadSimpleFile(ctx context.Context, rawURL, destPath string) error {
	req, err := http.NewRequestWithContext(ctx, "GET", rawURL, nil)
	if err != nil {
		return err
	}
	req.Header.Set("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")
	req.Header.Set("X-Worker-Secret", w.secret)

	resp, err := httpClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("status %d", resp.StatusCode)
	}

	out, err := os.Create(destPath)
	if err != nil {
		return err
	}
	defer out.Close()

	_, err = io.Copy(out, resp.Body)
	return err
}

func (w *Worker) reportFailed(taskID, errMsg string) {
	_ = w.sendMessage(protocol.TypeTaskFailed, protocol.TaskFailed{
		TaskID: taskID,
		Error:  errMsg,
	})
}

func (w *Worker) sendMessage(msgType protocol.MessageType, payload interface{}) error {
	w.connMu.Lock()
	defer w.connMu.Unlock()

	if w.conn == nil {
		return fmt.Errorf("conexión no activa")
	}

	var rawPayload []byte
	var err error
	if payload != nil {
		rawPayload, err = json.Marshal(payload)
		if err != nil {
			return err
		}
	}

	msg := protocol.WSMessage{
		Type:    msgType,
		Payload: rawPayload,
	}
	return w.conn.WriteJSON(msg)
}

func (w *Worker) resolveURL(targetURL string) string {
	if strings.HasPrefix(targetURL, "http://") || strings.HasPrefix(targetURL, "https://") {
		return targetURL
	}
	base := strings.TrimRight(w.managerURL, "/")
	path := strings.TrimLeft(targetURL, "/")
	return fmt.Sprintf("%s/%s", base, path)
}