"""
测试RXN工具函数模块

此模块包含用于测试rxn_tools模块中化学反应预测和分析工具函数的测试用例。
"""
import asyncio
import sys
import os
from pathlib import Path
from rich.console import Console

# 添加项目根目录到Python路径
sys.path.append('/home/ubuntu/sas0/lzy/multi_mcp_server')

from sci_mcp.chemistry_mcp.rxn_tools.rxn_tools import (
    predict_reaction_outcome_rxn,
    predict_reaction_topn_rxn,
    predict_reaction_properties_rxn,
    extract_reaction_actions_rxn
)

# 创建控制台对象用于格式化输出
console = Console()


async def test_predict_reaction_outcome():
    """测试反应结果预测功能"""
    console.print("[bold cyan]测试反应结果预测功能[/bold cyan]")

    # 使用固定参数：溴和蒽的反应
    fixed_reactants = "BrBr.c1ccc2cc3ccccc3cc2c1"
    console.print(f"固定反应物: {fixed_reactants}")

    try:
        result = await predict_reaction_outcome(fixed_reactants)
        console.print("[green]结果:[/green]")
        console.print(result)
    except Exception as e:
        console.print(f"[bold red]错误:[/bold red] {str(e)}")


async def test_predict_reaction_topn():
    """测试多产物预测功能"""
    console.print("\n[bold cyan]测试多产物预测功能[/bold cyan]")

    # 测试1：单个反应（字符串格式）
    fixed_reactants = "C=CC=O.CN"  # 丙烯醛和甲胺
    fixed_topn = 2
    console.print(f"测试1 - 单个反应（字符串格式）")
    console.print(f"固定反应物: {fixed_reactants}")
    console.print(f"固定预测产物数量: {fixed_topn}")

    try:
        result = await predict_reaction_topn(fixed_reactants, fixed_topn)
        console.print("[green]结果:[/green]")
        console.print(result)
    except Exception as e:
        console.print(f"[bold red]错误:[/bold red] {str(e)}")

    # 测试2：单个反应（列表格式）
    fixed_reactants_list = ["BrBr", "c1ccc2cc3ccccc3cc2c1"]  # 溴和蒽
    console.print(f"\n测试2 - 单个反应（列表格式）")
    console.print(f"固定反应物: {fixed_reactants_list}")
    console.print(f"固定预测产物数量: {fixed_topn}")

    try:
        result = await predict_reaction_topn(fixed_reactants_list, fixed_topn)
        console.print("[green]结果:[/green]")
        console.print(result)
    except Exception as e:
        console.print(f"[bold red]错误:[/bold red] {str(e)}")

    # 测试3：多个反应（列表的列表格式）
    fixed_reactants_batch = [
        ["BrBr", "c1ccc2cc3ccccc3cc2c1"],  # 溴和蒽
        ["BrBr", "c1ccc2cc3ccccc3cc2c1CCO"]  # 溴和修饰的蒽
    ]
    console.print(f"\n测试3 - 多个反应（列表的列表格式）")
    console.print(f"固定反应物批量: {fixed_reactants_batch}")
    console.print(f"固定预测产物数量: {fixed_topn}")

    try:
        result = await predict_reaction_topn(fixed_reactants_batch, fixed_topn)
        console.print("[green]结果:[/green]")
        console.print(result)
    except Exception as e:
        console.print(f"[bold red]错误:[/bold red] {str(e)}")
        console.print(f"[bold red]错误:[/bold red] {str(e)}")


async def test_predict_retrosynthesis():
    """测试逆合成分析功能"""
    console.print("\n[bold cyan]测试逆合成分析功能[/bold cyan]")

    # 使用固定参数：阿司匹林的逆合成分析
    fixed_target_molecule = "CC(=O)OC1=CC=CC=C1C(=O)O"  # 阿司匹林
    fixed_max_steps = 1
    console.print(f"固定目标分子: {fixed_target_molecule}")
    console.print(f"固定最大步骤数: {fixed_max_steps}")

    try:
        result = await predict_retrosynthesis(fixed_target_molecule, fixed_max_steps)
        console.print("[green]结果:[/green]")
        console.print(result)
    except Exception as e:
        console.print(f"[bold red]错误:[/bold red] {str(e)}")


