Task_Reporter/app/modules/realtime/router.py

"""WebSocket and REST API router for realtime messaging"""
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, HTTPException, Query
from fastapi.responses import JSONResponse
from sqlalchemy.orm import Session
from typing import Optional
from datetime import datetime
import json

from app.core.database import get_db
from app.modules.auth.dependencies import get_current_user
from app.modules.chat_room.models import RoomMember, MemberRole
from app.modules.realtime.websocket_manager import manager
from app.modules.realtime.services.message_service import MessageService
from app.modules.realtime.schemas import (
    WebSocketMessageIn,
    MessageBroadcast,
    SystemMessageBroadcast,
    MessageAck,
    ErrorMessage,
    MessageCreate,
    MessageUpdate,
    MessageResponse,
    MessageListResponse,
    ReactionCreate,
    WebSocketMessageType,
    SystemEventType,
    MessageTypeEnum
)
from app.modules.realtime.models import MessageType, Message
from sqlalchemy import and_

router = APIRouter(prefix="/api", tags=["realtime"])

SYSTEM_ADMIN_EMAIL = "ymirliu@panjit.com.tw"


def get_user_room_membership(db: Session, room_id: str, user_id: str) -> Optional[RoomMember]:
    """Check if user is a member of the room"""
    return db.query(RoomMember).filter(
        and_(
            RoomMember.room_id == room_id,
            RoomMember.user_id == user_id,
            RoomMember.removed_at.is_(None)
        )
    ).first()


def can_write_message(membership: Optional[RoomMember], user_id: str) -> bool:
    """Check if user has write permission (OWNER or EDITOR)"""
    if user_id == SYSTEM_ADMIN_EMAIL:
        return True

    if not membership:
        return False

    return membership.role in [MemberRole.OWNER, MemberRole.EDITOR]


