src/main.py

"""
霜云(Shimokumo) - 主入口文件

AI全功能ISP网络运营商巫女模型的启动入口。

支持命令行参数：
    --host          服务监听地址 (默认: 0.0.0.0)
    --port          服务监听端口 (默认: 8080)
    --api-key       API认证密钥
    --model-path    模型权重文件路径
    --tokenizer-path 分词器文件路径
    --device        运行设备 (auto/cuda/cpu)
    --log-level     日志级别 (DEBUG/INFO/WARNING/ERROR)
    --log-file      日志文件路径
    --no-search     禁用搜索功能
    --no-novel      禁用小说功能
    --no-video      禁用视频功能
    --no-chat       禁用角色聊天功能
    --no-media      禁用媒体功能
    --no-browser    禁用浏览功能
    --debug         启用调试模式
    --interactive   启动交互式命令行模式（不启动服务器）
    --help          显示帮助信息

用法：
    # 启动服务器
    python main.py --port 8080

    # 交互模式
    python main.py --interactive

    # 加载模型并启动
    python main.py --model-path ./models/shimokumo.bin --tokenizer-path ./models/tokenizer.model
"""

import argparse
import os
import signal
import sys
import time

# 将src目录加入Python路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from config import ShimokumoConfig
from utils.logger import get_logger
from utils.helpers import format_duration

logger = get_logger("Shimokumo.Main")


