api_server.py

"""
霜云(Shimokumo) - API服务器模块

基于Flask的REST API服务器，提供以下端点：
- POST /chat     : 聊天对话（支持流式）
- POST /search   : 联网搜索
- POST /browse   : 网页浏览
- POST /novel    : 小说创作管理
- POST /video    : 视频制作管理
- POST /media    : 多媒体播放管理
- GET  /health   : 健康检查
- GET  /info     : 模型信息

支持流式响应(SSE)、API Key认证和CORS。
"""

import json
import time
from functools import wraps
from typing import Any, Dict, Optional

from flask import Flask, Response, jsonify, request, stream_with_context

from configuration_shimokumo import ShimokumoConfig
from shimokumo_utils.logger import get_logger

logger = get_logger("Shimokumo.Server")


def create_api_key_auth(api_key: str):
    """创建API Key认证装饰器工厂"""
    def decorator(f):
        @wraps(f)
        def decorated_function(*args, **kwargs):
            if not api_key:
                return f(*args, **kwargs)

            # 从header或query参数获取API Key
            provided_key = request.headers.get("X-API-Key") or request.args.get("api_key")
            if not provided_key:
                return jsonify({
                    "error": "未提供API Key",
                    "message": "请在请求头中添加 X-API-Key 或在URL参数中添加 api_key",
                }), 401

            if provided_key != api_key:
                return jsonify({
                    "error": "API Key无效",
                    "message": "请检查您的API Key是否正确",
                }), 403

            return f(*args, **kwargs)
        return decorated_function
    return decorator


