"""
异步任务管理 MCP 工具

版权所有 (C) 吉林省左右软件开发有限公司
Copyright (C) Equilibrium Software Development Co., Ltd, Jilin
Update & Mod By Crystalxp (黑夜杀手 QQ:281309196)

提供异步任务的启动、查询和管理功能
"""

from typing import Any, Optional
from mcp.types import CallToolResult

from ..constants import DIR_DELPHI_KB, REG_KEY_EMBARCADERO_BDS
from ..services.knowledge_base.async_task_manager import get_task_manager, TaskStatus
from ..utils.logger import get_logger

logger = get_logger(__name__)


async def start_async_task(arguments: Any) -> CallToolResult:
    """
    启动异步任务

    Args:
        arguments: 包含以下参数:
             - task_type: 任务类型 (必需)
                - "build_knowledge_base": 构建Delphi知识库
                - "build_thirdparty_knowledge_base": 构建第三方库知识库
                - "init_project_knowledge_base": 初始化项目知识库
                - "build_document_knowledge_base": 构建文档知识库
                - "build_example_knowledge_base": 构建示例知识库
            - params: 任务参数 (可选, 根据任务类型不同)
                - build_document_knowledge_base:
                    - urls: 网页URL列表
                    - directory: 扫描目录路径
                    - extensions: 文件扩展名列表
            - show_progress: 是否显示进度 (可选, 默认 true)

    Returns:
        任务启动结果，包含任务ID
    """
    task_type = arguments.get("task_type")
    params = arguments.get("task_params", {})
    show_progress = arguments.get("show_progress", True)

    if not task_type:
        return CallToolResult(
            content=[{"type": "text", "text": "请提供任务类型"}],
            isError=True
        )

    task_manager = get_task_manager()

    # 根据任务类型创建任务函数
    if task_type == "build_knowledge_base":
        from ..services.knowledge_base.zvec_adapter import ZVecKnowledgeBaseAdapter
        from src.utils.delphi_env import get_delphi_version
        import winreg
        from pathlib import Path

        def _get_delphi_source_dirs():
            dirs = []
            try:
                key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, REG_KEY_EMBARCADERO_BDS)
                i = 0
                while True:
                    try:
                        vkey = winreg.EnumKey(key, i)
                        vpath = winreg.OpenKey(key, vkey)
                        try:
                            root = winreg.QueryValueEx(vpath, "RootDir")[0]
                            src = Path(root) / "source"
                            if src.exists():
                                dirs.append(str(src))
                        except OSError:
                            pass
                        finally:
                            winreg.CloseKey(vpath)
                        i += 1
                    except WindowsError:
                        break
                winreg.CloseKey(key)
            except Exception:
                pass
            return dirs

        def build_kb_task(**kwargs):
            rebuild = kwargs.get("rebuild", False)
            progress_callback = kwargs.get("_progress_callback")
            source_dirs = _get_delphi_source_dirs()
            kb_dir = str(Path(__file__).parent.parent.parent / "data" / "delphi-knowledge-base")
            service = ZVecKnowledgeBaseAdapter(kb_dir, source_dirs=source_dirs)
            service.progress_callback = progress_callback
            return service.build_knowledge_base(rebuild=rebuild)

        # 确定显示的版本号
        display_version = params.get("version")
        if not display_version:
            display_version = get_delphi_version()
        if not display_version:
            display_version = "最新"
        task_name = f"构建Delphi知识库 (版本: {display_version})"

    elif task_type == "build_thirdparty_knowledge_base":
        from ..services.knowledge_base.thirdparty_knowledge_base import ThirdPartyKnowledgeBase

        def build_thirdparty_task(**kwargs):
            version = kwargs.get("version")
            rebuild = kwargs.get("rebuild", False)
            progress_callback = kwargs.get("_progress_callback")

            service = ThirdPartyKnowledgeBase(progress_callback=progress_callback)
            return service.build_thirdparty_knowledge_base(version=version, rebuild=rebuild)

        task_name = "构建第三方库知识库"

    elif task_type == "init_project_knowledge_base":
        from ..services.knowledge_base.project_knowledge_base import ProjectKnowledgeBase
        import shutil
        import time as _time_mod

        def init_project_task(**kwargs):
            project_path = kwargs.get("project_path")
            rebuild = kwargs.get("rebuild", False)
            progress_callback = kwargs.get("_progress_callback")

            project_kb = ProjectKnowledgeBase(project_path, progress_callback)

            # ── 热切换：重建时构建到临时目录，旧 KB 仍可查 ──
            old_kb_dir = project_kb.kb_dir
            has_old_kb = old_kb_dir.exists() and any(old_kb_dir.iterdir())
            use_temp_build = rebuild and has_old_kb
            temp_dir = None

            if use_temp_build:
                # 用时间戳创建临时目录名，避免冲突
                temp_suffix = f"-tmp-{int(_time_mod.time())}"
                temp_dir = old_kb_dir.parent / f"{DIR_DELPHI_KB}{temp_suffix}"
                # 临时用 kb_dir 指向 temp，构建到临时目录
                project_kb.kb_dir = temp_dir
                logger.info(
                    "重建项目 KB 到临时目录: %s （旧 KB 位于 %s，仍可查）",
                    temp_dir, old_kb_dir,
                )

            results = {}
            try:
                _t0 = _time_mod.monotonic()
                success = project_kb.build_project_knowledge_base(
                    rebuild=rebuild)
                _elapsed = _time_mod.monotonic() - _t0
                logger.info(
                    "项目 KB 构建%s: %s (%.1f 秒, %s)",
                    "完成" if success else "失败",
                    project_path,
                    _elapsed,
                    "全量" if rebuild else "增量",
                )
                results["project"] = success

                # 构建成功 + 使用了临时目录 → 执行热切换
                if use_temp_build and success:
                    project_kb.close()
                    # 删除旧 KB
                    shutil.rmtree(old_kb_dir, ignore_errors=True)
                    # 临时目录 → 正式目录
                    shutil.move(str(temp_dir), str(old_kb_dir))
                    # 清缓存，下次搜索走新实例
                    from ..tools.knowledge_base import _clear_pkb_cache
                    _clear_pkb_cache(project_path)
                    logger.info("项目 KB 热切换完成: %s → %s", temp_dir, old_kb_dir)
                elif use_temp_build and not success:
                    # 构建失败：删除临时目录，旧 KB 保持不变
                    shutil.rmtree(temp_dir, ignore_errors=True)
                    project_kb.kb_dir = old_kb_dir  # 恢复指针
                    logger.warning("项目 KB 重建失败，已保留旧 KB: %s", old_kb_dir)

            except Exception:
                # 异常时也清理临时目录
                if use_temp_build and temp_dir and temp_dir.exists():
                    shutil.rmtree(temp_dir, ignore_errors=True)
                    project_kb.kb_dir = old_kb_dir
                raise

            stats = project_kb.get_statistics()
            results["statistics"] = stats

            return results

        task_name = f"初始化项目知识库 ({params.get('project_path', '未知项目')})"

    elif task_type == "build_embedding":
        from ..services.knowledge_base.project_knowledge_base import ProjectKnowledgeBase

        def build_embedding_task(**kwargs):
            project_path = kwargs.get("project_path")
            progress_callback = kwargs.get("_progress_callback")
            pkb = ProjectKnowledgeBase(project_path, progress_callback)
            pkb.load_knowledge_bases()
            counts = pkb.build_vectors(progress_callback=progress_callback)
            pkb.close()
            return counts

        task_name = f"构建 embedding 向量 ({params.get('project_path', '')})"

    elif task_type == "build_example_knowledge_base":
        from ..services.knowledge_base.example_knowledge_base import (
            ExampleKnowledgeBase,
        )

        def build_example_task(**kwargs):
            rebuild = kwargs.get("rebuild", False)
            progress_callback = kwargs.get("_progress_callback")
            ekb = ExampleKnowledgeBase(progress_callback=progress_callback)
            return ekb.build_example_knowledge_base(rebuild=rebuild)

        task_name = "构建示例知识库"

    elif task_type == "build_document_knowledge_base":
        from ..services.knowledge_base.scan_generic_documents import GenericDocumentScanner
        from pathlib import Path as FilePath
        
        def build_doc_task(**kwargs):
            urls = kwargs.get("urls", [])
            directory = kwargs.get("directory")
            extensions = kwargs.get("extensions")
            start_url = kwargs.get("start_url")
            max_pages = kwargs.get("max_pages", 100)
            max_depth = kwargs.get("max_depth", 3)
            domain_filter = kwargs.get("domain_filter")
            url_pattern = kwargs.get("url_pattern")
            exclude = kwargs.get("exclude")
            rebuild = kwargs.get("rebuild", False)
            progress_callback = kwargs.get("_progress_callback")
            
            server_root = FilePath(__file__).parent.parent.parent
            kb_dir = str(server_root / "data" / "document-knowledge-base")
            FilePath(kb_dir).mkdir(parents=True, exist_ok=True)
            
            scanner = GenericDocumentScanner(kb_dir, progress_callback=progress_callback)
            
            results = {"success": 0, "failed": 0, "total": 0, "total_size": 0}
            
            # 爬取网站（自动嵌套采集）
            if start_url:
                crawl_result = scanner.crawl_website(
                    start_url=start_url,
                    max_pages=max_pages,
                    max_depth=max_depth,
                    domain_filter=domain_filter,
                    url_pattern=url_pattern
                )
                results.update(crawl_result)
            
            # 添加指定 URL 列表
            if urls:
                for url in urls:
                    result = scanner.add_web_document(url)
                    if result and not (isinstance(result, dict) and 'error' in result):
                        results["success"] += 1
                    else:
                        results["failed"] += 1
                    results["total"] += 1
            
            # 扫描目录
            if directory:
                scan_result = scanner.scan_directory(directory, extensions=extensions, exclude=exclude,
                                                     rebuild=rebuild)
                results["scan"] = scan_result
                # 扫描器已直接写入 ZVec，无需额外导出步骤
                if scan_result and not scan_result.get("error"):
                    results["zvec_built"] = True

            return results
        
        if params.get("start_url"):
            task_name = f"爬取网站 (最多 {params.get('max_pages', 100)} 页)"
        elif params.get("directory"):
            task_name = f"构建文档知识库 (扫描目录: {params.get('directory')})"
        else:
            urls_count = len(params.get("urls", []))
            task_name = f"构建文档知识库 ({urls_count} 个URL)"

    else:
        return CallToolResult(
            content=[{"type": "text", "text": f"未知的任务类型: {task_type}"}],
            isError=True
        )

    # 取出 MCP 推送回调（由 server.py 注入）
    on_complete = arguments.pop('_on_complete', None)

    try:
        # 创建进度回调
        def progress_callback(progress, message):
            if show_progress:
                logger.info(f"任务进度: {progress:.1f}% - {message}")

        # 计算去重键，防止同类型构建重复提交
        dedup_key = params.pop("dedup_key", None)
        if dedup_key is None:
            if task_type == "init_project_knowledge_base":
                dedup_key = f"project_rebuild:{params.get('project_path', '')}"
            elif task_type == "build_thirdparty_knowledge_base":
                dedup_key = "thirdparty_rebuild"
            elif task_type == "build_knowledge_base":
                dedup_key = "delphi_rebuild"
            elif task_type == "build_document_knowledge_base":
                dedup_key = f"document_rebuild:{params.get('directory', '')}"
            elif task_type == "build_embedding":
                dedup_key = f"embedding:{params.get('project_path', '')}"
            elif task_type == "build_example_knowledge_base":
                dedup_key = "example_rebuild"

        # 提交异步任务
        task_func = (
            build_kb_task if task_type == "build_knowledge_base" else
            build_thirdparty_task if task_type == "build_thirdparty_knowledge_base" else
            init_project_task if task_type == "init_project_knowledge_base" else
            build_embedding_task if task_type == "build_embedding" else
            build_example_task if task_type == "build_example_knowledge_base" else
            build_doc_task
        )
        
        task_id = task_manager.submit_task(
            name=task_name,
            func=task_func,
            progress_callback=progress_callback if show_progress else None,
            dedup_key=dedup_key,
            on_complete=on_complete,
            **params
        )

        result_text = f"✓ 异步任务已启动\n\n"
        result_text += f"任务ID: {task_id}\n"
        result_text += f"任务名称: {task_name}\n"
        result_text += f"状态: 已提交到后台执行\n\n"
        result_text += f"使用以下命令查询任务状态:\n"
        result_text += f"  - get_task_status: 查询任务状态\n"
        result_text += f"  - get_task_result: 获取任务结果\n"
        result_text += f"  - list_tasks: 列出所有任务"

        return CallToolResult(content=[{"type": "text", "text": result_text}])

    except Exception as e:
        logger.error(f"启动异步任务失败: {e}", exc_info=True)
        return CallToolResult(
            content=[{"type": "text", "text": f"启动异步任务失败: {str(e)}"}],
            isError=True
        )


