radionica3d/backend/routers/chat.py

from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, Query, HTTPException
from services.chat_manager import manager
from services.global_manager import global_manager
import db
import auth_utils
import datetime
import schemas

router = APIRouter(tags=["chat"])

@router.get("/orders/{order_id}/messages")
async def get_order_messages(order_id: int, token: str = Depends(auth_utils.oauth2_scheme)):
    payload = auth_utils.decode_token(token)
    if not payload: raise HTTPException(status_code=401, detail="Invalid token")
    role = payload.get("role")
    user_id = payload.get("id")
    order = db.execute_query("SELECT user_id FROM orders WHERE id = %s", (order_id,))
    if not order: raise HTTPException(status_code=404, detail="Order not found")
    if role != 'admin' and order[0]['user_id'] != user_id: raise HTTPException(status_code=403, detail="Not authorized")
    messages = db.execute_query("SELECT id, is_from_admin, message, created_at FROM order_messages WHERE order_id = %s ORDER BY created_at ASC", (order_id,))
    for msg in messages:
        if msg.get('created_at'): msg['created_at'] = msg['created_at'].isoformat()
        
    # Mark messages as read
    if role == 'admin':
        db.execute_commit("UPDATE order_messages SET is_read = TRUE WHERE order_id = %s AND is_from_admin = FALSE AND is_read = FALSE", (order_id,))
        await global_manager.notify_admins()
    else:
        db.execute_commit("UPDATE order_messages SET is_read = TRUE WHERE order_id = %s AND is_from_admin = TRUE AND is_read = FALSE", (order_id,))
        await global_manager.notify_user(user_id)
        
    return messages

@router.post("/orders/{order_id}/messages")
async def post_order_message(order_id: int, data: schemas.MessageCreate, token: str = Depends(auth_utils.oauth2_scheme)):
    payload = auth_utils.decode_token(token)
    if not payload: raise HTTPException(status_code=401, detail="Invalid token")
    message = data.message.strip()
    if not message: raise HTTPException(status_code=400, detail="Empty message")
    role = payload.get("role")
    user_id = payload.get("id")
    is_admin = (role == 'admin')
    order = db.execute_query("SELECT user_id FROM orders WHERE id = %s", (order_id,))
    if not order: raise HTTPException(status_code=404, detail="Order not found")
    if not is_admin and order[0]['user_id'] != user_id: raise HTTPException(status_code=403, detail="Not authorized")
    query = "INSERT INTO order_messages (order_id, user_id, is_from_admin, message) VALUES (%s, %s, %s, %s)"
    msg_id = db.execute_commit(query, (order_id, user_id, is_admin, message))
    now = datetime.datetime.utcnow().isoformat()
    await manager.broadcast_to_order(order_id, {"id": msg_id, "is_from_admin": is_admin, "message": message, "created_at": now})
    if is_admin:
        await global_manager.notify_user(order[0]['user_id'])
    else:
        await global_manager.notify_admins()
    return {"id": msg_id, "status": "sent"}

@router.websocket("/ws/chat/{order_id}")
async def ws_chat(websocket: WebSocket, order_id: int, token: str = Query(...)):
    payload = auth_utils.decode_token(token)
    if not payload:
        await websocket.close(code=4001)
        return
    role = payload.get("role")
    user_id = payload.get("id")
    order = db.execute_query("SELECT user_id FROM orders WHERE id = %s", (order_id,))
    if not order:
        await websocket.close(code=4004)
        return
    if role != 'admin' and order[0]['user_id'] != user_id:
        await websocket.close(code=4003)
        return
    await manager.connect(websocket, order_id, role)
    try:
        while True:
            data = await websocket.receive_text()
            if data == "typing":
                await manager.broadcast_to_order(order_id, {"type": "typing", "is_admin": role == 'admin'})
            elif data == "stop_typing":
                await manager.broadcast_to_order(order_id, {"type": "stop_typing", "is_admin": role == 'admin'})
            elif data == "read":
                if role == 'admin':
                    db.execute_commit("UPDATE order_messages SET is_read = TRUE WHERE order_id = %s AND is_from_admin = FALSE AND is_read = FALSE", (order_id,))
                    await global_manager.notify_admins()
                else:
                    db.execute_commit("UPDATE order_messages SET is_read = TRUE WHERE order_id = %s AND is_from_admin = TRUE AND is_read = FALSE", (order_id,))
                    await global_manager.notify_user(user_id)
    except WebSocketDisconnect:
        manager.disconnect(websocket, order_id)