"""
诊微 AI 病历质控 MCP Server

标准 MCP 协议实现，支持 stdio 和 SSE 两种传输方式。
可在 Cursor、Claude Desktop、Cline 等 AI 工具中直接使用。

环境变量:
    ZHENWEI_API_URL   — 诊微平台 API 地址 (例: https://your-server.com/api/v1)
    ZHENWEI_API_KEY   — MCP API Key (在诊微平台 MCP 管理页面创建)
"""

from __future__ import annotations

import json
import os
import sys
from typing import Any

import httpx
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import (
    TextContent,
    Tool,
)
from pydantic import BaseModel

# ---- 配置 ----

API_URL = os.environ.get("ZHENWEI_API_URL", "http://localhost:8281/api/v1")
API_KEY = os.environ.get("ZHENWEI_API_KEY", "")

server = Server("zhenwei-mcp-server")


# ---- HTTP 客户端 ----

def _client() -> httpx.AsyncClient:
    return httpx.AsyncClient(
        base_url=API_URL,
        headers={"X-API-Key": API_KEY, "Content-Type": "application/json"},
        timeout=120.0,
    )


async def _post(path: str, data: dict) -> dict:
    async with _client() as c:
        r = await c.post(f"/mcp{path}", json=data)
        r.raise_for_status()
        return r.json()


async def _get(path: str, params: dict | None = None) -> dict:
    async with _client() as c:
        r = await c.get(f"/mcp{path}", params=params)
        r.raise_for_status()
        return r.json()


def _ok(result: Any) -> list[TextContent]:
    """统一格式化返回."""
    if isinstance(result, dict):
        text = json.dumps(result, ensure_ascii=False, indent=2)
    else:
        text = str(result)
    return [TextContent(type="text", text=text)]


def _err(msg: str) -> list[TextContent]:
    return [TextContent(type="text", text=f"❌ 错误: {msg}")]


# ============================================================
#  注册 12 个 Tool
# ============================================================