async def get_task_status(arguments: Any) -> CallToolResult:
    """
    获取任务状态（支持长轮询）

    Args:
        arguments: 包含以下参数:
            - task_id: 任务ID (必需)
            - long_poll_seconds: 长轮询等待秒数（可选，默认0即立即返回。MCP请求通道有超时限制，建议≤30秒，超时改用短轮询）

    Returns:
        任务状态信息
    """
    task_id = arguments.get("task_id")
    if not task_id:
        return CallToolResult(
            content=[{"type": "text", "text": "请提供任务ID"}],
            isError=True
        )

    task_manager = get_task_manager()
    
    # 长轮询：等待进度变化再返回
    long_poll = arguments.get("long_poll_seconds", 0)
    if long_poll > 0:
        task_info = task_manager.wait_for_progress_change(task_id, long_poll)
    else:
        task_info = task_manager.get_task_info(task_id)

    if not task_info:
        return CallToolResult(
            content=[{"type": "text", "text": f"未找到任务: {task_id}"}],
            isError=True
        )

    result_text = f"任务状态: {task_id}\n\n"
    result_text += f"任务名称: {task_info.name}\n"
    result_text += f"状态: {task_info.status.value}\n"
    result_text += f"进度: {task_info.progress:.1f}%\n"
    result_text += f"消息: {task_info.message}\n"
    result_text += f"创建时间: {task_info.created_at.strftime('%Y-%m-%d %H:%M:%S')}\n"

    if task_info.started_at:
        result_text += f"开始时间: {task_info.started_at.strftime('%Y-%m-%d %H:%M:%S')}\n"

    if task_info.completed_at:
        result_text += f"完成时间: {task_info.completed_at.strftime('%Y-%m-%d %H:%M:%S')}\n"

    if task_info.error:
        result_text += f"错误: {task_info.error}\n"

    return CallToolResult(content=[{"type": "text", "text": result_text}])