async def test_predict_biocatalytic_retrosynthesis():
    """测试生物催化逆合成分析功能"""
    console.print("\n[bold cyan]测试生物催化逆合成分析功能[/bold cyan]")

    # 使用固定参数：一个可能适合酶催化的分子
    fixed_target_molecule = "OC1C(O)C=C(Br)C=C1"
    console.print(f"固定目标分子: {fixed_target_molecule}")

    try:
        result = await predict_biocatalytic_retrosynthesis(fixed_target_molecule)
        console.print("[green]结果:[/green]")
        console.print(result)
    except Exception as e:
        console.print(f"[bold red]错误:[/bold red] {str(e)}")


async def test_predict_reaction_properties():
    """测试反应属性预测功能"""
    console.print("\n[bold cyan]测试反应属性预测功能[/bold cyan]")

    # 使用固定参数：原子映射
    fixed_reaction = "CC(C)S.CN(C)C=O.Fc1cccnc1F.O=C([O-])[O-].[K+].[K+]>>CC(C)Sc1ncccc1F"
    fixed_property_type = "atom-mapping"
    console.print(f"固定反应: {fixed_reaction}")
    console.print(f"固定属性类型: {fixed_property_type}")

    try:
        result = await predict_reaction_properties(fixed_reaction, fixed_property_type)
        console.print("[green]结果:[/green]")
        console.print(result)
    except Exception as e:
        console.print(f"[bold red]错误:[/bold red] {str(e)}")


async def test_extract_reaction_actions():
    """测试从文本提取反应步骤功能"""
    console.print("\n[bold cyan]测试从文本提取反应步骤功能[/bold cyan]")

    # 使用固定参数：从文本描述中提取反应步骤
    fixed_reaction_text = """To a stirred solution of 7-(difluoromethylsulfonyl)-4-fluoro-indan-1-one (110 mg, 0.42 mmol) in methanol (4 mL) was added sodium borohydride (24 mg, 0.62 mmol). The reaction mixture was stirred at ambient temperature for 1 hour."""
    console.print(f"固定反应文本: {fixed_reaction_text}")

    try:
        result = await extract_reaction_actions(fixed_reaction_text)
        console.print("[green]结果:[/green]")
        console.print(result)
    except Exception as e:
        console.print(f"[bold red]错误:[/bold red] {str(e)}")


async def test_all():
    """测试所有RXN工具函数"""
    console.print("[bold magenta]===== 开始测试RXN工具函数 =====[/bold magenta]\n")

    # 测试各个功能
    await test_predict_reaction_outcome()
    await test_predict_reaction_topn()
    await test_predict_retrosynthesis()
    await test_predict_biocatalytic_retrosynthesis()
    await test_predict_reaction_properties()
    await test_extract_reaction_actions()

    console.print("\n[bold magenta]===== RXN工具函数测试完成 =====[/bold magenta]")