def parse_args() -> argparse.Namespace:
    """解析命令行参数"""
    parser = argparse.ArgumentParser(
        description="霜云(Shimokumo) - AI全功能ISP网络运营商巫女模型",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            "示例:\n"
            "  python main.py --port 8080                # 启动API服务器\n"
            "  python main.py --interactive             # 启动交互式聊天\n"
            "  python main.py --model-path model.bin    # 加载模型权重\n"
            "  python main.py --debug --log-level DEBUG  # 调试模式\n"
        ),
    )

    # 服务配置
    parser.add_argument("--host", type=str, default="0.0.0.0", help="服务监听地址")
    parser.add_argument("--port", type=int, default=8080, help="服务监听端口")
    parser.add_argument("--api-key", type=str, default="", help="API认证密钥（为空则不需要认证）")

    # 模型配置
    parser.add_argument("--model-path", type=str, default="", help="模型权重文件路径")
    parser.add_argument("--tokenizer-path", type=str, default="", help="分词器模型文件路径")
    parser.add_argument("--device", type=str, default="auto", choices=["auto", "cuda", "cpu"], help="运行设备")

    # 日志配置
    parser.add_argument("--log-level", type=str, default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR"], help="日志级别")
    parser.add_argument("--log-file", type=str, default="", help="日志文件路径")

    # 功能开关
    parser.add_argument("--no-search", action="store_true", help="禁用搜索功能")
    parser.add_argument("--no-novel", action="store_true", help="禁用小说创作功能")
    parser.add_argument("--no-video", action="store_true", help="禁用视频制作功能")
    parser.add_argument("--no-chat", action="store_true", help="禁用角色聊天功能")
    parser.add_argument("--no-media", action="store_true", help="禁用媒体播放功能")
    parser.add_argument("--no-browser", action="store_true", help="禁用网页浏览功能")

    # 模式选择
    parser.add_argument("--debug", action="store_true", help="启用调试模式")
    parser.add_argument("--interactive", action="store_true", help="启动交互式命令行聊天模式")

    return parser.parse_args()


def load_config(args: argparse.Namespace) -> ShimokumoConfig:
    """从命令行参数创建配置"""
    config = ShimokokuConfig(
        host=args.host,
        port=args.port,
        api_key=args.api_key,
        device=args.device,
        model_path=args.model_path,
        tokenizer_path=args.tokenizer_path,
        log_level=args.log_level,
        debug=args.debug,
    )

    # 设置日志文件
    if args.log_file:
        config.log_file = args.log_file
    if args.debug and not config.log_file:
        config.log_file = "/data/user/work/shimokumo_source/logs/shimokumo_debug.log"

    # 功能开关
    if args.no_search:
        config.enable_web_search = False
    if args.no_novel:
        config.enable_novel_creation = False
    if args.no_video:
        config.enable_video_generation = False
    if args.no_chat:
        config.enable_character_chat = False
    if args.no_media:
        config.enable_media_player = False
    if args.no_browser:
        config.enable_web_browser = False

    return config


def load_model(config: ShimokumoConfig):
    """
    加载模型和分词器。

    Args:
        config: 模型配置

    Returns:
        (model, tokenizer) 或 (None, None)
    """
    try:
        import torch
        from model.shimokumo_model import ShimokumoModel
        from model.tokenizer import ShimokumoTokenizer

        # 加载分词器
        tokenizer_path = config.tokenizer_path or ""
        tokenizer = ShimokumoTokenizer(model_path=tokenizer_path if tokenizer_path else None)
        logger.info(f"分词器加载完成: {tokenizer}")

        # 加载模型
        if config.model_path and os.path.exists(config.model_path):
            logger.info(f"正在加载模型权重: {config.model_path}")
            model = ShimokumoModel.from_pretrained(config, config.model_path)
            logger.info(f"模型加载完成: {model.get_num_params() / 1e9:.2f}B 参数")
        else:
            logger.warning("未指定模型权重路径，将使用未训练的随机初始化模型")
            logger.info("如需加载预训练权重，请使用 --model-path 参数")
            model = ShimokumoModel(config)
            logger.info(f"已创建模型架构: {model.get_num_params() / 1e9:.2f}B 参数")

        return model, tokenizer

    except Exception as e:
        logger.error(f"模型加载失败: {e}")
        return None, None


def init_modules(config: ShimokokuConfig, inference=None):
    """
    初始化功能模块。

    Args:
        config: 模型配置
        inference: 推理引擎（可为None）

    Returns:
        模块字典
    """
    modules: Dict = {}

    # 联网搜索模块
    if config.enable_web_search:
        try:
            from modules.web_search import WebSearchModule
            modules["web_search"] = WebSearchModule()
            logger.info("联网搜索模块已加载")
        except Exception as e:
            logger.warning(f"搜索模块加载失败: {e}")
            config.enable_web_search = False

    # 网页浏览模块
    if config.enable_web_browser:
        try:
            from modules.web_browser import WebBrowserModule
            modules["web_browser"] = WebBrowserModule()
            logger.info("网页浏览模块已加载")
        except Exception as e:
            logger.warning(f"浏览模块加载失败: {e}")
            config.enable_web_browser = False

    # 多媒体播放模块
    if config.enable_media_player:
        try:
            from modules.media_player import MediaPlayerModule
            modules["media_player"] = MediaPlayerModule()
            logger.info("多媒体播放模块已加载")
        except Exception as e:
            logger.warning(f"媒体模块加载失败: {e}")
            config.enable_media_player = False

    # 小说创作模块
    if config.enable_novel_creation:
        try:
            from modules.novel_creator import NovelCreatorModule
            modules["novel_creator"] = NovelCreatorModule(inference_engine=inference)
            logger.info("小说创作模块已加载")
        except Exception as e:
            logger.warning(f"小说模块加载失败: {e}")
            config.enable_novel_creation = False

    # 视频制作模块
    if config.enable_video_generation:
        try:
            from modules.video_generator import VideoGeneratorModule
            modules["video_generator"] = VideoGeneratorModule(inference_engine=inference)
            logger.info("视频制作模块已加载")
        except Exception as e:
            logger.warning(f"视频模块加载失败: {e}")
            config.enable_video_generation = False

    # 角色聊天模块
    if config.enable_character_chat:
        try:
            from modules.character_chat import CharacterChatModule
            modules["character_chat"] = CharacterChatModule(inference_engine=inference)
            logger.info("角色聊天模块已加载")
        except Exception as e:
            logger.warning(f"角色聊天模块加载失败: {e}")
            config.enable_character_chat = False

    return modules


def interactive_mode(config: ShimokokuConfig, inference=None, modules: Dict = None):
    """
    交互式命令行聊天模式。

    Args:
        config: 模型配置
        inference: 推理引擎
        modules: 功能模块字典
    """
    print()
    print("=" * 50)
    print(f"  霜云(Shimokumo) - 交互式聊天模式")
    print(f"  {config.character_title}")
    print("=" * 50)
    print()
    print("命令:")
    print("  /chat <消息>    - 角色扮演聊天")
    print("  /search <关键词> - 联网搜索")
    print("  /browse <URL>   - 浏览网页")
    print("  /clear          - 清空对话历史")
    print("  /info           - 查看信息")
    print("  /quit           - 退出")
    print()
    print("直接输入消息即可开始普通聊天。")
    print()

    # 初始化角色聊天模块
    from modules.character_chat import CharacterChatModule
    chat_module = modules.get("character_chat") or CharacterChatModule(inference_engine=inference)

    while True:
        try:
            user_input = input("\n你: ").strip()
        except (KeyboardInterrupt, EOFError):
            print("\n\n再见~！")
            break

        if not user_input:
            continue

        # 命令处理
        if user_input.startswith("/"):
            parts = user_input.split(maxsplit=1)
            cmd = parts[0].lower()
            arg = parts[1] if len(parts) > 1 else ""

            if cmd == "/quit" or cmd == "/exit":
                print(f"\n{config.character_name}: 再见~期待下次见到主人的说...\n")
                break

            elif cmd == "/clear":
                if inference:
                    inference.clear_dialogue()
                chat_module.clear_history()
                print("对话历史已清空。")

            elif cmd == "/info":
                print(f"\n--- 霜云信息 ---")
                print(f"名称: {config.character_name} ({config.character_name_en})")
                print(f"头衔: {config.character_title}")
                print(f"设备: {inference.device if inference else 'N/A'}")
                if inference:
                    print(f"参数量: {inference.model.get_num_params() / 1e9:.2f}B")
                print(f"功能: {'、'.join(k.replace('_', '') for k, v in modules.items())}")
                print()

            elif cmd == "/search" and arg:
                if modules.get("web_search"):
                    results = modules["web_search"].search(arg)
                    print(f"\n{modules['web_search'].format_results(results)}")
                else:
                    print("搜索功能未启用。")

            elif cmd == "/browse" and arg:
                if modules.get("web_browser"):
                    page = modules["web_browser"].fetch(arg)
                    if page:
                        summary = modules["web_browser"].generate_summary(page)
                        print(f"\n标题: {page.title}")
                        print(f"摘要: {summary}")
                    else:
                        print("无法获取网页。")
                else:
                    print("浏览功能未启用。")

            elif cmd == "/chat" and arg:
                response = chat_module.chat(arg)
                print(f"\n{config.character_name}: {response.text}")

            else:
                print(f"未知命令: {cmd}。输入 /quit 退出。")
        else:
            # 普通聊天
            print(f"\n{config.character_name}: ", end="", flush=True)

            if inference:
                response = inference.chat_reply(user_input)
                print(response)
            else:
                # 回退：使用角色聊天模块
                response = chat_module.chat(user_input)
                print(response.text)


def main():
    """主入口函数"""
    start_time = time.time()

    # 打印启动横幅
    print()
    print("  ╔══════════════════════════════════════════════════╗")
    print("  ║                                                  ║")
    print("  ║        霜云 (Shimokumo)                          ║")
    print("  ║        AI全功能ISP网络运营商巫女模型              ║")
    print("  ║                                                  ║")
    print("  ╚══════════════════════════════════════════════════╝")
    print()

    # 解析参数
    args = parse_args()

    # 创建配置
    config = load_config(args)

    # 初始化日志
    logger = get_logger(
        "Shimokumo",
        log_level=config.log_level,
        log_file=config.log_file if config.log_file else None,
    )

    logger.info("霜云启动中...")
    logger.info(f"配置: host={config.host}, port={config.port}, device={config.device}")

    # 加载模型
    model, tokenizer = load_model(config)

    # 初始化推理引擎（如果模型加载成功）
    inference = None
    if model and tokenizer:
        try:
            from model.inference import ShimokumoInference
            inference = ShimokumoInference(config, model, tokenizer)
        except Exception as e:
            logger.error(f"推理引擎初始化失败: {e}")

    # 初始化功能模块
    modules = init_modules(config, inference)

    # 计算启动耗时
    init_duration = time.time() - start_time
    logger.info(f"初始化完成，耗时 {format_duration(init_duration)}")

    # 信号处理（优雅退出）
    def signal_handler(sig, frame):
        logger.info("收到退出信号，正在关闭...")
        print("\n\n霜云正在优雅退出...再见~")
        sys.exit(0)

    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    # 选择运行模式
    if args.interactive:
        # 交互式命令行模式
        interactive_mode(config, inference, modules)
    else:
        # 服务器模式
        try:
            from api.server import ShimokumoServer

            server = ShimokumoServer(
                config=config,
                inference_engine=inference,
                web_search_module=modules.get("web_search"),
                web_browser_module=modules.get("web_browser"),
                media_player_module=modules.get("media_player"),
                novel_creator_module=modules.get("novel_creator"),
                video_generator_module=modules.get("video_generator"),
                character_chat_module=modules.get("character_chat"),
            )

            server.run()
        except ImportError as e:
            logger.error(f"缺少依赖: {e}")
            logger.error("请安装Flask: pip install flask")
            print("错误: 需要安装Flask来运行服务器模式。")
            print("请运行: pip install flask")
            print("或使用 --interactive 参数进入交互模式。")
            sys.exit(1)


if __name__ == "__main__":
    main()