#!/usr/bin/env python3
"""
Quick test script to verify model selection updates

This script tests:
1. Model dataclass can be imported
2. Client has get_models method
3. AsyncClient has get_models method  
4. transcribe method accepts model_name parameter
"""

import sys

def test_imports():
    """Test that all imports work"""
    print("Testing imports...")
    try:
        from orbitalsai import Client, AsyncClient, Model
        from orbitalsai.models import Model as ModelDirect
        print("✅ All imports successful")
        return True
    except ImportError as e:
        print(f"❌ Import failed: {e}")
        return False

def test_model_dataclass():
    """Test Model dataclass"""
    print("\nTesting Model dataclass...")
    try:
        from orbitalsai import Model
        
        model = Model(
            id=1,
            model_name="Perigee-1",
            transcription_rate_per_second=0.000069444,
            transcription_rate_per_hour=0.25,
            is_active=True
        )
        
        assert model.model_name == "Perigee-1"
        assert model.transcription_rate_per_hour == 0.25
        assert model.is_active is True
        
        print("✅ Model dataclass works correctly")
        return True
    except Exception as e:
        print(f"❌ Model dataclass test failed: {e}")
        return False

def test_client_methods():
    """Test that Client has new methods"""
    print("\nTesting Client methods...")
    try:
        from orbitalsai import Client
        import inspect
        
        # Check get_models method exists
        assert hasattr(Client, 'get_models')
        print("  ✅ Client.get_models() method exists")
        
        # Check transcribe accepts model_name
        sig = inspect.signature(Client.transcribe)
        params = sig.parameters
        assert 'model_name' in params
        assert params['model_name'].default == "Perigee-1"
        print("  ✅ Client.transcribe() accepts model_name parameter")
        
        return True
    except Exception as e:
        print(f"❌ Client method test failed: {e}")
        return False

def test_async_client_methods():
    """Test that AsyncClient has new methods"""
    print("\nTesting AsyncClient methods...")
    try:
        from orbitalsai import AsyncClient
        import inspect
        
        # Check get_models method exists
        assert hasattr(AsyncClient, 'get_models')
        print("  ✅ AsyncClient.get_models() method exists")
        
        # Check transcribe accepts model_name
        sig = inspect.signature(AsyncClient.transcribe)
        params = sig.parameters
        assert 'model_name' in params
        assert params['model_name'].default == "Perigee-1"
        print("  ✅ AsyncClient.transcribe() accepts model_name parameter")
        
        return True
    except Exception as e:
        print(f"❌ AsyncClient method test failed: {e}")
        return False

def main():
    """Run all tests"""
    print("=" * 60)
    print("OrbitalsAI SDK - Model Selection Update Tests")
    print("=" * 60)
    
    tests = [
        test_imports,
        test_model_dataclass,
        test_client_methods,
        test_async_client_methods,
    ]
    
    results = []
    for test in tests:
        results.append(test())
    
    print("\n" + "=" * 60)
    print("Test Summary")
    print("=" * 60)
    passed = sum(results)
    total = len(results)
    print(f"Passed: {passed}/{total}")
    
    if passed == total:
        print("\n✅ All tests passed! SDK update is working correctly.")
        return 0
    else:
        print(f"\n❌ {total - passed} test(s) failed. Please review the errors above.")
        return 1

if __name__ == "__main__":
    sys.exit(main())

