#!/usr/bin/env python3
"""
Test script for the Universal Model Toggle System

This script demonstrates how the new universal model toggle system works
with any model, not just hardcoded ones.
"""

import sys
import os
import logging
from datetime import datetime

# Add the project root to the path
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def test_universal_model_toggles():
    """Test the universal model toggle system"""
    try:
        from core.orchestrator import TradingOrchestrator
        from core.data_provider import DataProvider
        from models import ModelInterface, get_model_registry
        
        logger.info("🧪 Testing Universal Model Toggle System")
        
        # Initialize components
        data_provider = DataProvider()
        orchestrator = TradingOrchestrator(data_provider=data_provider)
        
        # Test 1: Check existing models
        logger.info("\n📋 Test 1: Checking existing models")
        existing_models = orchestrator.get_all_registered_models()
        logger.info(f"Found {len(existing_models)} existing models: {list(existing_models.keys())}")
        
        # Test 2: Add a new model dynamically
        logger.info("\n➕ Test 2: Adding new model dynamically")
        
        class TestModel(ModelInterface):
            def __init__(self):
                super().__init__("test_model")
            
            def predict(self, data):
                return {"action": "TEST", "confidence": 0.85}
        
        test_model = TestModel()
        success = orchestrator.register_model_dynamically("test_model", test_model)
        logger.info(f"Dynamic model registration: {'✅ SUCCESS' if success else '❌ FAILED'}")
        
        # Test 3: Check toggle states
        logger.info("\n🔄 Test 3: Testing toggle states")
        
        # Test with existing model
        dqn_state = orchestrator.get_model_toggle_state("dqn")
        logger.info(f"DQN toggle state: {dqn_state}")
        
        # Test with new model
        test_model_state = orchestrator.get_model_toggle_state("test_model")
        logger.info(f"Test model toggle state: {test_model_state}")
        
        # Test 4: Update toggle states
        logger.info("\n⚙️ Test 4: Updating toggle states")
        
        # Disable inference for test model
        orchestrator.set_model_toggle_state("test_model", inference_enabled=False)
        updated_state = orchestrator.get_model_toggle_state("test_model")
        logger.info(f"Updated test model state: {updated_state}")
        
        # Test 5: Add another model without interface
        logger.info("\n➕ Test 5: Adding model without interface")
        orchestrator.set_model_toggle_state("custom_transformer", inference_enabled=True, training_enabled=True)
        transformer_state = orchestrator.get_model_toggle_state("custom_transformer")
        logger.info(f"Custom transformer state: {transformer_state}")
        
        # Test 6: Check all models after additions
        logger.info("\n📋 Test 6: Final model count")
        final_models = orchestrator.get_all_registered_models()
        logger.info(f"Final model count: {len(final_models)}")
        for model_name, model_info in final_models.items():
            toggle_state = orchestrator.get_model_toggle_state(model_name)
            logger.info(f"  - {model_name}: inf={toggle_state['inference_enabled']}, train={toggle_state['training_enabled']}")
        
        logger.info("\n✅ Universal Model Toggle System test completed successfully!")
        return True
        
    except Exception as e:
        logger.error(f"❌ Test failed: {e}")
        return False

def test_dashboard_integration():
    """Test dashboard integration with universal toggles"""
    try:
        logger.info("\n🖥️ Testing Dashboard Integration")
        
        from web.clean_dashboard import CleanTradingDashboard
        from core.orchestrator import TradingOrchestrator
        from core.data_provider import DataProvider
        
        # Initialize components
        data_provider = DataProvider()
        orchestrator = TradingOrchestrator(data_provider=data_provider)
        
        # Add some test models
        orchestrator.set_model_toggle_state("test_model_1", inference_enabled=True, training_enabled=False)
        orchestrator.set_model_toggle_state("test_model_2", inference_enabled=False, training_enabled=True)
        
        # Initialize dashboard (this will test the universal callback setup)
        dashboard = CleanTradingDashboard(
            data_provider=data_provider,
            orchestrator=orchestrator
        )
        
        # Test adding model dynamically through dashboard
        success = dashboard.add_model_dynamically("dynamic_test_model")
        logger.info(f"Dashboard dynamic model addition: {'✅ SUCCESS' if success else '❌ FAILED'}")
        
        # Check available models
        available_models = dashboard._get_available_models()
        logger.info(f"Dashboard sees {len(available_models)} models: {list(available_models.keys())}")
        
        logger.info("✅ Dashboard integration test completed!")
        return True
        
    except Exception as e:
        logger.error(f"❌ Dashboard integration test failed: {e}")
        return False

if __name__ == "__main__":
    logger.info("🚀 Starting Universal Model Toggle System Tests")
    logger.info("=" * 60)
    
    # Run tests
    test1_success = test_universal_model_toggles()
    test2_success = test_dashboard_integration()
    
    # Summary
    logger.info("\n" + "=" * 60)
    logger.info("📊 TEST SUMMARY")
    logger.info(f"Universal Toggle System: {'✅ PASS' if test1_success else '❌ FAIL'}")
    logger.info(f"Dashboard Integration: {'✅ PASS' if test2_success else '❌ FAIL'}")
    
    if test1_success and test2_success:
        logger.info("🎉 ALL TESTS PASSED! Universal model toggle system is working correctly.")
        sys.exit(0)
    else:
        logger.error("❌ Some tests failed. Check the logs above for details.")
        sys.exit(1)