"""
Test script for PubChem tools

This script tests the search_pubchem_advanced function from the chemistry_mcp module.
"""
import sys
sys.path.append('/home/ubuntu/sas0/lzy/multi_mcp_server')
import asyncio

from sci_mcp.chemistry_mcp.pubchem_tools.pubchem_tools import _search_by_formula


from sci_mcp.chemistry_mcp import search_pubchem_advanced

async def test_search_by_name():
    """Test searching compounds by name"""
    print("\n=== Testing search by name ===")
    result = await search_pubchem_advanced(name="Aspirin")
    print(result)

async def test_search_by_smiles():
    """Test searching compounds by SMILES notation"""
    print("\n=== Testing search by SMILES ===")
    # SMILES for Caffeine
    result = await search_pubchem_advanced(smiles="CN1C=NC2=C1C(=O)N(C(=O)N2C)C")
    print(result)

async def test_search_by_formula():
    """Test searching compounds by molecular formula"""
    print("\n=== Testing search by formula ===")
    # Formula for Aspirin
    result = await search_pubchem_advanced(formula="C9H8O4", max_results=2)
    print(result)

async def test_complex_formula():
    """Test searching with a more complex formula that might cause timeout"""
    print("\n=== Testing complex formula search ===")
    # A more complex formula that might return many results
    result = await search_pubchem_advanced(
        formula="C6H12O6",  # Glucose and isomers
        max_results=5
    )
    print(result)

async def test_complex_molecules():
    """Test searching for complex molecules with rich molecular features"""
    print("\n=== Testing complex molecules with rich features ===")

    # 1. Paclitaxel (Taxol) - Complex anticancer drug with many rotatable bonds and H-bond donors/acceptors
    print("\n--- Testing Paclitaxel (anticancer drug) ---")
    result = await search_pubchem_advanced(name="Paclitaxel")
    print(result)

    # 2. Vancomycin - Complex antibiotic with many H-bond donors/acceptors
    print("\n--- Testing Vancomycin (antibiotic) ---")
    result = await search_pubchem_advanced(name="Vancomycin")
    print(result)

    # 3. Cholesterol - Steroid with complex ring structure
    print("\n--- Testing Cholesterol (steroid) ---")
    result = await search_pubchem_advanced(name="Cholesterol")
    print(result)

    # 4. Ibuprofen - Common NSAID with rotatable bonds
    print("\n--- Testing Ibuprofen (NSAID) ---")
    result = await search_pubchem_advanced(name="Ibuprofen")
    print(result)

    # 5. Amoxicillin - Antibiotic with multiple functional groups
    print("\n--- Testing Amoxicillin (antibiotic) ---")
    result = await search_pubchem_advanced(name="Amoxicillin")
    print(result)

async def test_molecules_by_smiles():
    """Test searching for complex molecules using SMILES notation"""
    print("\n=== Testing complex molecules by SMILES ===")

    # 1. Atorvastatin (Lipitor) - Cholesterol-lowering drug with complex structure
    print("\n--- Testing Atorvastatin (Lipitor) ---")
    result = await search_pubchem_advanced(
        smiles="CC(C)C1=C(C(=C(C=C1)C(C)C)C2=CC(=C(C=C2)F)F)C(CC(CC(=O)O)O)NC(=O)C3=CC=C(C=C3)F"
    )
    print(result)

    # 2. Morphine - Opioid with multiple rings and H-bond features
    print("\n--- Testing Morphine ---")
    result = await search_pubchem_advanced(
        smiles="CN1CCC23C4C1CC5=C2C(=C(C=C5)O)OC3C(C=C4)O"
    )
    print(result)

async def test_invalid_search():
    """Test searching with invalid parameters"""
    print("\n=== Testing invalid search ===")
    # No parameters provided
    result = await search_pubchem_advanced()
    print(result)

    # Invalid SMILES
    print("\n=== Testing invalid SMILES ===")
    result = await search_pubchem_advanced(smiles="INVALID_SMILES_STRING")
    print(result)

async def run_all_tests():
    """Run all test functions"""
    await test_search_by_name()
    await test_search_by_smiles()
    await test_search_by_formula()
    await test_complex_formula()
    # await test_complex_molecules()
    # await test_molecules_by_smiles()
    #await test_invalid_search()
    # from sci_mcp.chemistry_mcp.pubchem_tools.pubchem_tools import _search_by_name
    # compounds=await _search_by_formula('C6H12O6')
    # print(compounds[0])
if __name__ == "__main__":
    print("Testing PubChem search tools...")
    asyncio.run(run_all_tests())
    print("\nAll tests completed.")
    # import pubchempy
    # compunnds = pubchempy.get_compounds('Aspirin', 'name')
    # print(compunnds[0].to_dict())