"""WebSocket message Pydantic models for the Aniworld web application. This module defines message models for WebSocket communication between the server and clients. Models ensure type safety and provide validation for real-time updates. Series Identifier Convention: - `key`: Primary identifier for series (provider-assigned, URL-safe) e.g., "attack-on-titan" - `folder`: Display metadata only (e.g., "Attack on Titan (2013)") All series-related WebSocket events should include `key` as the primary identifier in their data payload. The `folder` field is optional and used for display purposes only. """ from __future__ import annotations from datetime import datetime, timezone from enum import Enum from typing import Any, Dict, Optional from pydantic import BaseModel, Field class WebSocketMessageType(str, Enum): """Types of WebSocket messages.""" # Download-related messages DOWNLOAD_PROGRESS = "download_progress" DOWNLOAD_COMPLETE = "download_complete" DOWNLOAD_FAILED = "download_failed" DOWNLOAD_ADDED = "download_added" DOWNLOAD_REMOVED = "download_removed" # Queue-related messages QUEUE_STATUS = "queue_status" QUEUE_STARTED = "queue_started" QUEUE_STOPPED = "queue_stopped" QUEUE_PAUSED = "queue_paused" QUEUE_RESUMED = "queue_resumed" # Progress-related messages SCAN_PROGRESS = "scan_progress" SCAN_COMPLETE = "scan_complete" SCAN_FAILED = "scan_failed" # System messages SYSTEM_INFO = "system_info" SYSTEM_WARNING = "system_warning" SYSTEM_ERROR = "system_error" # Error messages ERROR = "error" # Connection messages CONNECTED = "connected" PING = "ping" PONG = "pong" class WebSocketMessage(BaseModel): """Base WebSocket message structure.""" type: WebSocketMessageType = Field( ..., description="Type of the message" ) timestamp: str = Field( default_factory=lambda: datetime.now(timezone.utc).isoformat(), description="ISO 8601 timestamp when message was created", ) data: Dict[str, Any] = Field( default_factory=dict, description="Message payload" ) class DownloadProgressMessage(BaseModel): """Download progress update message. Data payload should include: - download_id: Unique download identifier - key: Series identifier (primary, e.g., 'attack-on-titan') - folder: Series folder name (optional, display only) - percent: Download progress percentage - speed_mbps: Download speed - eta_seconds: Estimated time remaining """ type: WebSocketMessageType = Field( default=WebSocketMessageType.DOWNLOAD_PROGRESS, description="Message type", ) timestamp: str = Field( default_factory=lambda: datetime.now(timezone.utc).isoformat(), description="ISO 8601 timestamp", ) data: Dict[str, Any] = Field( ..., description=( "Progress data including download_id, key (series identifier), " "folder (display), percent, speed_mbps, eta_seconds" ), ) class DownloadCompleteMessage(BaseModel): """Download completion message. Data payload should include: - download_id: Unique download identifier - key: Series identifier (primary, e.g., 'attack-on-titan') - folder: Series folder name (optional, display only) - file_path: Path to downloaded file """ type: WebSocketMessageType = Field( default=WebSocketMessageType.DOWNLOAD_COMPLETE, description="Message type", ) timestamp: str = Field( default_factory=lambda: datetime.now(timezone.utc).isoformat(), description="ISO 8601 timestamp", ) data: Dict[str, Any] = Field( ..., description=( "Completion data including download_id, key (series identifier), " "folder (display), file_path" ), ) class DownloadFailedMessage(BaseModel): """Download failure message. Data payload should include: - download_id: Unique download identifier - key: Series identifier (primary, e.g., 'attack-on-titan') - folder: Series folder name (optional, display only) - error_message: Description of the failure """ type: WebSocketMessageType = Field( default=WebSocketMessageType.DOWNLOAD_FAILED, description="Message type", ) timestamp: str = Field( default_factory=lambda: datetime.now(timezone.utc).isoformat(), description="ISO 8601 timestamp", ) data: Dict[str, Any] = Field( ..., description=( "Error data including download_id, key (series identifier), " "folder (display), error_message" ), ) class QueueStatusMessage(BaseModel): """Queue status update message.""" type: WebSocketMessageType = Field( default=WebSocketMessageType.QUEUE_STATUS, description="Message type", ) timestamp: str = Field( default_factory=lambda: datetime.now(timezone.utc).isoformat(), description="ISO 8601 timestamp", ) data: Dict[str, Any] = Field( ..., description="Queue status including active, pending, completed counts", ) class SystemMessage(BaseModel): """System-level message (info, warning, error).""" type: WebSocketMessageType = Field( ..., description="System message type" ) timestamp: str = Field( default_factory=lambda: datetime.now(timezone.utc).isoformat(), description="ISO 8601 timestamp", ) data: Dict[str, Any] = Field( ..., description="System message data" ) class ErrorMessage(BaseModel): """Error message to client.""" type: WebSocketMessageType = Field( default=WebSocketMessageType.ERROR, description="Message type" ) timestamp: str = Field( default_factory=lambda: datetime.now(timezone.utc).isoformat(), description="ISO 8601 timestamp", ) data: Dict[str, Any] = Field( ..., description="Error data including code and message" ) class ConnectionMessage(BaseModel): """Connection-related message (connected, ping, pong).""" type: WebSocketMessageType = Field( ..., description="Connection message type" ) timestamp: str = Field( default_factory=lambda: datetime.now(timezone.utc).isoformat(), description="ISO 8601 timestamp", ) data: Dict[str, Any] = Field( default_factory=dict, description="Connection message data" ) class ClientMessage(BaseModel): """Inbound message from client to server.""" action: str = Field(..., description="Action requested by client") data: Optional[Dict[str, Any]] = Field( default_factory=dict, description="Action payload" ) class RoomSubscriptionRequest(BaseModel): """Request to join or leave a room.""" action: str = Field( ..., description="Action: 'join' or 'leave'" ) room: str = Field( ..., min_length=1, description="Room name to join or leave" ) class ScanProgressMessage(BaseModel): """Scan progress update message.""" type: WebSocketMessageType = Field( default=WebSocketMessageType.SCAN_PROGRESS, description="Message type", ) timestamp: str = Field( default_factory=lambda: datetime.now(timezone.utc).isoformat(), description="ISO 8601 timestamp", ) data: Dict[str, Any] = Field( ..., description="Scan progress data including current, total, percent", ) class ScanCompleteMessage(BaseModel): """Scan completion message.""" type: WebSocketMessageType = Field( default=WebSocketMessageType.SCAN_COMPLETE, description="Message type", ) timestamp: str = Field( default_factory=lambda: datetime.now(timezone.utc).isoformat(), description="ISO 8601 timestamp", ) data: Dict[str, Any] = Field( ..., description="Scan completion data including series_found, duration", ) class ScanFailedMessage(BaseModel): """Scan failure message.""" type: WebSocketMessageType = Field( default=WebSocketMessageType.SCAN_FAILED, description="Message type", ) timestamp: str = Field( default_factory=lambda: datetime.now(timezone.utc).isoformat(), description="ISO 8601 timestamp", ) data: Dict[str, Any] = Field( ..., description="Scan error data including error_message" ) class ErrorNotificationMessage(BaseModel): """Error notification message for critical errors.""" type: WebSocketMessageType = Field( default=WebSocketMessageType.SYSTEM_ERROR, description="Message type", ) timestamp: str = Field( default_factory=lambda: datetime.now(timezone.utc).isoformat(), description="ISO 8601 timestamp", ) data: Dict[str, Any] = Field( ..., description=( "Error notification data including severity, message, details" ), ) class ProgressUpdateMessage(BaseModel): """Generic progress update message. Can be used for any type of progress (download, scan, queue, etc.) """ type: WebSocketMessageType = Field( ..., description="Type of progress message" ) timestamp: str = Field( default_factory=lambda: datetime.now(timezone.utc).isoformat(), description="ISO 8601 timestamp", ) data: Dict[str, Any] = Field( ..., description=( "Progress data including id, status, percent, current, total" ), )