Task_Reporter/app/modules/realtime/schemas.py

"""Pydantic schemas for WebSocket messages and REST API"""
from pydantic import BaseModel, Field
from typing import Optional, Dict, Any, List
from datetime import datetime
from enum import Enum


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


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


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


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
    metadata: Optional[Dict[str, Any]] = Field(None, 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

    class Config:
        from_attributes = True
        populate_by_name = True  # Allow both 'metadata' and 'message_metadata'


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

    class Config:
        from_attributes = True


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


# File Upload WebSocket Schemas
class FileUploadedBroadcast(BaseModel):
    """Broadcast when a file is uploaded to a room"""
    type: str = "file_uploaded"
    file_id: str
    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
    uploaded_at: datetime

    def to_dict(self) -> dict:
        """Convert to dictionary for WebSocket broadcast"""
        return {
            "type": self.type,
            "file_id": self.file_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,
            "uploaded_at": self.uploaded_at.isoformat()
        }


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
    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,
            "room_id": self.room_id,
            "deleted_by": self.deleted_by,
            "deleted_at": self.deleted_at.isoformat()
        }