@router.websocket("/ws/{room_id}")
async def websocket_endpoint(
    websocket: WebSocket,
    room_id: str,
    token: Optional[str] = Query(None)
):
    """
    WebSocket endpoint for realtime messaging

    Authentication:
    - Token can be provided via query parameter: /ws/{room_id}?token=xxx
    - Or via WebSocket headers

    Connection flow:
    1. Client connects with room_id
    2. Server validates authentication and room membership
    3. Connection added to pool
    4. User joined event broadcast to room
    5. Client can send/receive messages
    """
    db: Session = next(get_db())

    try:
        # For now, we'll extract user from cookie or token
        # TODO: Implement proper WebSocket token authentication
        user_id = token if token else "anonymous@example.com"  # Placeholder

        # Check room membership
        membership = get_user_room_membership(db, room_id, user_id)
        if not membership and user_id != SYSTEM_ADMIN_EMAIL:
            await websocket.close(code=4001, reason="Not a member of this room")
            return

        # Connect to WebSocket manager
        conn_info = await manager.connect(websocket, room_id, user_id)

        # Broadcast user joined event
        await manager.broadcast_to_room(
            room_id,
            SystemMessageBroadcast(
                event=SystemEventType.USER_JOINED,
                user_id=user_id,
                room_id=room_id,
                timestamp=datetime.utcnow()
            ).dict(),
            exclude_user=user_id
        )

        try:
            while True:
                # Receive message from client
                data = await websocket.receive_text()
                message_data = json.loads(data)

                # Parse incoming message
                try:
                    ws_message = WebSocketMessageIn(**message_data)
                except Exception as e:
                    await websocket.send_json(
                        ErrorMessage(error=str(e), code="INVALID_MESSAGE").dict()
                    )
                    continue

                # Handle different message types
                if ws_message.type == WebSocketMessageType.MESSAGE:
                    # Check write permission
                    if not can_write_message(membership, user_id):
                        await websocket.send_json(
                            ErrorMessage(
                                error="Insufficient permissions",
                                code="PERMISSION_DENIED"
                            ).dict()
                        )
                        continue

                    # Create message in database
                    message = MessageService.create_message(
                        db=db,
                        room_id=room_id,
                        sender_id=user_id,
                        content=ws_message.content or "",
                        message_type=MessageType(ws_message.message_type.value) if ws_message.message_type else MessageType.TEXT,
                        metadata=ws_message.metadata
                    )

                    # Send acknowledgment to sender
                    await websocket.send_json(
                        MessageAck(
                            message_id=message.message_id,
                            sequence_number=message.sequence_number,
                            timestamp=message.created_at
                        ).dict()
                    )

                    # Broadcast message to all room members
                    await manager.broadcast_to_room(
                        room_id,
                        MessageBroadcast(
                            message_id=message.message_id,
                            room_id=message.room_id,
                            sender_id=message.sender_id,
                            content=message.content,
                            message_type=MessageTypeEnum(message.message_type.value),
                            metadata=message.message_metadata,
                            created_at=message.created_at,
                            sequence_number=message.sequence_number
                        ).dict()
                    )

                elif ws_message.type == WebSocketMessageType.EDIT_MESSAGE:
                    if not ws_message.message_id or not ws_message.content:
                        await websocket.send_json(
                            ErrorMessage(error="Missing message_id or content", code="INVALID_REQUEST").dict()
                        )
                        continue

                    # Edit message
                    edited_message = MessageService.edit_message(
                        db=db,
                        message_id=ws_message.message_id,
                        user_id=user_id,
                        new_content=ws_message.content
                    )

                    if not edited_message:
                        await websocket.send_json(
                            ErrorMessage(error="Cannot edit message", code="EDIT_FAILED").dict()
                        )
                        continue

                    # Broadcast edit to all room members
                    await manager.broadcast_to_room(
                        room_id,
                        MessageBroadcast(
                            type="edit_message",
                            message_id=edited_message.message_id,
                            room_id=edited_message.room_id,
                            sender_id=edited_message.sender_id,
                            content=edited_message.content,
                            message_type=MessageTypeEnum(edited_message.message_type.value),
                            metadata=edited_message.message_metadata,
                            created_at=edited_message.created_at,
                            edited_at=edited_message.edited_at,
                            sequence_number=edited_message.sequence_number
                        ).dict()
                    )

                elif ws_message.type == WebSocketMessageType.DELETE_MESSAGE:
                    if not ws_message.message_id:
                        await websocket.send_json(
                            ErrorMessage(error="Missing message_id", code="INVALID_REQUEST").dict()
                        )
                        continue

                    # Delete message
                    is_admin = user_id == SYSTEM_ADMIN_EMAIL
                    deleted_message = MessageService.delete_message(
                        db=db,
                        message_id=ws_message.message_id,
                        user_id=user_id,
                        is_admin=is_admin
                    )

                    if not deleted_message:
                        await websocket.send_json(
                            ErrorMessage(error="Cannot delete message", code="DELETE_FAILED").dict()
                        )
                        continue

                    # Broadcast deletion to all room members
                    await manager.broadcast_to_room(
                        room_id,
                        {"type": "delete_message", "message_id": deleted_message.message_id}
                    )

                elif ws_message.type == WebSocketMessageType.ADD_REACTION:
                    if not ws_message.message_id or not ws_message.emoji:
                        await websocket.send_json(
                            ErrorMessage(error="Missing message_id or emoji", code="INVALID_REQUEST").dict()
                        )
                        continue

                    # Add reaction
                    reaction = MessageService.add_reaction(
                        db=db,
                        message_id=ws_message.message_id,
                        user_id=user_id,
                        emoji=ws_message.emoji
                    )

                    if reaction:
                        # Broadcast reaction to all room members
                        await manager.broadcast_to_room(
                            room_id,
                            {
                                "type": "add_reaction",
                                "message_id": ws_message.message_id,
                                "user_id": user_id,
                                "emoji": ws_message.emoji
                            }
                        )

                elif ws_message.type == WebSocketMessageType.REMOVE_REACTION:
                    if not ws_message.message_id or not ws_message.emoji:
                        await websocket.send_json(
                            ErrorMessage(error="Missing message_id or emoji", code="INVALID_REQUEST").dict()
                        )
                        continue

                    # Remove reaction
                    removed = MessageService.remove_reaction(
                        db=db,
                        message_id=ws_message.message_id,
                        user_id=user_id,
                        emoji=ws_message.emoji
                    )

                    if removed:
                        # Broadcast reaction removal to all room members
                        await manager.broadcast_to_room(
                            room_id,
                            {
                                "type": "remove_reaction",
                                "message_id": ws_message.message_id,
                                "user_id": user_id,
                                "emoji": ws_message.emoji
                            }
                        )

                elif ws_message.type == WebSocketMessageType.TYPING:
                    # Set typing status
                    is_typing = message_data.get("is_typing", True)
                    await manager.set_typing(room_id, user_id, is_typing)

                    # Broadcast typing status to other room members
                    await manager.broadcast_to_room(
                        room_id,
                        {"type": "typing", "user_id": user_id, "is_typing": is_typing},
                        exclude_user=user_id
                    )

        except WebSocketDisconnect:
            pass
        finally:
            # Disconnect and broadcast user left event
            await manager.disconnect(conn_info)
            await manager.broadcast_to_room(
                room_id,
                SystemMessageBroadcast(
                    event=SystemEventType.USER_LEFT,
                    user_id=user_id,
                    room_id=room_id,
                    timestamp=datetime.utcnow()
                ).dict()
            )

    finally:
        db.close()