async def get_task_result(arguments: Any) -> CallToolResult:
    """
    获取任务结果

    Args:
        arguments: 包含以下参数:
            - task_id: 任务ID (必需)

    Returns:
        任务结果
    """
    task_id = arguments.get("task_id")
    if not task_id:
        return CallToolResult(
            content=[{"type": "text", "text": "请提供任务ID"}],
            isError=True
        )

    task_manager = get_task_manager()
    task_info = task_manager.get_task_info(task_id)

    if not task_info:
        return CallToolResult(
            content=[{"type": "text", "text": f"未找到任务: {task_id}"}],
            isError=True
        )

    if task_info.status != TaskStatus.COMPLETED:
        result_text = f"任务尚未完成\n\n"
        result_text += f"任务ID: {task_id}\n"
        result_text += f"当前状态: {task_info.status.value}\n"
        result_text += f"进度: {task_info.progress:.1f}%\n"
        result_text += f"消息: {task_info.message}\n"

        return CallToolResult(content=[{"type": "text", "text": result_text}])

    if task_info.error:
        return CallToolResult(
            content=[{"type": "text", "text": f"任务执行失败: {task_info.error}"}],
            isError=True
        )

    # 格式化结果
    result_text = f"任务执行成功!\n\n"
    if isinstance(task_info.result, dict):
        for key, value in task_info.result.items():
            result_text += f"{key}: {value}\n"
    else:
        result_text += f"结果: {task_info.result}\n"

    return CallToolResult(content=[{"type": "text", "text": result_text}])


