pit-router/app/__init__.py

"""
PIT Router Flask 应用工厂
"""
from flask import Flask
from flask_socketio import SocketIO
import logging

from app.config import config
from app.extensions import (
    db, migrate, jwt, login_manager, cors, limiter, init_redis
)

# Socket.IO 实例 - 避免与 app/socketio 目录冲突
socketio_app = SocketIO(cors_allowed_origins="*", async_mode='threading')


def create_app(config_name='default'):
    """创建 Flask 应用实例"""
    app = Flask(__name__)
    app.config.from_object(config[config_name])
    
    # 初始化扩展
    _init_extensions(app)
    
    # 注册蓝图
    _register_blueprints(app)
    
    # 注册 Socket.IO 事件
    _register_socketio_events()
    
    # 配置日志
    _configure_logging(app)
    
    # 注册错误处理
    _register_error_handlers(app)
    
    # 根路由
    @app.route('/')
    def index():
        return {
            'service': '智队中枢',
            'version': '0.6.1',
            'status': 'running',
            'endpoints': {
                'health': '/health',
                'auth': '/api/auth',
                'sessions': '/api/sessions',
                'agents': '/api/agents',
                'gateways': '/api/gateways',
                'messages': '/api/messages',
                'stats': '/api/stats',
                'websocket': 'ws://host:9000'
            }
        }
    
    # 健康检查端点
    @app.route('/health')
    def health_check():
        return {'status': 'ok', 'service': 'pit-router'}
    
    return app


def _init_extensions(app):
    """初始化 Flask 扩展"""
    db.init_app(app)
    migrate.init_app(app, db)
    jwt.init_app(app)
    login_manager.init_app(app)
    cors.init_app(app)
    limiter.init_app(app)
    socketio_app.init_app(app)
    
    # 配置 Flask-Login user_loader
    from app.models import User
    @login_manager.user_loader
    def load_user(user_id):
        return User.query.get(user_id)
    
    # 初始化 Redis
    init_redis(app)
    
    # 创建数据库表
    with app.app_context():
        db.create_all()


def _register_blueprints(app):
    """注册蓝图"""
    from app.routes.auth import auth_bp
    from app.routes.sessions import sessions_bp
    from app.routes.agents import agents_bp
    from app.routes.gateways import gateways_bp
    from app.routes.messages import messages_bp
    from app.routes.stats import stats_bp
    from app.web.routes import web_bp
    from app.web.api import web_api_bp
    
    app.register_blueprint(auth_bp, url_prefix='/api/auth')
    app.register_blueprint(sessions_bp, url_prefix='/api/sessions')
    app.register_blueprint(agents_bp, url_prefix='/api/agents')
    app.register_blueprint(gateways_bp, url_prefix='/api/gateways')
    app.register_blueprint(messages_bp, url_prefix='/api/messages')
    app.register_blueprint(stats_bp, url_prefix='/api/stats')
    app.register_blueprint(web_bp, url_prefix='/web')
    app.register_blueprint(web_api_bp)


def _register_socketio_events():
    """注册 Socket.IO 事件处理器"""
    from app.socketio import handlers as socketio_handlers
    socketio_handlers.register_handlers(socketio_app)


def _configure_logging(app):
    """配置日志"""
    log_level = app.config.get('LOG_LEVEL', 'INFO')
    logging.basicConfig(
        level=getattr(logging, log_level),
        format='%(asctime)s [%(levelname)s] %(name)s: %(message)s'
    )


def _register_error_handlers(app):
    """注册错误处理器"""
    from flask import jsonify, render_template, request, render_template
    
    @app.errorhandler(400)
    def bad_request(error):
        if request.accepts('text/html'):
            return render_template('errors/400.html', error=str(error)), 400
        return jsonify({'error': 'Bad Request', 'message': str(error)}), 400
    
    @app.errorhandler(401)
    def unauthorized(error):
        if request.accepts('text/html'):
            return render_template('errors/401.html', error=str(error)), 401
        return jsonify({'error': 'Unauthorized', 'message': str(error)}), 401
    
    @app.errorhandler(403)
    def forbidden(error):
        if request.accepts('text/html'):
            return render_template('errors/403.html', error=str(error)), 403
        return jsonify({'error': 'Forbidden', 'message': str(error)}), 403
    
    @app.errorhandler(404)
    def not_found(error):
        if request.accepts('text/html'):
            return render_template('errors/404.html', error=str(error)), 404
        return jsonify({'error': 'Not Found', 'message': str(error)}), 404
    
    @app.errorhandler(500)
    def internal_error(error):
        if request.accepts('text/html'):
            return render_template('errors/500.html', error=str(error)), 500
        return jsonify({'error': 'Internal Server Error', 'message': str(error)}), 500