web/backend/routes/messaging.py

from datetime import datetime
import logging
from fastapi import APIRouter, Depends, HTTPException, WebSocket, WebSocketDisconnect
from fastapi.security import HTTPAuthorizationCredentials
from sqlalchemy.orm import Session
from dependencies import get_current_user, get_db
from models import Message, SendMessageRequest, User

router = APIRouter()
logger = logging.getLogger("uvicorn.error")

def convert_message(msg: Message) -> dict:
    return {
        "id": msg.id,
        "content": msg.content,
        "timestamp": msg.timestamp.isoformat(),
        "is_read": msg.is_read,
        "username": msg.author.username,
        "profile_picture": msg.author.profile_picture
    }

async def get_messages_inner(db: Session):
    messages = db.query(Message).order_by(Message.timestamp.asc()).all()

    messages_data = []
    for msg in messages:
        messages_data.append(convert_message(msg))

    return {
        "status": "success",
        "messages": messages_data
    }

async def send_message_inner(request: SendMessageRequest, current_user: User, db: Session):
    if not request.content.strip():
        raise HTTPException(
            status_code=400,
            detail="No content provided"
        )

    new_message = Message(
        content=request.content.strip(),
        user_id=current_user.id,
        timestamp=datetime.now()
    )

    db.add(new_message)
    db.commit()
    db.refresh(new_message)

    return {"status": "success", "message": convert_message(new_message)}

@router.post("/send_message")
async def send_message(
    request: SendMessageRequest,
    current_user: User = Depends(get_current_user),
    db: Session = Depends(get_db)
):
    return await send_message_inner(request, current_user, db)


@router.get("/get_messages")
async def get_messages(db: Session = Depends(get_db)):
    return await get_messages_inner(db)


class MessaggingSocketManager:
    def __init__(self) -> None:
        self.connections: list[WebSocket] = []

    async def send_error(self, websocket: WebSocket, type: str, e: HTTPException):
        await websocket.send_json({"type": type, "error": {"code": e.status_code, "detail": e.detail}})

    async def handle_connection(self, websocket: WebSocket, db: Session):
        while True:
            data = await websocket.receive_json()
            type = data["type"]

            def get_current_user_inner() -> User | None:
                if data["credentials"]:
                    return get_current_user(
                        HTTPAuthorizationCredentials(
                            scheme=data["credentials"]["scheme"], 
                            credentials=data["credentials"]["credentials"]
                        ), 
                        db
                    )
                else:
                    return None

            if type == "ping":
                await websocket.send_json({"type": "ping", "data": {"status": "success"}})
            elif type == "getMessages":
                try:
                    current_user = get_current_user_inner()
                    if not current_user:
                        raise HTTPException(401)

                    await websocket.send_json({"type": type, "data": await get_messages_inner(current_user, db)})
                except HTTPException as e:
                    await self.send_error(websocket, type, e)
            elif type == "sendMessage":
                try:
                    current_user = get_current_user_inner()
                    if not current_user:
                        raise HTTPException(401)
                    
                    request: SendMessageRequest = SendMessageRequest.model_validate(data["data"])

                    response = await send_message_inner(request, current_user, db)
                    await self.broadcast({
                        "type": "newMessage",
                        "data": response["message"]
                    })

                    await websocket.send_json({"type": type, "data": response})
                except HTTPException as e:
                    await self.send_error(websocket, type, e)
            else:
                await websocket.send_json({"type": type, "error": {"code": 400, "detail": "Invalid type"}})

    async def disconnect(self, websocket: WebSocket, code: int = 1000, message: str | None = None):
        try:
            await websocket.close(code=code, reason=message)
        finally: 
            self.connections.remove(websocket)
    
    async def connect(self, websocket: WebSocket, db: Session):
        await websocket.accept()
        self.connections.append(websocket)
        try:
            await self.handle_connection(websocket, db)
        except WebSocketDisconnect as e:
            logger.info(f"WebSocket disconnected with code {e.code}: {e.reason}")
        finally:
            self.connections.remove(websocket)

    async def broadcast(self, message: dict):
        for websocket in self.connections:
            await websocket.send_json(message)

messagingManager = MessaggingSocketManager()

@router.websocket("/chat/ws")
async def chat_websocket(
    websocket: WebSocket,
    db: Session = Depends(get_db)
):
    await messagingManager.connect(websocket, db)