"""
Message data structures for WebSocket communication.
Based on worker.md protocol specification.
"""
from typing import Dict, Any, List, Optional, Union, Literal
from pydantic import BaseModel, Field
from datetime import datetime


class SubscriptionObject(BaseModel):
    """Individual camera subscription configuration."""
    subscriptionIdentifier: str = Field(..., description="Format: displayId;cameraId")
    rtspUrl: Optional[str] = Field(None, description="RTSP stream URL")
    snapshotUrl: Optional[str] = Field(None, description="HTTP snapshot URL")
    snapshotInterval: Optional[int] = Field(None, description="Snapshot interval in milliseconds")
    modelUrl: str = Field(..., description="Pre-signed URL to .mpta file")
    modelId: int = Field(..., description="Unique model identifier")
    modelName: str = Field(..., description="Human-readable model name")
    cropX1: Optional[int] = Field(None, description="Crop region X1 coordinate")
    cropY1: Optional[int] = Field(None, description="Crop region Y1 coordinate")
    cropX2: Optional[int] = Field(None, description="Crop region X2 coordinate")
    cropY2: Optional[int] = Field(None, description="Crop region Y2 coordinate")


class CameraConnection(BaseModel):
    """Camera connection status for state reporting."""
    subscriptionIdentifier: str
    modelId: int
    modelName: str
    online: bool
    cropX1: Optional[int] = None
    cropY1: Optional[int] = None
    cropX2: Optional[int] = None
    cropY2: Optional[int] = None


class DetectionData(BaseModel):
    """Detection result data structure."""
    detection: Dict[str, Any] = Field(..., description="Flat key-value detection results")
    modelId: int
    modelName: str


# Incoming Messages from Backend to Worker

class SetSubscriptionListMessage(BaseModel):
    """Complete subscription list for declarative state management."""
    type: Literal["setSubscriptionList"] = "setSubscriptionList"
    subscriptions: List[SubscriptionObject]


class SetSessionIdPayload(BaseModel):
    """Session ID association payload."""
    displayIdentifier: str
    sessionId: Optional[int] = None


class SetSessionIdMessage(BaseModel):
    """Associate session ID with display."""
    type: Literal["setSessionId"] = "setSessionId"
    payload: SetSessionIdPayload


class SetProgressionStagePayload(BaseModel):
    """Progression stage payload."""
    displayIdentifier: str
    progressionStage: Optional[str] = None


class SetProgressionStageMessage(BaseModel):
    """Set progression stage for display."""
    type: Literal["setProgressionStage"] = "setProgressionStage"
    payload: SetProgressionStagePayload


class RequestStateMessage(BaseModel):
    """Request current worker state."""
    type: Literal["requestState"] = "requestState"


class PatchSessionResultPayload(BaseModel):
    """Patch session result payload."""
    sessionId: int
    success: bool
    message: str


class PatchSessionResultMessage(BaseModel):
    """Response to patch session request."""
    type: Literal["patchSessionResult"] = "patchSessionResult"
    payload: PatchSessionResultPayload


# Outgoing Messages from Worker to Backend

class StateReportMessage(BaseModel):
    """Periodic heartbeat with system metrics."""
    type: Literal["stateReport"] = "stateReport"
    cpuUsage: float
    memoryUsage: float
    gpuUsage: Optional[float] = None
    gpuMemoryUsage: Optional[float] = None
    cameraConnections: List[CameraConnection]


class ImageDetectionMessage(BaseModel):
    """Detection event message."""
    type: Literal["imageDetection"] = "imageDetection"
    subscriptionIdentifier: str
    timestamp: str = Field(default_factory=lambda: datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%S.%fZ"))
    data: DetectionData


class PatchSessionMessage(BaseModel):
    """Request to modify session data."""
    type: Literal["patchSession"] = "patchSession"
    sessionId: int
    data: Dict[str, Any] = Field(..., description="Partial DisplayPersistentData structure")


# Union type for all incoming messages
IncomingMessage = Union[
    SetSubscriptionListMessage,
    SetSessionIdMessage,
    SetProgressionStageMessage,
    RequestStateMessage,
    PatchSessionResultMessage
]

# Union type for all outgoing messages
OutgoingMessage = Union[
    StateReportMessage,
    ImageDetectionMessage,
    PatchSessionMessage
]