"""
WebSocket message models for real-time updates.
"""

from datetime import datetime
from enum import Enum
from typing import Any, Dict, Optional

from pydantic import BaseModel, Field


class WSMessageType(str, Enum):
    """WebSocket message types."""

    DEVICE_UPDATE = "device_update"
    TASK_STARTED = "task_started"
    TASK_STEP = "task_step"
    TASK_COMPLETED = "task_completed"
    TASK_FAILED = "task_failed"
    TASK_STOPPED = "task_stopped"
    SCREENSHOT = "screenshot"
    ERROR = "error"
    PING = "ping"
    PONG = "pong"


class WSMessage(BaseModel):
    """Base WebSocket message."""

    type: WSMessageType = Field(..., description="Message type")
    data: Dict[str, Any] = Field(default_factory=dict, description="Message data")
    timestamp: datetime = Field(default_factory=datetime.now, description="Message timestamp")

    class Config:
        json_schema_extra = {
            "example": {
                "type": "task_step",
                "data": {
                    "task_id": "task_123",
                    "device_id": "emulator-5554",
                    "step": 5,
                    "action": {"action": "Tap"},
                    "thinking": "Tapping on button",
                },
                "timestamp": "2024-01-09T10:00:00",
            }
        }


class StepUpdate(BaseModel):
    """Step update message data."""

    task_id: str = Field(..., description="Task ID")
    device_id: str = Field(..., description="Device ID")
    step: int = Field(..., description="Step number")
    action: Optional[Dict[str, Any]] = Field(None, description="Action taken")
    thinking: Optional[str] = Field(None, description="AI reasoning")
    finished: bool = Field(False, description="Whether task is finished")
    success: bool = Field(True, description="Whether step succeeded")
    message: Optional[str] = Field(None, description="Status message")


class ScreenshotUpdate(BaseModel):
    """Screenshot update message data."""

    device_id: str = Field(..., description="Device ID")
    screenshot: str = Field(..., description="Base64 encoded screenshot")
    width: int = Field(..., description="Screenshot width")
    height: int = Field(..., description="Screenshot height")
    timestamp: datetime = Field(default_factory=datetime.now)