async def cancel_task(arguments: Any) -> CallToolResult:
    """
    取消任务

    Args:
        arguments: 包含以下参数:
            - task_id: 任务ID (必需)

    Returns:
        操作结果
    """
    task_id = arguments.get("task_id")
    if not task_id:
        return CallToolResult(
            content=[{"type": "text", "text": "请提供任务ID"}],
            isError=True
        )

    task_manager = get_task_manager()
    success = task_manager.cancel_task(task_id)

    if success:
        return CallToolResult(content=[{"type": "text", "text": f"任务 {task_id} 已取消"}])
    else:
        return CallToolResult(
            content=[{"type": "text", "text": f"无法取消任务 {task_id}（任务不存在或已完成）"}],
            isError=True
        )


async def list_tasks(arguments: Any) -> CallToolResult:
    """
    列出所有任务

    Args:
        arguments: 空参数

    Returns:
        所有任务列表
    """
    task_manager = get_task_manager()
    tasks = task_manager.get_all_tasks()

    if not tasks:
        return CallToolResult(content=[{"type": "text", "text": "当前没有任务"}])

    result_text = f"任务列表 (共 {len(tasks)} 个):\n\n"

    for task_id, task_info in tasks.items():
        result_text += f"任务ID: {task_id}\n"
        result_text += f"  名称: {task_info.name}\n"
        result_text += f"  状态: {task_info.status.value}\n"
        result_text += f"  进度: {task_info.progress:.1f}%\n"
        result_text += f"  消息: {task_info.message}\n"
        result_text += f"  创建时间: {task_info.created_at.strftime('%Y-%m-%d %H:%M:%S')}\n"
        result_text += "\n"

    return CallToolResult(content=[{"type": "text", "text": result_text}])