@server.list_tools()
async def list_tools() -> list[Tool]:
    return [
        Tool(
            name="check_medical_record",
            description="病历质控检查 — 提交病历文本，AI 引擎自动检测缺陷（完整性、逻辑性、时效性、内涵质量），返回评分和缺陷列表。支持规则引擎 + 大模型语义双重检测。",
            inputSchema={
                "type": "object",
                "properties": {
                    "record_text": {
                        "type": "string",
                        "description": "病历全文文本（包含主诉、现病史、既往史、体格检查、诊断等）",
                    },
                    "record_type": {
                        "type": "string",
                        "description": "病历类型: admission(入院记录) / discharge(出院记录) / progress(病程记录) / surgery(手术记录)",
                        "default": "admission",
                    },
                    "check_level": {
                        "type": "string",
                        "description": "检查级别: full(全面检查) / basic(基础检查) / quick(快速检查)",
                        "default": "full",
                    },
                },
                "required": ["record_text"],
            },
        ),
        Tool(
            name="batch_check_records",
            description="批量病历质控 — 一次提交最多 50 份病历进行批量质控，返回每份病历的评分和缺陷数量。",
            inputSchema={
                "type": "object",
                "properties": {
                    "records": {
                        "type": "array",
                        "description": "病历数组",
                        "items": {
                            "type": "object",
                            "properties": {
                                "record_text": {"type": "string", "description": "病历文本"},
                                "record_type": {"type": "string", "default": "admission"},
                                "external_record_id": {"type": "string", "description": "外部系统的病历ID"},
                            },
                            "required": ["record_text"],
                        },
                        "maxItems": 50,
                    },
                    "check_level": {"type": "string", "default": "full"},
                },
                "required": ["records"],
            },
        ),
        Tool(
            name="get_qc_result",
            description="获取质控结果 — 通过任务 ID 查询之前提交的质控任务的详细结果。",
            inputSchema={
                "type": "object",
                "properties": {
                    "task_id": {"type": "integer", "description": "质控任务 ID"},
                },
                "required": ["task_id"],
            },
        ),
        Tool(
            name="get_qc_rules",
            description="查询质控规则 — 获取平台当前启用的质控规则列表，支持按类别、严重级别筛选。可用于了解质控标准。",
            inputSchema={
                "type": "object",
                "properties": {
                    "category": {
                        "type": "string",
                        "description": "规则类别: veto(单项否决) / completeness(完整性) / consistency(一致性) / timeliness(时效性)",
                    },
                    "severity": {
                        "type": "string",
                        "description": "严重级别: VETO(否决) / MAJOR(重大) / MINOR(一般) / SUGGESTION(建议)",
                    },
                    "keyword": {"type": "string", "description": "关键词搜索"},
                    "page": {"type": "integer", "default": 1},
                    "page_size": {"type": "integer", "default": 20},
                },
            },
        ),
        Tool(
            name="get_qc_statistics",
            description="质控统计数据 — 获取指定时间段的质控统计概况（平均分、甲乙丙级分布、合格率等）。",
            inputSchema={
                "type": "object",
                "properties": {
                    "days": {
                        "type": "integer",
                        "description": "统计天数（默认30天）",
                        "default": 30,
                    },
                },
            },
        ),
        Tool(
            name="search_knowledge_base",
            description="知识库搜索 — 搜索病历书写规范、卫生法规、ICD编码、药物知识等参考资料。",
            inputSchema={
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "搜索关键词"},
                    "category": {
                        "type": "string",
                        "description": "知识类别: regulation(法规) / standard(标准) / icd(编码) / drug(药物)",
                    },
                    "top_k": {"type": "integer", "default": 5, "description": "返回条数"},
                },
                "required": ["query"],
            },
        ),
        Tool(
            name="get_record_detail",
            description="获取病历详情 — 获取指定病历的结构化内容和最新质控结果。",
            inputSchema={
                "type": "object",
                "properties": {
                    "record_id": {"type": "integer", "description": "病历 ID"},
                },
                "required": ["record_id"],
            },
        ),
        Tool(
            name="submit_appeal",
            description="提交缺陷申诉 — 对质控发现的缺陷提交申诉，附带理由和佐证。",
            inputSchema={
                "type": "object",
                "properties": {
                    "workorder_id": {"type": "integer", "description": "工单 ID"},
                    "defect_id": {"type": "integer", "description": "缺陷 ID"},
                    "appeal_reason": {"type": "string", "description": "申诉理由"},
                    "appeal_evidence": {"type": "string", "description": "佐证材料"},
                },
                "required": ["workorder_id", "appeal_reason"],
            },
        ),
        Tool(
            name="get_suggestions",
            description="AI 修改建议 — 针对质控缺陷获取 AI 生成的修改建议和示范文本。",
            inputSchema={
                "type": "object",
                "properties": {
                    "defect_description": {"type": "string", "description": "缺陷描述"},
                    "record_type": {"type": "string", "description": "病历类型"},
                    "record_text": {"type": "string", "description": "病历原文（提供则建议更精准）"},
                },
                "required": ["defect_description"],
            },
        ),
        Tool(
            name="check_icd_code",
            description="ICD 编码校验 — AI 辅助 ICD-10 编码校验和推荐，输入诊断文本自动匹配正确编码。",
            inputSchema={
                "type": "object",
                "properties": {
                    "diagnosis_text": {"type": "string", "description": "诊断文本"},
                    "icd_code": {"type": "string", "description": "待校验的 ICD 编码（可选）"},
                },
                "required": ["diagnosis_text"],
            },
        ),
        Tool(
            name="check_drg_grouping",
            description="DRG 分组预测 — 根据诊断和手术编码预测 DRG 分组，检测编码风险和高编低编。",
            inputSchema={
                "type": "object",
                "properties": {
                    "principal_diagnosis_code": {"type": "string", "description": "主要诊断 ICD 编码"},
                    "procedure_codes": {"type": "array", "items": {"type": "string"}, "description": "手术/操作编码列表"},
                    "age": {"type": "integer"},
                    "gender": {"type": "string"},
                    "complications": {"type": "array", "items": {"type": "string"}, "description": "并发症编码列表"},
                    "los": {"type": "integer", "description": "住院天数"},
                },
                "required": ["principal_diagnosis_code"],
            },
        ),
        Tool(
            name="check_drug_safety",
            description="用药安全检查 — 检查过敏禁忌、配伍禁忌、剂量合理性、给药途径等用药安全问题。",
            inputSchema={
                "type": "object",
                "properties": {
                    "drug_name": {"type": "string", "description": "药物名称"},
                    "dose": {"type": "number", "description": "剂量"},
                    "unit": {"type": "string", "default": "mg"},
                    "route": {"type": "string", "default": "口服", "description": "给药途径"},
                    "frequency": {"type": "string", "default": "qd", "description": "频次"},
                    "allergy_history": {"type": "string", "description": "过敏史"},
                    "other_drugs": {"type": "array", "items": {"type": "string"}, "description": "联用药物"},
                    "patient_age": {"type": "integer"},
                    "diagnosis": {"type": "string"},
                },
                "required": ["drug_name"],
            },
        ),
    ]


