"""
Test script for domain-specific tool retrieval functions
"""

import json
from pprint import pprint
import sys
sys.path.append("/home/ubuntu/sas0/lzy/multi_mcp_server")
from sci_mcp.core.llm_tools import get_domain_tools, get_domain_tool_schemas, get_all_tool_schemas,get_all_tools

def test_get_domain_tools():
    """Test retrieving tools from specific domains"""
    print("\n=== Testing get_domain_tools ===")

    # Test with material and general domains
    domains = ['material', 'general']
    print(f"Getting tools for domains: {domains}")

    domain_tools = get_domain_tools(domains)

    # Print results
    for domain, tools in domain_tools.items():
        print(f"\nDomain: {domain}")
        print(f"Number of tools: {len(tools)}")
        print("Tool names:")
        for tool_name in tools.keys():
            print(f"  - {tool_name}")

def test_get_domain_tool_schemas():
    """Test retrieving tool schemas from specific domains"""
    print("\n=== Testing get_domain_tool_schemas (OpenAI format) ===")

    # Test with material and general domains
    domains = ['material', 'general']
    print(f"Getting tool schemas for domains: {domains}")

    domain_schemas = get_domain_tool_schemas(domains)

    # Print results
    for domain, schemas in domain_schemas.items():
        print(f"\nDomain: {domain}")
        print(f"Number of schemas: {len(schemas)}")
        print("Tool names in schemas:")
        for schema in schemas:
            print(f"  - {schema['function']['name']}")


def test_all_tool_schemas():
    """Test retrieving all tool schemas in both formats"""
    print("\n=== Testing get_all_tool_schemas ===")

    # OpenAI format
    print("Getting all tool schemas in OpenAI format")
    openai_schemas = get_all_tool_schemas()
    print(f"Number of schemas: {len(openai_schemas)}")
    print("Tool names:")
    for schema in openai_schemas:
        print(f"  - {schema['function']['name']}")

    # MCP format
    print("\nGetting all tool schemas in MCP format")
    mcp_schemas = get_all_tool_schemas(use_mcp_format=True)
    print(f"Number of schemas: {len(mcp_schemas)}")
    print("Tool names:")
    for schema in mcp_schemas:
        print(f"  - {schema['name']}")

def main():
    """Main function to run tests"""
    print("Testing domain-specific tool retrieval functions")

    #test_get_domain_tools()
    test_get_domain_tool_schemas()

    test_all_tool_schemas()

if __name__ == "__main__":
    #print(get_all_tool_schemas())
    from  rich.console import Console
    console = Console()
    console.print("Testing domain-specific tool retrieval functions", style="bold green")
    console.print(get_domain_tool_schemas(['chemistry']))