"""Pydantic schemas for WebSocket messages and REST API""" from pydantic import BaseModel, Field, ConfigDict, field_serializer from typing import Optional, Dict, Any, List from datetime import datetime from enum import Enum def serialize_datetime_utc(dt: datetime) -> str: """Serialize datetime with 'Z' suffix to indicate UTC""" if dt is None: return None return dt.isoformat() + "Z" class MessageTypeEnum(str, Enum): """Message type enumeration for validation""" TEXT = "text" IMAGE_REF = "image_ref" FILE_REF = "file_ref" SYSTEM = "system" INCIDENT_DATA = "incident_data" class WebSocketMessageType(str, Enum): """WebSocket message type for protocol""" MESSAGE = "message" EDIT_MESSAGE = "edit_message" DELETE_MESSAGE = "delete_message" ADD_REACTION = "add_reaction" REMOVE_REACTION = "remove_reaction" TYPING = "typing" SYSTEM = "system" class SystemEventType(str, Enum): """System event types""" USER_JOINED = "user_joined" USER_LEFT = "user_left" ROOM_STATUS_CHANGED = "room_status_changed" MEMBER_ADDED = "member_added" MEMBER_REMOVED = "member_removed" FILE_UPLOADED = "file_uploaded" FILE_DELETED = "file_deleted" # WebSocket Incoming Messages (from client) class WebSocketMessageIn(BaseModel): """Incoming WebSocket message from client""" type: WebSocketMessageType content: Optional[str] = None message_type: Optional[MessageTypeEnum] = MessageTypeEnum.TEXT message_id: Optional[str] = None # For edit/delete/reaction operations emoji: Optional[str] = None # For reactions metadata: Optional[Dict[str, Any]] = None # For mentions, file refs, etc. class TextMessageIn(BaseModel): """Text message input""" content: str = Field(..., min_length=1, max_length=10000) mentions: Optional[List[str]] = None class ImageRefMessageIn(BaseModel): """Image reference message input""" content: str # Description file_id: str file_url: str class FileRefMessageIn(BaseModel): """File reference message input""" content: str # Description file_id: str file_url: str file_name: str class IncidentDataMessageIn(BaseModel): """Structured incident data message input""" content: Dict[str, Any] # Structured data (temperature, pressure, etc.) # WebSocket Outgoing Messages (to client) class MessageBroadcast(BaseModel): """Message broadcast to all room members""" type: str = "message" message_id: str room_id: str sender_id: str sender_display_name: Optional[str] = None # Display name from users table content: str message_type: MessageTypeEnum metadata: Optional[Dict[str, Any]] = None created_at: datetime edited_at: Optional[datetime] = None deleted_at: Optional[datetime] = None sequence_number: int @field_serializer("created_at", "edited_at", "deleted_at") def serialize_datetime(self, dt: Optional[datetime]) -> Optional[str]: """Serialize datetime with 'Z' suffix to indicate UTC""" if dt is None: return None return dt.isoformat() + "Z" class SystemMessageBroadcast(BaseModel): """System message broadcast""" type: str = "system" event: SystemEventType user_id: Optional[str] = None room_id: Optional[str] = None timestamp: datetime data: Optional[Dict[str, Any]] = None @field_serializer("timestamp") def serialize_datetime(self, dt: datetime) -> str: """Serialize datetime with 'Z' suffix to indicate UTC""" return dt.isoformat() + "Z" class TypingBroadcast(BaseModel): """Typing indicator broadcast""" type: str = "typing" room_id: str user_id: str is_typing: bool class MessageAck(BaseModel): """Message acknowledgment""" type: str = "ack" message_id: str sequence_number: int timestamp: datetime @field_serializer("timestamp") def serialize_datetime(self, dt: datetime) -> str: """Serialize datetime with 'Z' suffix to indicate UTC""" return dt.isoformat() + "Z" class ErrorMessage(BaseModel): """Error message""" type: str = "error" error: str code: str details: Optional[Dict[str, Any]] = None # REST API Schemas class MessageCreate(BaseModel): """Create message via REST API""" content: str = Field(..., min_length=1, max_length=10000) message_type: MessageTypeEnum = MessageTypeEnum.TEXT metadata: Optional[Dict[str, Any]] = None class MessageUpdate(BaseModel): """Update message content""" content: str = Field(..., min_length=1, max_length=10000) class MessageResponse(BaseModel): """Message response""" message_id: str room_id: str sender_id: str sender_display_name: Optional[str] = None # Display name from users table content: str message_type: MessageTypeEnum # Use validation_alias to read from ORM's message_metadata, but serialize as "metadata" metadata: Optional[Dict[str, Any]] = Field(None, validation_alias="message_metadata") created_at: datetime edited_at: Optional[datetime] = None deleted_at: Optional[datetime] = None sequence_number: int reaction_counts: Optional[Dict[str, int]] = None # emoji -> count model_config = ConfigDict( from_attributes=True, populate_by_name=True, ) @field_serializer("created_at", "edited_at", "deleted_at") def serialize_datetime(self, dt: Optional[datetime]) -> Optional[str]: """Serialize datetime with 'Z' suffix to indicate UTC""" if dt is None: return None return dt.isoformat() + "Z" class MessageListResponse(BaseModel): """Paginated message list response""" messages: List[MessageResponse] total: int limit: int offset: int has_more: bool class ReactionCreate(BaseModel): """Add reaction to message""" emoji: str = Field(..., min_length=1, max_length=10) class ReactionResponse(BaseModel): """Reaction response""" reaction_id: int message_id: str user_id: str emoji: str created_at: datetime model_config = ConfigDict(from_attributes=True) @field_serializer("created_at") def serialize_datetime(self, dt: datetime) -> str: """Serialize datetime with 'Z' suffix to indicate UTC""" return dt.isoformat() + "Z" class ReactionSummary(BaseModel): """Reaction summary for a message""" emoji: str count: int users: List[str] # List of user IDs who reacted class OnlineUser(BaseModel): """Online user in a room""" user_id: str connected_at: datetime @field_serializer("connected_at") def serialize_datetime(self, dt: datetime) -> str: """Serialize datetime with 'Z' suffix to indicate UTC""" return dt.isoformat() + "Z" # File Upload WebSocket Schemas class FileUploadedBroadcast(BaseModel): """Broadcast when a file is uploaded to a room""" type: str = "file_uploaded" file_id: str message_id: Optional[str] = None # Associated chat message ID room_id: str uploader_id: str filename: str file_type: str # image, document, log file_size: int mime_type: str download_url: Optional[str] = None thumbnail_url: Optional[str] = None # Thumbnail URL for images uploaded_at: datetime def to_dict(self) -> dict: """Convert to dictionary for WebSocket broadcast""" return { "type": self.type, "file_id": self.file_id, "message_id": self.message_id, "room_id": self.room_id, "uploader_id": self.uploader_id, "filename": self.filename, "file_type": self.file_type, "file_size": self.file_size, "mime_type": self.mime_type, "download_url": self.download_url, "thumbnail_url": self.thumbnail_url, "uploaded_at": self.uploaded_at.isoformat() + "Z" } class FileUploadAck(BaseModel): """Acknowledgment sent to uploader after successful upload""" type: str = "file_upload_ack" file_id: str status: str # success, error download_url: Optional[str] = None error_message: Optional[str] = None def to_dict(self) -> dict: """Convert to dictionary for WebSocket message""" return { "type": self.type, "file_id": self.file_id, "status": self.status, "download_url": self.download_url, "error_message": self.error_message } class FileDeletedBroadcast(BaseModel): """Broadcast when a file is deleted from a room""" type: str = "file_deleted" file_id: str message_id: Optional[str] = None # Associated chat message ID (also deleted) room_id: str deleted_by: str deleted_at: datetime def to_dict(self) -> dict: """Convert to dictionary for WebSocket broadcast""" return { "type": self.type, "file_id": self.file_id, "message_id": self.message_id, "room_id": self.room_id, "deleted_by": self.deleted_by, "deleted_at": self.deleted_at.isoformat() + "Z" } class MessageDeletedBroadcast(BaseModel): """Broadcast when a message is deleted""" type: str = "message_deleted" message_id: str room_id: str deleted_by: str deleted_at: datetime def to_dict(self) -> dict: """Convert to dictionary for WebSocket broadcast""" return { "type": self.type, "message_id": self.message_id, "room_id": self.room_id, "deleted_by": self.deleted_by, "deleted_at": self.deleted_at.isoformat() + "Z" }