chat.py 4.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788
  1. from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, Query, HTTPException
  2. from services.chat_manager import manager
  3. from services.global_manager import global_manager
  4. import db
  5. import auth_utils
  6. import datetime
  7. import schemas
  8. import locales
  9. from dependencies import get_current_user
  10. router = APIRouter(tags=["chat"])
  11. # In-memory storage for flood control: {user_id: timestamp}
  12. last_message_times = {}
  13. @router.get("/orders/{order_id}/messages")
  14. async def get_order_messages(order_id: int, user: dict = Depends(get_current_user)):
  15. role = user.get("role")
  16. user_id = user.get("id")
  17. # Fetch user chat status
  18. user_info = db.execute_query("SELECT can_chat FROM users WHERE id = %s", (user_id,))
  19. can_chat = user_info[0]['can_chat'] if user_info else False
  20. order = db.execute_query("SELECT user_id FROM orders WHERE id = %s", (order_id,))
  21. if not order: raise HTTPException(status_code=404, detail="Order not found")
  22. if role != 'admin':
  23. if order[0]['user_id'] != user_id: raise HTTPException(status_code=403, detail="Not authorized")
  24. if not can_chat: raise HTTPException(status_code=403, detail="Chat access disabled for your account")
  25. 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,))
  26. for msg in messages:
  27. if msg.get('created_at'): msg['created_at'] = msg['created_at'].isoformat()
  28. msg['is_from_admin'] = bool(msg['is_from_admin'])
  29. # Mark messages as read
  30. if role == 'admin':
  31. 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,))
  32. await global_manager.notify_admins()
  33. await global_manager.notify_order_read(order_id)
  34. else:
  35. 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,))
  36. await global_manager.notify_user(user_id)
  37. return messages
  38. @router.post("/orders/{order_id}/messages")
  39. async def post_order_message(order_id: int, data: schemas.MessageCreate, user: dict = Depends(get_current_user), lang: str = "en"):
  40. role = user.get("role")
  41. user_id = user.get("id")
  42. # Flood control for non-admin users
  43. if role != 'admin':
  44. now = datetime.datetime.utcnow().timestamp()
  45. last_time = last_message_times.get(user_id, 0)
  46. if now - last_time < 10:
  47. raise HTTPException(status_code=429, detail=locales.translate_error("flood_control", lang))
  48. last_message_times[user_id] = now
  49. message = data.message.strip()
  50. if not message: raise HTTPException(status_code=400, detail="Empty message")
  51. role = payload.get("role")
  52. user_id = payload.get("id")
  53. is_admin = (role == 'admin')
  54. if not is_admin:
  55. user_info = db.execute_query("SELECT can_chat FROM users WHERE id = %s", (user_id,))
  56. if not user_info or not user_info[0]['can_chat']:
  57. raise HTTPException(status_code=403, detail="Chat access disabled")
  58. order = db.execute_query("SELECT user_id FROM orders WHERE id = %s", (order_id,))
  59. if not order: raise HTTPException(status_code=404, detail="Order not found")
  60. if not is_admin and order[0]['user_id'] != user_id: raise HTTPException(status_code=403, detail="Not authorized")
  61. query = "INSERT INTO order_messages (order_id, user_id, is_from_admin, message) VALUES (%s, %s, %s, %s)"
  62. msg_id = db.execute_commit(query, (order_id, user_id, is_admin, message))
  63. now = datetime.datetime.utcnow().isoformat()
  64. await manager.broadcast_to_order(order_id, {"id": msg_id, "is_from_admin": is_admin, "message": message, "created_at": now})
  65. if is_admin:
  66. await global_manager.notify_user(order[0]['user_id'])
  67. else:
  68. from services import event_hooks
  69. await global_manager.notify_admins()
  70. await global_manager.notify_admins_new_message(order_id, message)
  71. event_hooks.on_message_received(order_id, user_id, message)
  72. return {"id": msg_id, "status": "sent"}
  73. # WebSocket moved to main.py
  74. # @router.websocket("/ws/chat/{order_id}")
  75. # async def ws_chat(websocket: WebSocket, order_id: int, token: str = Query(...)):
  76. # ...