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 import locales from dependencies import get_current_user router = APIRouter(tags=["chat"]) # In-memory storage for flood control: {user_id: timestamp} last_message_times = {} @router.get("/orders/{order_id}/messages") async def get_order_messages(order_id: int, user: dict = Depends(get_current_user)): role = user.get("role") user_id = user.get("id") # Fetch user chat status user_info = db.execute_query("SELECT can_chat FROM users WHERE id = %s", (user_id,)) can_chat = user_info[0]['can_chat'] if user_info else False 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': if order[0]['user_id'] != user_id: raise HTTPException(status_code=403, detail="Not authorized") if not can_chat: raise HTTPException(status_code=403, detail="Chat access disabled for your account") 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() msg['is_from_admin'] = bool(msg['is_from_admin']) # 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() await global_manager.notify_order_read(order_id) 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, user: dict = Depends(get_current_user), lang: str = "en"): role = user.get("role") user_id = user.get("id") # Flood control for non-admin users if role != 'admin': now = datetime.datetime.utcnow().timestamp() last_time = last_message_times.get(user_id, 0) if now - last_time < 10: raise HTTPException(status_code=429, detail=locales.translate_error("flood_control", lang)) last_message_times[user_id] = now 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') if not is_admin: user_info = db.execute_query("SELECT can_chat FROM users WHERE id = %s", (user_id,)) if not user_info or not user_info[0]['can_chat']: raise HTTPException(status_code=403, detail="Chat access disabled") 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: from services import event_hooks await global_manager.notify_admins() await global_manager.notify_admins_new_message(order_id, message) event_hooks.on_message_received(order_id, user_id, message) 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") if role != 'admin': user_info = db.execute_query("SELECT can_chat FROM users WHERE id = %s", (user_id,)) if not user_info or not user_info[0]['can_chat']: await websocket.close(code=4003) return 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() await global_manager.notify_order_read(order_id) 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)