from sales_assistant.agent_tools.describe_table import describe_table
from sales_assistant.agent_tools.get_samples_data import get_sample_data
from sales_assistant.agent_tools.execute_advanced_query import execute_advanced_query
from sales_assistant.agent_tools.get_distinct_values import get_distinct_values
from sales_assistant.agent_tools.get_product_by_criteria import search_products_by_criteria
from sales_assistant.agent_tools.get_table_statistics import get_table_statistics

def test_describe_table_tool():
    """
    Test the describe_table tool to ensure it retrieves schema information correctly.
    """
    print("\n=== Testing describe_table tool ===")
    
    # Test default query
    print("Test 1: Default DESCRIBE query")
    # Since describe_table is a LangChain tool, we need to invoke it properly
    result = describe_table.invoke({})
    print("Result:")
    print(result)
    
    # Test with custom query
    print("\nTest 2: Custom query")
    custom_query = "SHOW TABLES FROM streamnet;"
    result2 = describe_table.invoke({"sql_query": custom_query})
    print("Result:")
    print(result2)

def test_get_sample_data_tool():
    """
    Test the get_sample_data tool to ensure it retrieves sample data correctly with different methods.
    """
    print("\n=== Testing get_sample_data tool ===")
    
    # Test default method (limit with default sample_size=5)
    print("Test 1: Default method (limit, sample_size=5)")
    result = get_sample_data.invoke({})
    print("Result:")
    print(result)
    
    # Test limit method with custom sample_size
    print("Test 2: Limit method with sample_size=3")
    result2 = get_sample_data.invoke({
        "sql_query": "SELECT * FROM streamnet.products_list WHERE category = 'camera'",
        "method": "limit",
        "sample_size": 3
    })
    print("Result:")
    print(result2)
    
    # Test random method
    print("Test 3: Random method with sample_size=2")
    result3 = get_sample_data.invoke({
        "sql_query": "SELECT manufacturer, model_name FROM streamnet.products_list",
        "method": "random",
        "sample_size": 2
    })
    print("Result:")
    print(result3)
    
    # Test None method (all results)
    print("Test 4: None method (all results)")
    result4 = get_sample_data.invoke({
        "sql_query": "SELECT manufacturer, COUNT(*) as count FROM streamnet.products_list GROUP BY manufacturer",
        "method": None,
        "sample_size": 10  # This is ignored
    })
    print("Result:")
    print(result4)
    
    # Test "none" string method (alternative way)
    print("Test 5: 'none' string method (all results)")
    result5 = get_sample_data.invoke({
        "sql_query": "SELECT category, COUNT(*) as count FROM streamnet.products_list GROUP BY category",
        "method": "none",
        "sample_size": 5 # This is ignored
    })
    print("Result:")
    print(result5)

def test_execute_advanced_query_tool():
    """
    Test the execute_advanced_query tool to ensure it executes complex SQL queries correctly.
    """
    print("\n=== Testing execute_advanced_query tool ===")
    
    # Test basic advanced query
    print("Test 1: Basic advanced query with aggregation")
    query1 = "SELECT manufacturer, COUNT(*) as product_count, AVG(msrp) as avg_price FROM streamnet.products_list WHERE msrp > 0 GROUP BY manufacturer LIMIT 5"
    result1 = execute_advanced_query.invoke({"sql_query": query1})
    print("Result:")
    print(result1)
    
    # Test complex query with joins/subqueries concept
    print("Test 2: Complex query with filtering")
    query2 = "SELECT manufacturer, category, model_name, msrp FROM streamnet.products_list WHERE msrp = (SELECT MAX(msrp) FROM streamnet.products_list WHERE manufacturer = 'Apple') LIMIT 3"
    result2 = execute_advanced_query.invoke({"sql_query": query2})
    print("Result:")
    print(result2)

def test_get_distinct_values_tool():
    """
    Test the get_distinct_values tool to ensure it retrieves distinct values and patterns correctly.
    """
    print("\n=== Testing get_distinct_values tool ===")
    
    # Test default query
    print("Test 1: Default query (manufacturers by count)")
    result1 = get_distinct_values.invoke({})
    print("Result:")
    print(result1)
    
    # Test custom query for categories
    print("Test 2: Custom query for categories")
    query2 = "SELECT category, COUNT(*) as count FROM streamnet.products_list GROUP BY category ORDER BY count DESC LIMIT 5"
    result2 = get_distinct_values.invoke({"sql_query": query2})
    print("Result:")
    print(result2)
    
    # Test distinct model names
    print("Test 3: Distinct model names for specific category")
    query3 = "SELECT DISTINCT model_name FROM streamnet.products_list WHERE category = 'camera' ORDER BY model_name LIMIT 5"
    result3 = get_distinct_values.invoke({"sql_query": query3})
    print("Result:")
    print(result3)

def test_search_products_by_criteria_tool():
    """
    Test the search_products_by_criteria tool to ensure it searches products correctly.
    """
    print("\n=== Testing search_products_by_criteria tool ===")
    
    # Test default query
    print("Test 1: Default query (Samsung products)")
    result1 = search_products_by_criteria.invoke({})
    print("Result:")
    print(result1)
    
    # Test custom search with price range
    print("Test 2: Custom search with price range")
    query2 = "SELECT * FROM streamnet.products_list WHERE msrp BETWEEN 100 AND 500 AND manufacturer = 'Apple' LIMIT 5"
    result2 = search_products_by_criteria.invoke({"sql_query": query2})
    print("Result:")
    print(result2)
    
    # Test search with text matching
    print("Test 3: Search with text matching")
    query3 = "SELECT manufacturer, model_name, msrp FROM streamnet.products_list WHERE model_name LIKE '%Pro%' ORDER BY msrp DESC LIMIT 5"
    result3 = search_products_by_criteria.invoke({"sql_query": query3})
    print("Result:")
    print(result3)

def test_get_table_statistics_tool():
    """
    Test the get_table_statistics tool to ensure it generates statistics correctly.
    """
    print("\n=== Testing get_table_statistics tool ===")
    
    # Test default query
    print("Test 1: Default query (basic table statistics)")
    result1 = get_table_statistics.invoke({})
    print("Result:")
    print(result1)
    
    # Test price statistics
    print("Test 2: Price statistics")
    query2 = "SELECT COUNT(*) as total_products, AVG(msrp) as avg_price, MIN(msrp) as min_price, MAX(msrp) as max_price FROM streamnet.products_list WHERE msrp > 0"
    result2 = get_table_statistics.invoke({"sql_query": query2})
    print("Result:")
    print(result2)
    
    # Test manufacturer statistics
    print("Test 3: Manufacturer statistics")
    query3 = "SELECT manufacturer, COUNT(*) as product_count, AVG(msrp) as avg_price FROM streamnet.products_list WHERE msrp > 0 GROUP BY manufacturer ORDER BY product_count DESC LIMIT 5"
    result3 = get_table_statistics.invoke({"sql_query": query3})
    print("Result:")
    print(result3)

if __name__ == "__main__":
    test_describe_table_tool()
    test_get_sample_data_tool()
    test_execute_advanced_query_tool()
    test_get_distinct_values_tool() 
    test_search_products_by_criteria_tool()
    test_get_table_statistics_tool()