chat.py 5.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144
  1. from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, Query, HTTPException, File, UploadFile, Form
  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. import os
  11. import uuid
  12. import config
  13. import json
  14. from typing import Optional
  15. router = APIRouter(tags=["chat"])
  16. # In-memory storage for flood control: {user_id: timestamp}
  17. last_message_times = {}
  18. @router.get("/orders/{order_id}/messages")
  19. async def get_order_messages(order_id: int, user: dict = Depends(get_current_user)):
  20. role = user.get("role")
  21. user_id = user.get("id")
  22. # Fetch user chat status
  23. user_info = db.execute_query("SELECT can_chat FROM users WHERE id = %s", (user_id,))
  24. can_chat = user_info[0]['can_chat'] if user_info else False
  25. order = db.execute_query("SELECT user_id FROM orders WHERE id = %s", (order_id,))
  26. if not order: raise HTTPException(status_code=404, detail="Order not found")
  27. if role != 'admin':
  28. if order[0]['user_id'] != user_id: raise HTTPException(status_code=403, detail="Not authorized")
  29. if not can_chat: raise HTTPException(status_code=403, detail="Chat access disabled for your account")
  30. query = """
  31. SELECT m.id, m.is_from_admin, m.message, m.created_at, u.first_name, u.email
  32. FROM order_messages m
  33. LEFT JOIN users u ON m.user_id = u.id
  34. WHERE m.order_id = %s ORDER BY m.created_at ASC
  35. """
  36. messages = db.execute_query(query, (order_id,))
  37. for msg in messages:
  38. if msg.get('created_at'): msg['created_at'] = msg['created_at'].isoformat()
  39. msg['is_from_admin'] = bool(msg['is_from_admin'])
  40. # Mark messages as read
  41. if role == 'admin':
  42. 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,))
  43. await global_manager.notify_admins()
  44. await global_manager.notify_order_read(order_id)
  45. else:
  46. 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,))
  47. await global_manager.notify_user(user_id)
  48. return messages
  49. @router.post("/orders/{order_id}/messages")
  50. async def post_order_message(
  51. order_id: int,
  52. message: Optional[str] = Form(None),
  53. file: Optional[UploadFile] = File(None),
  54. user: dict = Depends(get_current_user),
  55. lang: str = "en"
  56. ):
  57. role = user.get("role")
  58. user_id = user.get("id")
  59. # Flood control for non-admin users
  60. if role != 'admin':
  61. now = datetime.datetime.utcnow().timestamp()
  62. last_time = last_message_times.get(user_id, 0)
  63. if now - last_time < 10:
  64. raise HTTPException(status_code=429, detail=locales.translate_error("flood_control", lang))
  65. last_message_times[user_id] = now
  66. if not message and not file:
  67. raise HTTPException(status_code=400, detail="Empty message")
  68. if message:
  69. message = message.strip()
  70. is_admin = (role == 'admin')
  71. if not is_admin:
  72. user_info = db.execute_query("SELECT can_chat FROM users WHERE id = %s", (user_id,))
  73. if not user_info or not user_info[0]['can_chat']:
  74. raise HTTPException(status_code=403, detail="Chat access disabled")
  75. order = db.execute_query("SELECT user_id FROM orders WHERE id = %s", (order_id,))
  76. if not order: raise HTTPException(status_code=404, detail="Order not found")
  77. if not is_admin and order[0]['user_id'] != user_id: raise HTTPException(status_code=403, detail="Not authorized")
  78. final_message = ""
  79. db_file_path = None
  80. if file and file.filename:
  81. file_ext = os.path.splitext(file.filename)[1]
  82. unique_filename = f"{uuid.uuid4()}{file_ext}"
  83. file_path = os.path.join(config.CHAT_UPLOADS_DIR, unique_filename)
  84. db_file_path = f"uploads/chats/{unique_filename}"
  85. with open(file_path, "wb") as buffer:
  86. while chunk := file.file.read(8192):
  87. buffer.write(chunk)
  88. msg_obj = {}
  89. if db_file_path:
  90. msg_obj["image_url"] = db_file_path
  91. if message:
  92. msg_obj["text"] = message
  93. final_message = json.dumps(msg_obj)
  94. query = "INSERT INTO order_messages (order_id, user_id, is_from_admin, message) VALUES (%s, %s, %s, %s)"
  95. msg_id = db.execute_commit(query, (order_id, user_id, is_admin, final_message))
  96. now = datetime.datetime.utcnow().isoformat()
  97. user_info = db.execute_query("SELECT first_name, email FROM users WHERE id = %s", (user_id,))
  98. first_name = user_info[0]['first_name'] if user_info else None
  99. email = user_info[0]['email'] if user_info else None
  100. await manager.broadcast_to_order(order_id, {
  101. "id": msg_id,
  102. "is_from_admin": is_admin,
  103. "message": final_message,
  104. "created_at": now,
  105. "first_name": first_name,
  106. "email": email
  107. })
  108. if is_admin:
  109. await global_manager.notify_user(order[0]['user_id'])
  110. else:
  111. from services import event_hooks
  112. await global_manager.notify_admins()
  113. await global_manager.notify_admins_new_message(order_id, message or "Image uploaded")
  114. event_hooks.on_message_received(order_id, user_id, message or "Image uploaded")
  115. return {
  116. "id": msg_id,
  117. "status": "sent",
  118. "message": final_message,
  119. "first_name": first_name,
  120. "email": email
  121. }