"""
Definición de protocolo WebSocket y serialización de mensajes para comunicación con el Manager.
100% idéntico a internal/protocol/messages.go.
"""

import json
import base64
from typing import Any, Tuple, Dict, List, Optional
from dataclasses import dataclass, asdict

from config import logger

# Tipos de mensajes
TYPE_HANDSHAKE = "handshake"
TYPE_HANDSHAKE_ACK = "handshake_ack"
TYPE_PING = "ping"
TYPE_PONG = "pong"
TYPE_TASK_ASSIGN = "task_assign"
TYPE_TASK_PROGRESS = "task_progress"
TYPE_TASK_COMPLETE = "task_complete"
TYPE_TASK_FAILED = "task_failed"
TYPE_TASK_CANCEL = "task_cancel"


@dataclass
class VideoSettings:
    resolution: str = ""
    crf: str = ""
    audio_bitrate: str = ""
    fps: str = ""
    preset: str = ""
    codec: str = ""
    watermark: str = ""
    watermark_pos: str = ""
    watermark_image_file_id: str = ""
    format: str = ""
    thumbnail_file_id: str = ""


@dataclass
class TaskState:
    task_id: str
    status: str  # "converting", "completed"
    upload_url: str = ""


@dataclass
class Handshake:
    worker_id: str
    secret: str
    cpu_cores: int
    max_jobs: int
    active_tasks: List[Dict[str, Any]]


@dataclass
class TaskProgress:
    task_id: str
    percent: float
    speed: str
    time: str
    eta: str


@dataclass
class TaskComplete:
    task_id: str
    original_size: int
    compressed_size: int
    duration_sec: float


@dataclass
class TaskFailed:
    task_id: str
    error: str


@dataclass
class TaskCancel:
    task_id: str


def build_ws_packet(msg_type: str, payload_obj: Any = None) -> str:
    """
    Empaqueta el mensaje en formato WebSocket esperado por el Manager en Go:
    type WSMessage struct {
        Type    MessageType `json:"type"`
        Payload []byte      `json:"payload,omitempty"`
    }
    En Go, un campo `[]byte` se codifica/decodifica como una cadena Base64.
    """
    if payload_obj is not None:
        if hasattr(payload_obj, "__dataclass_fields__"):
            payload_data = asdict(payload_obj)
        elif isinstance(payload_obj, dict):
            payload_data = payload_obj
        else:
            payload_data = payload_obj

        raw_json = json.dumps(payload_data).encode("utf-8")
        b64_payload = base64.b64encode(raw_json).decode("ascii")
        return json.dumps({"type": msg_type, "payload": b64_payload})
    return json.dumps({"type": msg_type})


def parse_ws_packet(raw_text: str) -> Tuple[str, Any]:
    """
    Decodifica el mensaje JSON entrante soportando tanto payload en Base64
    como payload embebido directamente como diccionario/objeto.
    """
    try:
        data = json.loads(raw_text)
        msg_type = data.get("type", "")
        payload = data.get("payload")
        if isinstance(payload, str):
            try:
                decoded_bytes = base64.b64decode(payload)
                payload = json.loads(decoded_bytes.decode("utf-8"))
            except Exception:
                pass
        return msg_type, payload
    except Exception as e:
        logger.error(f"Error decodificando paquete WS: {e}")
        return "", None