from typing import Dict, List, Optional
import mcp.types as types
from mcp.server.lowlevel import Server


def create_messages(
    properties: Dict[str, str] = None,
    batch_size: int = 2,
) -> list[types.PromptMessage]:
    """
    创建用于材料生成和合成的提示词消息。

    Args:
        properties: 材料性质及其值的字典，例如 {"dft_band_gap": "2.0"}
        batch_size: 生成材料的数量，默认为2

    Returns:
        提示词消息列表
    """
    messages = []

    # 系统消息，定义助手的角色和任务
    system_message = """你是一位专业的材料科学家，擅长材料生成和合成方案设计。
你的任务是：
1. 根据用户提供的材料性质要求，使用mars_toolkit中的generate_materials工具生成符合要求的材料
2. 系统地分析生成的材料的四要素：成分、结构、工艺和性能
3. 为生成的材料设计科学合理的合成方案
4. 使用mermaid语法绘制材料的合成流程图

请确保你的回答包含以下内容：
- 对用户需求的理解和分析
- 使用generate_material工具生成的材料结构
- 生成材料的四要素详细分析：
  * 成分(Composition)：详细的化学成分、元素比例、化学计量比
  * 结构(Structure)：晶体结构、空间群、晶格参数、原子位置、配位环境
  * 工艺(Processing)：可行的合成路线、工艺参数、关键控制因素
  * 性能(Properties)：预期的物理、化学、机械性能及其与结构的关系
- 详细的合成方案，包括：
  * 原料选择及纯度要求
  * 精确的反应条件（温度、压力、时间、气氛）
  * 分步骤的合成流程及每步的理论依据
  * 可能的挑战及解决方案
  * 表征方法建议
- 使用mermaid语法绘制的合成流程图，清晰展示从原料到最终产品的全过程
"""

    messages.append(
        types.PromptMessage(
            role="system",
            content=types.TextContent(type="text", text=system_message),
        )
    )

    # 构建主提示词
    if properties and len(properties) > 0:
        properties_text = "\n".join([f"- {key}: {value}" for key, value in properties.items()])
        prompt = f"""请根据以下材料性质要求，生成{batch_size}个合适的材料并设计其合成方案：

{properties_text}

请按照以下步骤进行：

1. 使用mars_toolkit中的generate_material工具生成材料，参数设置为batch_size={batch_size}
2. 对生成的每种材料进行系统的四要素分析：
   - 成分：详细分析元素组成、化学计量比及其理论依据
   - 结构：分析晶体结构、空间群、晶格参数、原子排布及其稳定性
   - 工艺：探讨可行的合成路线、工艺参数及其科学依据
   - 性能：预测材料可能具有的物理、化学、机械性能及其应用前景

3. 为每种材料设计详细的合成方案，包括：
   - 原料选择及纯度要求
   - 精确的反应条件参数（温度、压力、时间、气氛等）
   - 分步骤的合成流程及每步的理论依据
   - 可能遇到的挑战及解决方案
   - 推荐的表征方法

4. 使用mermaid语法绘制材料的合成流程图，清晰展示从原料到最终产品的全过程，包括关键工艺参数。"""
    else:
        prompt = f"""请生成{batch_size}种具有创新性的新型材料并设计其合成方案。

请按照以下步骤进行：

1. 使用mars_toolkit中的generate_material工具生成材料，参数设置为batch_size={batch_size}
2. 对生成的每种材料进行系统的四要素分析：
   - 成分：详细分析元素组成、化学计量比及其理论依据
   - 结构：分析晶体结构、空间群、晶格参数、原子排布及其稳定性
   - 工艺：探讨可行的合成路线、工艺参数及其科学依据
   - 性能：预测材料可能具有的物理、化学、机械性能及其应用前景

3. 为每种材料设计详细的合成方案，包括：
   - 原料选择及纯度要求
   - 精确的反应条件参数（温度、压力、时间、气氛等）
   - 分步骤的合成流程及每步的理论依据
   - 可能遇到的挑战及解决方案
   - 推荐的表征方法

4. 使用mermaid语法绘制材料的合成流程图，清晰展示从原料到最终产品的全过程，包括关键工艺参数。"""

    messages.append(
        types.PromptMessage(
            role="user",
            content=types.TextContent(type="text", text=prompt)
        )
    )

    return messages


def register_prompt_handlers(app: Server):
    """
    注册提示词处理器到MCP服务器。

    Args:
        app: MCP服务器实例
    """
    @app.list_prompts()
    async def list_prompts() -> list[types.Prompt]:
        return [
            types.Prompt(
                name="material_synthesis",
                description="基于材料四要素（成分、结构、工艺、性能）生成材料并设计合成方案，使用mermaid绘制合成流程图",
                arguments=[
                    types.PromptArgument(
                        name="properties",
                        description="材料性质及其值的JSON字符串，例如 {\"dft_band_gap\": \"2.0\"}",
                        required=False,
                    ),
                    types.PromptArgument(
                        name="batch_size",
                        description="生成材料的数量，默认为2",
                        required=False,
                    ),
                ],
            )
        ]

    @app.get_prompt()
    async def get_prompt(
        name: str, arguments: dict[str, str] | None = None
    ) -> types.GetPromptResult:
        if name != "material_synthesis":
            raise ValueError(f"未知的提示词: {name}")

        if arguments is None:
            arguments = {}

        # 解析properties参数
        properties = {}
        if "properties" in arguments and arguments["properties"]:
            try:
                import json
                properties = json.loads(arguments["properties"])
            except json.JSONDecodeError:
                properties = {}

        # 解析batch_size参数
        batch_size = 2  # 默认值
        if "batch_size" in arguments and arguments["batch_size"]:
            try:
                batch_size = int(arguments["batch_size"])
            except ValueError:
                pass  # 使用默认值

        return types.GetPromptResult(
            messages=create_messages(properties=properties, batch_size=batch_size),
            description="基于材料四要素（成分、结构、工艺、性能）生成材料并设计合成方案，使用mermaid绘制合成流程图",
        )