class ShimokumoServer:
    """霜云API服务器

    封装Flask应用，提供REST API接口。

    用法：
        server = ShimokumoServer(config, inference, ...)
        server.run()

    或使用工厂方法：
        app = create_app(config, inference)
        app.run(host="0.0.0.0", port=8080)
    """

    def __init__(
        self,
        config: ShimokumoConfig,
        inference_engine=None,
        web_search_module=None,
        web_browser_module=None,
        media_player_module=None,
        novel_creator_module=None,
        video_generator_module=None,
        character_chat_module=None,
    ):
        """
        初始化API服务器。

        Args:
            config: 模型配置
            inference_engine: 推理引擎
            web_search_module: 联网搜索模块
            web_browser_module: 网页浏览模块
            media_player_module: 多媒体播放模块
            novel_creator_module: 小说创作模块
            video_generator_module: 视频制作模块
            character_chat_module: 角色聊天模块
        """
        self.config = config
        self.inference = inference_engine
        self.web_search = web_search_module
        self.web_browser = web_browser_module
        self.media_player = media_player_module
        self.novel_creator = novel_creator_module
        self.video_generator = video_generator_module
        self.character_chat = character_chat_module

        self.app = self._create_flask_app()
        self.start_time = time.time()

    def _create_flask_app(self) -> Flask:
        """创建并配置Flask应用"""
        app = Flask(
            __name__,
            static_folder=None,
            template_folder=None,
        )

        # CORS支持
        self._setup_cors(app)

        # API Key认证装饰器
        auth = create_api_key_auth(self.config.api_key)

        # ==================== 健康检查 ====================

        @app.route("/health", methods=["GET"])
        def health_check():
            """健康检查端点"""
            uptime = time.time() - self.start_time
            return jsonify({
                "status": "healthy",
                "service": "Shimokumo",
                "uptime": uptime,
                "version": "1.0.0",
            })

        @app.route("/info", methods=["GET"])
        def model_info():
            """模型信息端点"""
            info = {
                "name": self.config.character_name,
                "name_en": self.config.character_name_en,
                "title": self.config.character_title,
                "model_config": {
                    "hidden_size": self.config.hidden_size,
                    "num_layers": self.config.num_layers,
                    "num_heads": self.config.num_heads,
                    "num_kv_heads": self.config.num_kv_heads,
                    "vocab_size": self.config.vocab_size,
                    "max_seq_len": self.config.max_seq_len,
                },
                "capabilities": {
                    "web_search": self.config.enable_web_search,
                    "novel_creation": self.config.enable_novel_creation,
                    "video_generation": self.config.enable_video_generation,
                    "character_chat": self.config.enable_character_chat,
                    "media_player": self.config.enable_media_player,
                    "web_browser": self.config.enable_web_browser,
                },
            }

            if self.inference:
                model = self.inference.model
                info["model_params"] = f"{model.get_num_params() / 1e9:.2f}B"
                info["model_size_mb"] = f"{model.get_model_size_mb():.1f}MB"

            return jsonify(info)

        # ==================== 聊天接口 ====================

        @app.route("/chat", methods=["POST"])
        @auth
        def chat():
            """聊天对话接口

            请求体:
                message (str): 用户消息
                stream (bool): 是否流式输出
                user_id (str): 用户ID（多轮对话）
                max_tokens (int): 最大生成token数
                temperature (float): 采样温度
                top_p (float): Top-P采样
            """
            data = request.get_json(silent=True) or {}
            message = data.get("message", "")
            is_stream = data.get("stream", False)
            user_id = data.get("user_id", "default")
            max_tokens = data.get("max_tokens", self.config.max_new_tokens)
            temperature = data.get("temperature", self.config.temperature)
            top_p = data.get("top_p", self.config.top_p)

            if not message:
                return jsonify({"error": "message不能为空"}), 400

            # 角色聊天模式
            if data.get("roleplay", False) and self.character_chat:
                response = self.character_chat.chat(
                    message,
                    user_id=user_id,
                    max_new_tokens=max_tokens,
                )
                return jsonify(response.to_dict())

            # 普通聊天
            if not self.inference:
                return jsonify({
                    "response": f"（{self.config.character_name}还没有加载模型的说...请稍后再试）",
                    "error": "模型未加载",
                }), 503

            if is_stream:
                return Response(
                    stream_with_context(self._stream_chat(
                        message, user_id, max_tokens, temperature, top_p
                    )),
                    mimetype="text/event-stream",
                    headers={
                        "Cache-Control": "no-cache",
                        "X-Accel-Buffering": "no",
                    },
                )
            else:
                response = self.inference.chat_reply(
                    user_message=message,
                    max_new_tokens=max_tokens,
                    temperature=temperature,
                    top_p=top_p,
                )
                return jsonify({
                    "response": response,
                    "user_id": user_id,
                    "model": self.config.character_name,
                })

        # ==================== 搜索接口 ====================

        @app.route("/search", methods=["POST"])
        @auth
        def search():
            """联网搜索接口

            请求体:
                query (str): 搜索关键词
                engines (list): 搜索引擎列表
                max_results (int): 最大结果数
            """
            if not self.config.enable_web_search:
                return jsonify({"error": "搜索功能未启用"}), 403

            data = request.get_json(silent=True) or {}
            query = data.get("query", "")
            engines = data.get("engines")
            max_results = data.get("max_results", 10)

            if not query:
                return jsonify({"error": "query不能为空"}), 400

            results = self.web_search.search(
                query,
                engines=engines,
                max_total_results=max_results,
            )

            return jsonify({
                "query": query,
                "results": [r.to_dict() for r in results],
                "total": len(results),
            })

        # ==================== 浏览接口 ====================

        @app.route("/browse", methods=["POST"])
        @auth
        def browse():
            """网页浏览接口

            请求体:
                url (str): 网页URL
                extract_summary (bool): 是否提取摘要
            """
            if not self.config.enable_web_browser:
                return jsonify({"error": "浏览功能未启用"}), 403

            data = request.get_json(silent=True) or {}
            url = data.get("url", "")
            extract_summary = data.get("extract_summary", True)

            if not url:
                return jsonify({"error": "url不能为空"}), 400

            page = self.web_browser.fetch(url)

            if not page:
                return jsonify({"error": f"无法获取网页: {url}"}), 404

            result = page.to_dict()
            if extract_summary:
                result["summary"] = self.web_browser.generate_summary(page)

            return jsonify(result)

        # ==================== 小说接口 ====================

        @app.route("/novel", methods=["POST"])
        @auth
        def novel():
            """小说创作接口

            请求体:
                action (str): 操作类型 (create_project/set_world/add_character/
                                   generate_outline/generate_chapter/check/save/load)
                ... (各操作的具体参数)
            """
            if not self.config.enable_novel_creation:
                return jsonify({"error": "小说功能未启用"}), 403

            data = request.get_json(silent=True) or {}
            action = data.get("action", "")

            if not self.novel_creator:
                return jsonify({"error": "小说模块未初始化"}), 503

            try:
                result = self._handle_novel_action(data)
                return jsonify(result)
            except Exception as e:
                logger.error(f"小说操作失败: {e}")
                return jsonify({"error": str(e)}), 500

        # ==================== 视频接口 ====================

        @app.route("/video", methods=["POST"])
        @auth
        def video():
            """视频制作接口

            请求体:
                action (str): 操作类型 (create/generate_storyboard/generate_subtitles/export_srt)
                ... (各操作的具体参数)
            """
            if not self.config.enable_video_generation:
                return jsonify({"error": "视频功能未启用"}), 403

            data = request.get_json(silent=True) or {}
            action = data.get("action", "")

            if not self.video_generator:
                return jsonify({"error": "视频模块未初始化"}), 503

            try:
                result = self._handle_video_action(data)
                return jsonify(result)
            except Exception as e:
                logger.error(f"视频操作失败: {e}")
                return jsonify({"error": str(e)}), 500

        # ==================== 媒体接口 ====================

        @app.route("/media", methods=["POST"])
        @auth
        def media():
            """多媒体播放接口

            请求体:
                action (str): 操作类型 (parse/add_playlist/get_playlist/resolve_stream)
                url (str): 媒体URL
            """
            if not self.config.enable_media_player:
                return jsonify({"error": "媒体功能未启用"}), 403

            data = request.get_json(silent=True) or {}
            action = data.get("action", "")
            url = data.get("url", "")

            if not self.media_player:
                return jsonify({"error": "媒体模块未初始化"}), 503

            if action == "parse" and url:
                item = self.media_player.parse_url(url)
                return jsonify(item.to_dict())
            elif action == "add_playlist" and url:
                item = self.media_player.add_to_playlist(url, data.get("title", ""))
                return jsonify(item.to_dict())
            elif action == "get_playlist":
                return jsonify({
                    "name": self.media_player.playlist.name,
                    "items": len(self.media_player.playlist),
                    "info": self.media_player.get_playlist_info(),
                })
            elif action == "resolve_stream" and url:
                stream_url = self.media_player.resolve_stream_url(url)
                return jsonify({"stream_url": stream_url})
            else:
                return jsonify({"error": "未知的操作类型或缺少参数"}), 400

        # ==================== 错误处理 ====================

        @app.errorhandler(404)
        def not_found(e):
            return jsonify({
                "error": "未找到请求的资源",
                "message": f"请检查您的请求路径。可用端点: /chat, /search, /browse, /novel, /video, /media, /health, /info",
            }), 404

        @app.errorhandler(500)
        def server_error(e):
            return jsonify({
                "error": "服务器内部错误",
                "message": str(e),
            }), 500

        return app

    def _setup_cors(self, app: Flask) -> None:
        """配置CORS跨域支持"""
        @app.after_request
        def add_cors_headers(response):
            response.headers.add("Access-Control-Allow-Origin", "*")
            response.headers.add("Access-Control-Allow-Headers", "Content-Type, X-API-Key")
            response.headers.add("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
            return response

        # 处理OPTIONS预检请求
        @app.before_request
        def handle_options():
            if request.method == "OPTIONS":
                response = app.make_default_options_response()
                response.headers.add("Access-Control-Allow-Origin", "*")
                response.headers.add("Access-Control-Allow-Headers", "Content-Type, X-API-Key")
                response.headers.add("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
                return response

    def _stream_chat(
        self,
        message: str,
        user_id: str,
        max_tokens: int,
        temperature: float,
        top_p: float,
    ):
        """SSE流式聊天生成器"""
        try:
            for token_text in self.inference.chat_stream(
                user_message=message,
                max_new_tokens=max_tokens,
                temperature=temperature,
                top_p=top_p,
            ):
                # SSE格式: data: {json}\n\n
                chunk = json.dumps({"token": token_text}, ensure_ascii=False)
                yield f"data: {chunk}\n\n"

            # 发送结束标记
            yield "data: [DONE]\n\n"
        except Exception as e:
            error_chunk = json.dumps({"error": str(e)}, ensure_ascii=False)
            yield f"data: {error_chunk}\n\n"

    def _handle_novel_action(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """处理小说创作操作"""
        action = data.get("action")
        creator = self.novel_creator

        if action == "create_project":
            project = creator.create_project(
                title=data.get("title", "未命名"),
                genre=data.get("genre", "奇幻"),
            )
            return {"success": True, "project_id": project.id, "title": project.title}

        elif action == "set_world":
            project_id = data.get("project_id")
            # 查找项目
            project = None
            for p in creator.projects.values():
                if p.id == project_id:
                    project = p
                    break
            if not project:
                return {"error": "项目不存在"}
            creator.set_world(project, **data)
            return {"success": True}

        elif action == "generate_outline":
            return {"error": "需要提供project_id"}

        elif action == "save":
            return {"success": True}

        else:
            return {"error": f"未知的操作: {action}"}

    def _handle_video_action(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """处理视频制作操作"""
        action = data.get("action")
        generator = self.video_generator

        if action == "create":
            project = generator.create_project(
                title=data.get("title", "未命名视频"),
                width=data.get("width", 1920),
                height=data.get("height", 1080),
                fps=data.get("fps", 24),
            )
            return {"success": True, "project_id": project.id}

        elif action == "generate_subtitles":
            return {"error": "需要提供project_id"}

        elif action == "export_srt":
            return {"error": "需要提供project_id"}

        else:
            return {"error": f"未知的操作: {action}"}

    def run(self, host: Optional[str] = None, port: Optional[int] = None, debug: bool = False):
        """
        启动服务器。

        Args:
            host: 监听地址（覆盖配置）
            port: 监听端口（覆盖配置）
            debug: 调试模式
        """
        host = host or self.config.host
        port = port or self.config.port

        logger.info(f"霜云API服务器启动中...")
        logger.info(f"地址: http://{host}:{port}")
        logger.info(f"API文档: http://{host}:{port}/info")
        logger.info(f"健康检查: http://{host}:{port}/health")

        self.app.run(
            host=host,
            port=port,
            debug=debug or self.config.debug,
            threaded=True,
        )