# REST API endpoints
@router.get("/rooms/{room_id}/messages", response_model=MessageListResponse)
async def get_messages(
    room_id: str,
    limit: int = Query(50, ge=1, le=100),
    before: Optional[datetime] = None,
    offset: int = Query(0, ge=0),
    current_user: dict = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """Get message history for a room"""
    user_id = current_user["username"]

    # Check room membership
    membership = get_user_room_membership(db, room_id, user_id)
    if not membership and user_id != SYSTEM_ADMIN_EMAIL:
        raise HTTPException(status_code=403, detail="Not a member of this room")

    return MessageService.get_messages(
        db=db,
        room_id=room_id,
        limit=limit,
        before_timestamp=before,
        offset=offset
    )


@router.post("/rooms/{room_id}/messages", response_model=MessageResponse, status_code=201)
async def create_message(
    room_id: str,
    message: MessageCreate,
    current_user: dict = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """Create a message via REST API (alternative to WebSocket)"""
    user_id = current_user["username"]

    # Check room membership and write permission
    membership = get_user_room_membership(db, room_id, user_id)
    if not can_write_message(membership, user_id):
        raise HTTPException(status_code=403, detail="Insufficient permissions")

    # Create message
    created_message = MessageService.create_message(
        db=db,
        room_id=room_id,
        sender_id=user_id,
        content=message.content,
        message_type=MessageType(message.message_type.value),
        metadata=message.metadata
    )

    # Broadcast to WebSocket connections
    await manager.broadcast_to_room(
        room_id,
        MessageBroadcast(
            message_id=created_message.message_id,
            room_id=created_message.room_id,
            sender_id=created_message.sender_id,
            content=created_message.content,
            message_type=MessageTypeEnum(created_message.message_type.value),
            metadata=created_message.message_metadata,
            created_at=created_message.created_at,
            sequence_number=created_message.sequence_number
        ).dict()
    )

    return MessageResponse.from_orm(created_message)


@router.get("/rooms/{room_id}/messages/search", response_model=MessageListResponse)
async def search_messages(
    room_id: str,
    q: str = Query(..., min_length=1),
    limit: int = Query(50, ge=1, le=100),
    offset: int = Query(0, ge=0),
    current_user: dict = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """Search messages in a room"""
    user_id = current_user["username"]

    # Check room membership
    membership = get_user_room_membership(db, room_id, user_id)
    if not membership and user_id != SYSTEM_ADMIN_EMAIL:
        raise HTTPException(status_code=403, detail="Not a member of this room")

    return MessageService.search_messages(
        db=db,
        room_id=room_id,
        query=q,
        limit=limit,
        offset=offset
    )


@router.get("/rooms/{room_id}/online")
async def get_online_users(
    room_id: str,
    current_user: dict = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """Get list of online users in a room"""
    user_id = current_user["username"]

    # Check room membership
    membership = get_user_room_membership(db, room_id, user_id)
    if not membership and user_id != SYSTEM_ADMIN_EMAIL:
        raise HTTPException(status_code=403, detail="Not a member of this room")

    online_users = manager.get_online_users(room_id)
    return {"room_id": room_id, "online_users": online_users, "count": len(online_users)}


@router.get("/rooms/{room_id}/typing")
async def get_typing_users(
    room_id: str,
    current_user: dict = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    """Get list of users currently typing in a room"""
    user_id = current_user["username"]

    # Check room membership
    membership = get_user_room_membership(db, room_id, user_id)
    if not membership and user_id != SYSTEM_ADMIN_EMAIL:
        raise HTTPException(status_code=403, detail="Not a member of this room")

    typing_users = manager.get_typing_users(room_id)
    return {"room_id": room_id, "typing_users": typing_users, "count": len(typing_users)}