#!/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, List, Optional

# 设置标准输出编码为 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
    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工具
# sys 和 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,
    export_measurement_csv_tool,
    get_all_studies_and_series_tool
)

# 配置日志 - 使用 INFO 级别以显示关键信息
logging.basicConfig(
    level=logging.INFO,  # 使用 INFO 级别以显示启动和关键操作信息
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stderr)
    ]
)
logger = logging.getLogger(__name__)

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

# 创建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参数相同

class ExportMeasurementCsvArgs(BaseModel):
    studyInstanceUids: List[str]

class GetAllStudiesAndSeriesArgs(BaseModel):
    searchStr: Optional[str] = None

@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"]
            }
        ),
        Tool(
            name="export-measurement-csv",
            description="从账号服务器（云端）导出指定StudyInstanceUID数组的测量数值CSV文件，返回下载URL。注意：这是账号服务器操作，不是本地文件操作。使用前请确认用户要导出的是账号服务器上的数据，而不是本地文件。",
            inputSchema={
                "type": "object",
                "properties": {
                    "studyInstanceUids": {
                        "type": "array",
                        "items": {
                            "type": "string"
                        },
                        "description": "账号服务器（云端）上患者影像的StudyInstanceUID数组，用于导出测量数值到CSV文件，不是本地文件的UID"
                    }
                },
                "required": ["studyInstanceUids"]
            }
        ),
        Tool(
            name="get-all-studies-and-series",
            description="从账号服务器（云端）获取全部已上传的病人影像序列信息，包括每个病例的所有序列的详细信息。注意：这是账号服务器操作，不是本地文件操作。使用前请确认用户要查询的是账号服务器上的数据，而不是本地目录。如果用户提到'本地'、'当前目录'、'工作目录'等，应该使用scan-dicom-directory等本地操作工具。",
            inputSchema={
                "type": "object",
                "properties": {
                    "searchStr": {
                        "type": "string",
                        "description": "可选，患者姓名、病例号、备注的模糊查询字符串，用于在账号服务器上搜索"
                    }
                },
                "required": []
            }
        ),
      
    ]


@server.call_tool()
async def call_tool(name: str, arguments: Dict[str, Any]) -> list[TextContent]:
    """处理工具调用请求"""
    try:
        print(f"\n🔨 收到工具调用请求: {name}", file=sys.stderr)
        print(f"📝 参数: {json.dumps(arguments, ensure_ascii=False, indent=2)}", file=sys.stderr)
        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)
        elif name == "export-measurement-csv":
            args = ExportMeasurementCsvArgs(**arguments)
            result = await export_measurement_csv_tool(args.studyInstanceUids)
        elif name == "get-all-studies-and-series":
            args = GetAllStudiesAndSeriesArgs(**arguments)
            result = await get_all_studies_and_series_tool(args.searchStr)
        else:
            raise ValueError(f"未知工具: {name}")

        # 验证结果格式并转换为MCP标准格式
        if not isinstance(result, dict):
            raise ValueError(f"工具 {name} 返回了无效的结果格式: 期望字典，得到 {type(result)}")
        
        if "content" not in result:
            raise ValueError(f"工具 {name} 返回的结果缺少 'content' 字段")
        
        if not isinstance(result["content"], list):
            raise ValueError(f"工具 {name} 返回的 'content' 字段不是列表类型")
        
        # 转换结果格式为MCP标准格式
        text_contents = []
        for content in result["content"]:
            if isinstance(content, dict) and content.get("type") == "text":
                text_contents.append(
                    TextContent(
                        type="text",
                        text=content.get("text", "")
                    )
                )
        
        # 如果没有找到任何文本内容，返回错误信息
        if not text_contents:
            logger.warning(f"工具 {name} 返回的结果中没有找到文本内容")
            print(f"⚠️  工具 {name} 未返回有效内容", file=sys.stderr)
            return [
                TextContent(
                    type="text",
                    text=json.dumps({
                        "error": True,
                        "message": f"工具 {name} 返回的结果中没有有效的文本内容",
                        "raw_result": str(result)
                    }, ensure_ascii=False)
                )
            ]
        
        print(f"✅ 工具 {name} 执行成功", file=sys.stderr)
        return text_contents

    except Exception as e:
        logger.error(f"工具调用失败: {name}, 错误: {e}", exc_info=True)
        print(f"❌ 工具 {name} 执行失败: {str(e)}", file=sys.stderr)

        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:
        # 打印启动信息到 stderr
        print("=" * 60, file=sys.stderr)
        print("🚀 DICOM 工具 MCP 服务器正在启动...", file=sys.stderr)
        print("=" * 60, file=sys.stderr)
        print(f"📋 服务器名称: dicom-tools-python", file=sys.stderr)
        print(f"📡 通信协议: MCP (Model Context Protocol) via stdio", file=sys.stderr)
        print(f"🔧 可用工具数量: 5", file=sys.stderr)
        print(f"   - fileforsep (DICOM文件拆分)", file=sys.stderr)
        print(f"   - Analysis_dicom_directory (DICOM序列分析)", file=sys.stderr)
        print(f"   - get_separate_task_status (任务状态查询)", file=sys.stderr)
        print(f"   - export-measurement-csv (导出测量数据)", file=sys.stderr)
        print(f"   - get-all-studies-and-series (获取所有病例)", file=sys.stderr)
        print("=" * 60, file=sys.stderr)
        print("✅ 服务器已启动，等待客户端连接...", file=sys.stderr)
        print("💡 提示: 这是一个 MCP 服务器，需要通过 MCP 客户端连接使用", file=sys.stderr)
        print("   (例如: Claude Desktop, Cursor 等)", file=sys.stderr)
        print("=" * 60, file=sys.stderr)
        print("", file=sys.stderr)
        
        logger.info("启动 DICOM 工具 MCP 服务器 ...")

        # 使用stdio传输启动服务器
        async with stdio_server() as (read_stream, write_stream):
            logger.info("stdio 连接已建立，服务器运行中...")
            await server.run(
                read_stream,
                write_stream,
                server.create_initialization_options()
            )

    except KeyboardInterrupt:
        print("\n" + "=" * 60, file=sys.stderr)
        print("⚠️  收到中断信号 (Ctrl+C)，正在关闭服务器...", file=sys.stderr)
        logger.info("收到中断信号，正在关闭服务器...")
    except Exception as e:
        print("\n" + "=" * 60, file=sys.stderr)
        print(f"❌ 服务器运行失败: {e}", file=sys.stderr)
        print("=" * 60, file=sys.stderr)
        logger.error(f"服务器运行失败: {e}", exc_info=True)
        sys.exit(1)
    finally:
        # 清理资源
        print("🧹 正在清理资源...", file=sys.stderr)
        _cleanup_resources()
        print("👋 服务器已关闭", file=sys.stderr)
        print("=" * 60, file=sys.stderr)


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


def run():
    """同步入口函数，用于 uvx 调用"""
    # 设置事件循环策略（Windows兼容性）
    if sys.platform == "win32":
        asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
        print("🪟 检测到 Windows 系统，已设置事件循环策略", file=sys.stderr)

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


if __name__ == "__main__":
    run()