Task_Reporter/app/modules/realtime/schemas.py

"""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"
        }