"""
测试晶体结构优化工具函数

此脚本测试改进后的optimize_crystal_structure函数，
该函数接受单一的file_name_or_content_string参数，可以是文件路径或直接的结构内容。
"""

import sys
import asyncio
import os
import tempfile
sys.path.append("/home/ubuntu/sas0/lzy/multi_mcp_server")

from sci_mcp.material_mcp.fairchem_structure_opt.structure_opt_tools import optimize_crystal_structure
from sci_mcp.core.config import material_config

# 简单的CIF结构示例
SAMPLE_CIF = """
data_SrTiO3
_cell_length_a 3.905
_cell_length_b 3.905
_cell_length_c 3.905
_cell_angle_alpha 90
_cell_angle_beta 90
_cell_angle_gamma 90
_symmetry_space_group_name_H-M 'P m -3 m'
_symmetry_Int_Tables_number 221
loop_
_atom_site_label
_atom_site_type_symbol
_atom_site_fract_x
_atom_site_fract_y
_atom_site_fract_z
Sr1 Sr 0.0 0.0 0.0
Ti1 Ti 0.5 0.5 0.5
O1 O 0.5 0.5 0.0
O2 O 0.5 0.0 0.5
O3 O 0.0 0.5 0.5
"""

async def test_with_content():
    """测试使用直接结构内容"""
    print("\n=== 测试使用直接结构内容 ===")
    result = await optimize_crystal_structure(
        file_name_or_content_string=SAMPLE_CIF,
        format_type="cif",
        optimization_level="quick"
    )
    print(result)

async def test_with_file():
    """测试使用文件路径（如果文件存在）"""
    print("\n=== 测试使用文件路径 ===")
    # 创建临时CIF文件

    with tempfile.NamedTemporaryFile(suffix=".cif", mode="w", delete=False) as tmp_file:
        tmp_file.write(SAMPLE_CIF)
        tmp_path = tmp_file.name

    try:
        result = await optimize_crystal_structure(
            file_name_or_content_string=tmp_path,
            format_type="auto",
            optimization_level="quick"
        )
        print(result)
    finally:
        # 清理临时文件
        if os.path.exists(tmp_path):
            os.unlink(tmp_path)

async def test_with_temp_file():
    """测试使用临时目录中的文件名"""
    print("\n=== 测试使用临时目录中的文件名 ===")

    # 确保临时目录存在
    os.makedirs(material_config.TEMP_ROOT, exist_ok=True)

    # 在临时目录中创建文件
    temp_filename = "test_structure.cif"
    temp_filepath = os.path.join(material_config.TEMP_ROOT, temp_filename)

    with open(temp_filepath, 'w', encoding='utf-8') as f:
        f.write(SAMPLE_CIF)

    try:
        # 只传递文件名，而不是完整路径
        result = await optimize_crystal_structure(
            file_name_or_content_string=temp_filename,
            format_type="auto",
            optimization_level="quick"
        )
        print(result)
    finally:
        # 清理临时文件
        if os.path.exists(temp_filepath):
            os.unlink(temp_filepath)

async def test_auto_format():
    """测试自动格式检测"""
    print("\n=== 测试自动格式检测 ===")
    result = await optimize_crystal_structure(
        file_name_or_content_string=SAMPLE_CIF,
        format_type="auto"
    )
    print(result)

async def main():
    """运行所有测试"""
    print("测试改进后的optimize_crystal_structure函数")

    await test_with_content()
    await test_with_file()
    await test_with_temp_file()
    await test_auto_format()

if __name__ == "__main__":
    asyncio.run(main())