def get_rxn_tool_questions():
    """
    获取为RXN工具函数生成的问题列表
    这些问题设计为能够引导大模型调用相应的工具函数

    Returns:
        包含工具名称和对应问题的字典
    """
    questions = {
        "predict_reaction_outcome": [
            "如果我将溴和蒽混合在一起，会形成什么产物？",
            "乙酸和乙醇反应会生成什么？",
            "预测一下丙烯醛和甲胺反应的结果"
        ],


        "predict_reaction_topn": [
            "丙烯醛和甲胺反应可能生成哪几种主要产物？",
            "预测溴和蒽反应可能的前3个产物",
            "乙酸和乙醇反应可能有哪些不同的结果？请给出最可能的几种产物"
        ],

        "predict_retrosynthesis": [
            "如何合成阿司匹林？请给出可能的合成路线",
            "对于分子CC(=O)OC1=CC=CC=C1C(=O)O，有哪些可能的合成路径？",
            "请分析一下布洛芬的可能合成路线"
        ],

        "predict_biocatalytic_retrosynthesis": [
            "有没有可能用酶催化合成OC1C(O)C=C(Br)C=C1这个分子？",
            "请提供一种使用生物催化方法合成对羟基苯甲醇的路线",
            "我想用酶催化方法合成一些复杂分子，能否分析一下OC1C(O)C=C(Br)C=C1的可能合成路径？"
        ],

        "predict_reaction_properties": [
            "在这个反应中：CC(C)S.CN(C)C=O.Fc1cccnc1F.O=C([O-])[O-].[K+].[K+]>>CC(C)Sc1ncccc1F，原子是如何映射的？",
            "这个反应的产率可能是多少：Clc1ccccn1.Cc1ccc(N)cc1>>Cc1ccc(Nc2ccccn2)cc1",
            "能分析一下这个反应中原子的去向吗：CC(=O)O.CCO>>CC(=O)OCC"
        ],

        "extract_reaction_actions": [
            "能否将这段实验描述转换为结构化的反应步骤？'To a stirred solution of 7-(difluoromethylsulfonyl)-4-fluoro-indan-1-one (110 mg, 0.42 mmol) in methanol (4 mL) was added sodium borohydride (24 mg, 0.62 mmol). The reaction mixture was stirred at ambient temperature for 1 hour.'",
            "请从这段文本中提取出具体的反应操作步骤：'A solution of benzoic acid (1.0 g, 8.2 mmol) in thionyl chloride (10 mL) was heated under reflux for 2 hours. The excess thionyl chloride was removed under reduced pressure to give benzoyl chloride as a colorless liquid.'",
            "帮我解析这个实验步骤，提取出关键操作：'The aldehyde (5 mmol) was dissolved in methanol (20 mL) and sodium borohydride (7.5 mmol) was added portionwise at 0°C. The mixture was allowed to warm to room temperature and stirred for 3 hours.'"
        ]
    }

    return questions


def update_agent_test_questions():
    """
    更新agent_test.py中的工具问题字典，添加RXN工具函数的问题
    """
    try:
        # 获取agent_test.py文件路径
        agent_test_path = Path('/home/ubuntu/sas0/lzy/multi_mcp_server/test_tools/agent_test.py')

        # 读取文件内容
        with open(agent_test_path, 'r') as f:
            content = f.read()

        # 获取RXN工具函数的问题
        rxn_questions = get_rxn_tool_questions()

        # 检查文件中是否已包含RXN工具函数的问题
        rxn_tools_exist = any(tool in content for tool in rxn_questions.keys())

        if not rxn_tools_exist:
            # 找到questions字典的结束位置
            dict_end_pos = content.find('    return questions')

            if dict_end_pos != -1:
                # 构建要插入的RXN工具函数问题
                rxn_questions_str = ""
                for tool_name, questions_list in rxn_questions.items():
                    rxn_questions_str += f'\n        "{tool_name}": [\n'
                    for q in questions_list:
                        rxn_questions_str += f'            "{q}",\n'
                    rxn_questions_str += '        ],'

                # 在字典结束前插入RXN工具函数问题
                new_content = content[:dict_end_pos] + rxn_questions_str + content[dict_end_pos:]

                # 写回文件
                with open(agent_test_path, 'w') as f:
                    f.write(new_content)

                console.print("[green]成功更新agent_test.py，添加了RXN工具函数的测试问题[/green]")
            else:
                console.print("[yellow]无法找到questions字典的结束位置，未更新agent_test.py[/yellow]")
        else:
            console.print("[yellow]agent_test.py中已包含RXN工具函数的问题，无需更新[/yellow]")

    except Exception as e:
        console.print(f"[bold red]更新agent_test.py时出错:[/bold red] {str(e)}")


if __name__ == "__main__":
    # 运行所有测试
    # asyncio.run(test_all())

    # # 更新agent_test.py中的工具问题
    # update_agent_test_questions()

     api_key = 'apk-8928522a146c2503f30b16d9909222d7583f412ee8f1049f08d32a089ba88d34'
     from rxn4chemistry import RXN4ChemistryWrapper

     rxn4chemistry_wrapper = RXN4ChemistryWrapper(api_key=api_key)
     rxn4chemistry_wrapper.create_project('test_wrapper')
     response = rxn4chemistry_wrapper.predict_automatic_retrosynthesis(
    'Brc1c2ccccc2c(Br)c2ccccc12')
     results = rxn4chemistry_wrapper.get_predict_automatic_retrosynthesis_results(response['prediction_id'])
     print(results['status'])
    # NOTE: upon 'SUCCESS' you can inspect the predicted retrosynthetic paths.
     print(results['retrosynthetic_paths'][0])