# ============================================================
#  Tool 调用路由
# ============================================================

@server.call_tool()
async def call_tool(name: str, arguments: dict) -> list[TextContent]:
    try:
        if name == "check_medical_record":
            result = await _post("/check", arguments)
            return _ok(result)

        elif name == "batch_check_records":
            result = await _post("/batch-check", arguments)
            return _ok(result)

        elif name == "get_qc_result":
            result = await _get(f"/tasks/{arguments['task_id']}")
            return _ok(result)

        elif name == "get_qc_rules":
            params = {k: v for k, v in arguments.items() if v is not None}
            result = await _get("/rules", params)
            return _ok(result)

        elif name == "get_qc_statistics":
            result = await _get("/statistics", {"days": arguments.get("days", 30)})
            return _ok(result)

        elif name == "search_knowledge_base":
            result = await _post("/knowledge/search", arguments)
            return _ok(result)

        elif name == "get_record_detail":
            result = await _get(f"/records/{arguments['record_id']}")
            return _ok(result)

        elif name == "submit_appeal":
            result = await _post("/appeals", arguments)
            return _ok(result)

        elif name == "get_suggestions":
            result = await _post("/suggestions", arguments)
            return _ok(result)

        elif name == "check_icd_code":
            result = await _post("/icd/check", arguments)
            return _ok(result)

        elif name == "check_drg_grouping":
            result = await _post("/drg/check", arguments)
            return _ok(result)

        elif name == "check_drug_safety":
            result = await _post("/drug/check", arguments)
            return _ok(result)

        else:
            return _err(f"未知工具: {name}")

    except httpx.HTTPStatusError as e:
        body = e.response.text
        return _err(f"API 返回 {e.response.status_code}: {body}")
    except httpx.ConnectError:
        return _err(f"无法连接到诊微平台 ({API_URL})，请检查 ZHENWEI_API_URL 配置")
    except Exception as e:
        return _err(str(e))


# ============================================================
#  入口
# ============================================================

def main():
    import asyncio

    if not API_KEY:
        print("⚠️  未设置 ZHENWEI_API_KEY 环境变量", file=sys.stderr)
        print("   请在诊微平台 MCP 管理页面创建 API Key，然后设置:", file=sys.stderr)
        print("   export ZHENWEI_API_KEY=mqc_your_key_here", file=sys.stderr)
        sys.exit(1)

    async def run():
        async with stdio_server() as (read_stream, write_stream):
            await server.run(read_stream, write_stream, server.create_initialization_options())

    asyncio.run(run())


if __name__ == "__main__":
    main()
