#!/usr/bin/env python3
"""
Test script for Databricks MCP Server tools
Tests each tool to verify they work correctly
"""

import asyncio
import os
import json
import sys

# Add the src directory to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

from databricks_mcp_server.databricks_client import (
    who_am_i,
    list_catalogs,
    list_schemas,
    list_tables,
    get_table_details,
    get_sample_data
)


async def test_who_am_i():
    """Test the who_am_i tool."""
    print("\n" + "="*60)
    print("Testing: who_am_i")
    print("="*60)
    
    result = await who_am_i()
    print(result)
    return json.loads(result) if not result.startswith("❌") else None


async def test_list_catalogs():
    """Test the list_catalogs tool."""
    print("\n" + "="*60)
    print("Testing: list_catalogs")
    print("="*60)
    
    result = await list_catalogs()
    print(result)
    return json.loads(result) if not result.startswith("❌") else None


async def test_list_schemas(catalog_name: str):
    """Test the list_schemas tool."""
    print("\n" + "="*60)
    print(f"Testing: list_schemas (catalog: {catalog_name})")
    print("="*60)
    
    result = await list_schemas(catalog_name)
    print(result)
    return json.loads(result) if not result.startswith("❌") else None


async def test_list_tables(catalog_name: str, schema_name: str):
    """Test the list_tables tool."""
    print("\n" + "="*60)
    print(f"Testing: list_tables (catalog.schema: {catalog_name}.{schema_name})")
    print("="*60)
    
    result = await list_tables(catalog_name, schema_name)
    print(result)
    return json.loads(result) if not result.startswith("❌") else None


async def test_get_table_details(catalog_name: str, schema_name: str, table_name: str):
    """Test the get_table_details tool."""
    print("\n" + "="*60)
    print(f"Testing: get_table_details ({catalog_name}.{schema_name}.{table_name})")
    print("="*60)
    
    result = await get_table_details(catalog_name, schema_name, table_name)
    print(result)
    return json.loads(result) if not result.startswith("❌") else None


async def test_get_sample_data(catalog_name: str, schema_name: str, table_name: str, limit: int = 5):
    """Test the get_sample_data tool."""
    print("\n" + "="*60)
    print(f"Testing: get_sample_data ({catalog_name}.{schema_name}.{table_name}, limit={limit})")
    print("="*60)
    
    result = await get_sample_data(catalog_name, schema_name, table_name, limit)
    print(result)
    return json.loads(result) if not result.startswith("❌") else None


async def main():
    """Run all tests."""
    print("="*60)
    print("Databricks MCP Server - Tool Tests")
    print("="*60)
    
    # Set the profile
    profile = os.getenv("DATABRICKS_PROFILE", "mareana_vsm")
    print(f"\nUsing Databricks profile: {profile}")
    os.environ["DATABRICKS_PROFILE"] = profile
    
    # Test 1: who_am_i
    user = await test_who_am_i()
    if not user:
        print("\n❌ Authentication failed. Please run: databricks auth login --profile " + profile)
        return
    
    print(f"\n✅ Authenticated as: {user.get('username')}")
    
    # Test 2: list_catalogs
    catalogs_result = await test_list_catalogs()
    if not catalogs_result or 'error' in catalogs_result:
        print("\n❌ Failed to list catalogs")
        return
    
    catalogs = catalogs_result.get('catalogs', [])
    if not catalogs:
        print("\n⚠️ No catalogs found")
        return
    
    print(f"\n✅ Found {len(catalogs)} catalogs")
    
    # Use the first catalog or 'mareana_vsm' if available
    test_catalog = next((c['name'] for c in catalogs if c['name'] == 'mareana_vsm'), catalogs[0]['name'])
    print(f"   Using catalog: {test_catalog}")
    
    # Test 3: list_schemas
    schemas_result = await test_list_schemas(test_catalog)
    if not schemas_result or 'error' in schemas_result:
        print("\n❌ Failed to list schemas")
        return
    
    schemas = schemas_result.get('schemas', [])
    if not schemas:
        print("\n⚠️ No schemas found")
        return
    
    print(f"\n✅ Found {len(schemas)} schemas")
    
    # Use 'harmonized' schema if available, otherwise first schema
    test_schema = next((s['name'] for s in schemas if s['name'] == 'harmonized'), schemas[0]['name'])
    print(f"   Using schema: {test_schema}")
    
    # Test 4: list_tables
    tables_result = await test_list_tables(test_catalog, test_schema)
    if not tables_result or 'error' in tables_result:
        print("\n❌ Failed to list tables")
        return
    
    tables = tables_result.get('tables', [])
    if not tables:
        print("\n⚠️ No tables found")
        return
    
    print(f"\n✅ Found {len(tables)} tables")
    
    # Use first table
    test_table = tables[0]['name']
    print(f"   Using table: {test_table}")
    
    # Test 5: get_table_details
    details_result = await test_get_table_details(test_catalog, test_schema, test_table)
    if not details_result or 'error' in details_result:
        print("\n❌ Failed to get table details")
    else:
        print(f"\n✅ Got table details: {details_result.get('column_count')} columns")
    
    # Test 6: get_sample_data
    sample_result = await test_get_sample_data(test_catalog, test_schema, test_table, limit=5)
    if not sample_result or 'error' in sample_result:
        print(f"\n⚠️ Could not get sample data: {sample_result.get('message', 'Unknown error')}")
    else:
        print(f"\n✅ Got {sample_result.get('row_count')} sample rows")
    
    print("\n" + "="*60)
    print("All tests completed!")
    print("="*60)


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