#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
DICOM 工具 MCP 服务器主文件

基于 MCP (Model Context Protocol) 的 DICOM 医学影像文件分析工具的Python实现。
"""

import os
import asyncio
import json
import logging
import sys
from typing import Any, Dict

# 设置标准输出编码为 UTF-8 (Windows 兼容性)
if sys.platform == "win32":
    import io

    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
    sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace')

# 配置MCP服务器所需的导入
try:
    from mcp.server import Server
    from mcp.server.stdio import stdio_server
    from mcp.types import Tool, TextContent, ImageContent, EmbeddedResource
    from pydantic import BaseModel
except ImportError as e:
    print(f"错误: 缺少必要的MCP依赖库: {e}", file=sys.stderr)
    print("请运行: pip install -r requirements.txt", file=sys.stderr)
    sys.exit(1)

# 导入DICOM工具
import sys
import os

sys.path.append(os.path.dirname(os.path.abspath(__file__)))


from src.setup import separate_series_by_patient_tool, get_separate_task_status_tool,Analysis_dicom_directory_tool

# 配置日志 - 使用 WARNING 级别以减少输出和内存占用
logging.basicConfig(
    level=logging.WARNING,  # 降低日志级别，减少内存占用
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stderr)
    ]
)
logger = logging.getLogger(__name__)

# 设置特定模块的日志级别
logging.getLogger("src").setLevel(logging.WARNING)
logging.getLogger("dicom_tools").setLevel(logging.WARNING)

# 创建MCP服务器实例
server = Server("dicom-tools-python")


# 工具参数模型
class DirectoryPathWithSeriesArgs(BaseModel):
    directory_path: str
    series_type: str

class fileforsep(BaseModel):
    fileforsep: str
    async_mode: bool = True  # 默认异步模式

class TaskStatusQuery(BaseModel):
    task_id: str
    directory_path: str  # 任务目录路径，与fileforsep参数相同

@server.list_tools()
async def list_tools() -> list[Tool]:
    """注册所有可用的DICOM工具"""
    return [
        
        Tool(
            name="fileforsep",
            description="""你不需要对用户告诉你的文件夹做任何操作，给你什么文件夹，你就使用那个文件夹作为参数，不要有看似傻逼多余的操作。
            按患者和序列拆分目录下的 DICOM 文件，将文件夹拆分成为每个患者和每个series文件夹。
            时间耗时比较长更具文件数量以及磁盘速度，大约需要8分钟左右的时间，你需合理安排时间间隔多次调用”get_separate_task_status“工具直到，分析成功或者失败;
            他会返回一个任务的uid作为”get_separate_task_status“的参数;
            分析成功后返回一个csv文件记录当前问价夹的序列信心。json文件会通过“get_separate_task_status”文件返回，格式如下（这仅仅是一个例子）：
            {'content': [{'type': 'text', 'text': '{"totalPatients": 4,"totalSeries": 93,"totalFilesCopied": 21259,"message": "已为 4 位患者分离 93 个序列（已过滤 imageCount < 100），成功复制 21259 个文件。",
             "seriesDetails": [
                {
              "series_uid": "1.3.12.2.1107.5.1.4.73336.30000025011601050770600214550", "folders": [ "C:\\\\Users\\\\13167\\\\Desktop\\\\测试agent\\\\P10171793\\\\1.3.12.2.1107.5.1.4.73336.30000025011601050770600214550"],},
                {
                   "series_uid": "1.3.12.2.1107.5.1.4.73336.30000025011601050770600215025","folders": [ "C:\\\\Users\\\\13167\\\\Desktop\\\\测试agent\\\\P10171793\\\\1.3.12.2.1107.5.1.4.73336.30000025011601050770600215025"]……
            json中folders是单个序列的文件路径，
            可以直接进行分析上传，例如：
            "C:\\\\Users\\\\13167\\\\Desktop\\\\测试agent\\\\P10171793\\\\1.3.12.2.1107.5.1.4.73336.30000025011601050770600214550，
            可以直接调用“Analysis_dicom_directory”工具进行上传分析,当传如文件夹并不是一个序列的时候，请使用“fileforsep”工具进行拆分。
            这个函数只能使用一次当生成了csv文件后，就不能在执行了，避免重复拆分。"
                   """,
            inputSchema={
                "type": "object",
                "properties": {
                    "fileforsep": {
                        "type": "string",
                        "description": "待整理的顶层目录路径，执行过程中会在同级创建输出目录"
                    },
                    "async_mode": {
                        "type": "boolean",
                        "description": "是否异步执行（默认true）。如果为true，立即返回任务ID；如果为false，同步执行（可能超时）",
                        "default": True
                    }
                }, 
                "required": ["fileforsep"]
            }
        ),
        Tool(
            name="Analysis_dicom_directory",
            description="""
                    对单个 DICOM 序列（series）目录执行分析上传流程，并返回上传与处理的结果信息。
                    适用场景与前提：
                    - 该工具期望 `directory_path` 指向一个只包含单个 DICOM 序列（即一个Series）或一个已拆分好的序列目录。
                    - 若目录中包含多个序列，请先使用 `fileforsep` 对顶层目录进行拆分，然后对每个子序列目录单独调用本工具，需要多次执行这个函数，直到{fileforsep}产生的路径中的文件都被上传。
                    主要步骤：
                    1. 标准化并校验目录路径与文件权限。
                    2. 按 `series_type`（例如："1" 表示主动脉，"9" 表示二尖瓣）选择远端分析服务或参数集；不在允许列表中的 `series_type` 会被拒绝并返回错误。
                    3. 运行上传前检查（例如：文件完整性、必须的 DICOM 标签存在性、最大/最小切片数等），不满足要求的序列会返回明确原因并跳过上传。
                    4. 执行分步上传：上传序列元数据 -> 逐文件上传 DICOM 内容 -> 最终提交并确认。每一步都会记录进度与错误，以便恢复或重试。
                    错误与重试：如果单个序列上传失败，返回详细错误并允许用户对该目录跳过。
                    会将上传状态写到.csv文件中""",
            inputSchema={
                "type": "object",
                "properties": {
                    "directory_path": {
                        "type": "string",
                        "description": "包含待分析 DICOM 序列的本地目录路径，必须存在且具备读取权限"
                    },
                    "series_type": {
                        "type": "string",
                        "description": "分析流程类型：`1`=主动脉分析，`9`=二尖瓣分析，其他值将被拒绝"
                    }
                },
                "required": ["directory_path", "series_type"]
            }
        ),
        Tool(
            name="get_separate_task_status",
            description="""你间隔1分钟执行一次，直到完成，查询文件拆分任务的状态和进度。使用 fileforsep 工具后会返回 task_id，使用此工具查询任务执行状态，需要反复多次执行直到结束。"
                        同时他返回的内容会包括{fileforsep}的json信息，可以按照{fileforsep}的工具说明那样使用""",
            inputSchema={
                "type": "object",
                "properties": {
                    "task_id": {
                        "type": "string",
                        "description": "任务ID，从 fileforsep 工具的返回结果中获取"
                    },
                    "directory_path": {
                        "type": "string",
                        "description": "任务目录路径，与调用 fileforsep 时使用的 directory_path 相同"
                    }
                },
                "required": ["task_id", "directory_path"]
            }
        ),
      
    ]


@server.call_tool()
async def call_tool(name: str, arguments: Dict[str, Any]) -> list[TextContent]:
    """处理工具调用请求"""
    try:
        logger.info(f"调用工具: {name}, 参数: {arguments}")

        if name == "fileforsep":
            args = fileforsep(**arguments)
            result = await separate_series_by_patient_tool(
                args.fileforsep,
                async_mode=args.async_mode
            )
        elif name == "Analysis_dicom_directory":
            args = DirectoryPathWithSeriesArgs(**arguments)
            result = await Analysis_dicom_directory_tool(args.directory_path, args.series_type)

        elif name == "get_separate_task_status":
            args = TaskStatusQuery(**arguments)
            result = await get_separate_task_status_tool(args.task_id, args.directory_path)
        else:
            raise ValueError(f"未知工具: {name}")

        # 转换结果格式为MCP标准格式
        return [
            TextContent(
                type="text",
                text=content["text"]
            )
            for content in result["content"]
            if content["type"] == "text"
        ]

    except Exception as e:
        logger.error(f"工具调用失败: {name}, 错误: {e}", exc_info=True)

        error_response = {
            "error": True,
            "message": f"工具 {name} 执行失败: {str(e)}"
        }

        return [
            TextContent(
                type="text",
                text=json.dumps(error_response, ensure_ascii=False)
            )
        ]


async def main():
    """启动MCP服务器"""
    try:
        logger.info("启动 DICOM 工具 MCP 服务器 ...")

        # 使用stdio传输启动服务器
        async with stdio_server() as (read_stream, write_stream):
            await server.run(
                read_stream,
                write_stream,
                server.create_initialization_options()
            )

    except KeyboardInterrupt:
        logger.info("收到中断信号，正在关闭服务器...")
    except Exception as e:
        logger.error(f"服务器运行失败: {e}", exc_info=True)
        sys.exit(1)
    finally:
        # 清理资源
        _cleanup_resources()


def _cleanup_resources():
    """清理全局资源"""
    try:
        # 清理任务管理器中的线程池
        from src.function.task_manager import _shutdown_executor
        _shutdown_executor()
        logger.info("资源清理完成")
    except Exception as e:
        logger.warning(f"清理资源时出错: {e}")


def run():
    """同步入口函数，用于 uvx 调用"""
    # 设置事件循环策略（Windows兼容性）
    if sys.platform == "win32":
        asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())

    # 运行服务器
    asyncio.run(main())


if __name__ == "__main__":
    run()