app.py

import os
from flask import Flask, request, session
from flask_cors import CORS

from utils.config import Config
from utils.extensions import ext
from routes.auth import auth_bp
from routes.conversations import conversations_bp
from routes.query import query_bp
from routes.admin import admin_bp
from routes.pages import pages_bp

# Khởi tạo Flask app
app = Flask(__name__)
CORS(app)

# Load config
app.config.from_object(Config)

# Đảm bảo thư mục session tồn tại
os.makedirs(app.config['SESSION_FILE_DIR'], exist_ok=True)

# Khởi tạo các extension
if not ext.initialize_all(app):
    print("Failed to initialize extensions. Exiting...")
    exit(1)

# Đăng ký blueprints
app.register_blueprint(auth_bp)
app.register_blueprint(conversations_bp)
app.register_blueprint(query_bp)
app.register_blueprint(admin_bp)
app.register_blueprint(pages_bp)

# WebSocket handlers
@ext.socketio.on('connect')
def handle_connect(auth=None):
    user_id = session.get('user_id')
    if user_id:
        ext.connected_clients[user_id] = request.sid
        try:
            from bson import ObjectId
            user = ext.db.users.find_one({'_id': ObjectId(user_id)})
            if user:
                ext.socketio.emit('query_update', {
                    'query_count': user.get('query_count', 0),
                    'query_limit': user.get('query_limit', 3 if user.get('account_type') == 'limited' else None)
                }, room=request.sid)
                print(f"User {user_id} connected via WebSocket with SID {request.sid}")
        except Exception as e:
            print(f"Error in WebSocket connect for user {user_id}: {e}")
    else:
        ext.socketio.disconnect(request.sid)
        print("Unauthorized WebSocket connection attempt")

@ext.socketio.on('disconnect')
def handle_disconnect():
    user_id = session.get('user_id')
    if user_id in ext.connected_clients and ext.connected_clients[user_id] == request.sid:
        del ext.connected_clients[user_id]
        print(f"User {user_id} disconnected from WebSocket")

if __name__ == '__main__':
    ext.socketio.run(
        app,
        host='0.0.0.0',
        port=int(os.getenv('PORT', 7860)),
        debug=True,
        allow_unsafe_werkzeug=True
    )