#!/usr/bin/env python3
"""
Test script to validate that Lightning LoRA is always loaded as base model
"""
import sys
import os

# Add the current directory to the Python path
sys.path.insert(0, '/config/workspace/hf/Qwen-Image-Edit-2509-Turbo-Lightning')

def test_lightning_always_on():
    """Test that Lightning LoRA is configured as always-loaded base"""
    print("Testing Lightning LoRA always-on configuration...")
    
    try:
        # Read the app.py file
        with open('/config/workspace/hf/Qwen-Image-Edit-2509-Turbo-Lightning/app.py', 'r') as f:
            content = f.read()
        
        # Check Lightning LoRA configuration
        if 'Lightning (4-Step)' not in content:
            print("❌ Lightning LoRA not found in configuration")
            return False
        
        print("✅ Found Lightning LoRA in configuration")
        
        # Check for always_load flag
        if '"always_load": True' not in content:
            print("❌ Lightning LoRA missing always_load flag")
            return False
        
        print("✅ Lightning LoRA has always_load flag")
        
        # Check for Lightning-specific loading logic
        lightning_loading_patterns = [
            'print(f"Loading always-active Lightning LoRA: {LIGHTNING_LORA_NAME}")',
            'lora_manager.load_lora(LIGHTNING_LORA_NAME)',
            'lora_manager.fuse_lora(LIGHTNING_LORA_NAME)',
            'Lightning will remain active'
        ]
        
        for pattern in lightning_loading_patterns:
            if pattern not in content:
                print(f"❌ Missing Lightning loading pattern: {pattern}")
                return False
            print(f"✅ Found Lightning loading pattern: {pattern}")
        
        # Check for multi-LoRA combination support
        if 'adapter_names' not in content:
            print("⚠️  Multi-LoRA combination not found (this might be expected)")
        else:
            print("✅ Multi-LoRA combination supported")
        
        # Check UI updates to reflect always-on Lightning
        if 'Lightning LoRA always active' not in content:
            print("❌ Missing UI indication of Lightning always-on")
            return False
        
        print("✅ UI shows Lightning LoRA always active")
        
        print("✅ Lightning LoRA always-on test passed!")
        return True
        
    except Exception as e:
        print(f"❌ Lightning LoRA test failed: {e}")
        return False

def test_configuration_structure():
    """Test that LoRA configurations are properly structured"""
    print("\nTesting LoRA configuration structure...")
    
    try:
        # Read the app.py file
        with open('/config/workspace/hf/Qwen-Image-Edit-2509-Turbo-Lightning/app.py', 'r') as f:
            content = f.read()
        
        # Check for proper configuration structure
        required_configs = [
            '"Lightning (4-Step)"',
            '"repo_id": "lightx2v/Qwen-Image-Lightning"',
            '"type": "base"',
            '"method": "standard"',
            '"Qwen-Image-Lightning-4steps-V2.0.safetensors"'
        ]
        
        for config in required_configs:
            if config not in content:
                print(f"❌ Missing configuration: {config}")
                return False
            print(f"✅ Found configuration: {config}")
        
        print("✅ Configuration structure test passed!")
        return True
        
    except Exception as e:
        print(f"❌ Configuration structure test failed: {e}")
        return False

def test_inference_flow():
    """Test that inference flow preserves Lightning LoRA"""
    print("\nTesting inference flow with Lightning always-on...")
    
    try:
        # Read the app.py file
        with open('/config/workspace/hf/Qwen-Image-Edit-2509-Turbo-Lightning/app.py', 'r') as f:
            content = f.read()
        
        # Check for Lightning preservation in inference
        inference_patterns = [
            'if lora_name == LIGHTNING_LORA_NAME:',
            'Lightning LoRA is already active.',
            'pipe.set_adapters([LIGHTNING_LORA_NAME])',
            "print(f\"LoRA: {lora_name} (with Lightning always active)\")",
            "Don't unfuse Lightning",
            'pipe.disable_adapters()'
        ]
        
        for pattern in inference_patterns:
            if pattern not in content:
                print(f"❌ Missing inference pattern: {pattern}")
                return False
            print(f"✅ Found inference pattern: {pattern}")
        
        print("✅ Inference flow test passed!")
        return True
        
    except Exception as e:
        print(f"❌ Inference flow test failed: {e}")
        return False

def test_loading_sequence():
    """Test the Lightning LoRA loading sequence"""
    print("\nTesting Lightning LoRA loading sequence...")
    
    try:
        # Read the app.py file
        with open('/config/workspace/hf/Qwen-Image-Edit-2509-Turbo-Lightning/app.py', 'r') as f:
            content = f.read()
        
        # Check for proper loading sequence
        sequence_patterns = [
            'print(f"Loading always-active Lightning LoRA: {LIGHTNING_LORA_NAME}")',
            'lightning_config = LORA_CONFIG[LIGHTNING_LORA_NAME]',
            'hf_hub_download(',
            'repo_id=lightning_config["repo_id"],',
            'filename=lightning_config["filename"]',
            'lora_manager.register_lora(LIGHTNING_LORA_NAME, lightning_lora_path, **lightning_config)',
            'lora_manager.load_lora(LIGHTNING_LORA_NAME)',
            'lora_manager.fuse_lora(LIGHTNING_LORA_NAME)'
        ]
        
        for pattern in sequence_patterns:
            if pattern not in content:
                print(f"❌ Missing loading sequence: {pattern}")
                return False
            print(f"✅ Found loading sequence: {pattern}")
        
        print("✅ Loading sequence test passed!")
        return True
        
    except Exception as e:
        print(f"❌ Loading sequence test failed: {e}")
        return False

def main():
    """Run all Lightning LoRA tests"""
    print("=" * 60)
    print("Lightning LoRA Always-On Validation")
    print("=" * 60)
    
    tests = [
        test_lightning_always_on,
        test_configuration_structure,
        test_inference_flow,
        test_loading_sequence
    ]
    
    passed = 0
    failed = 0
    
    for test in tests:
        try:
            if test():
                passed += 1
            else:
                failed += 1
        except Exception as e:
            print(f"❌ {test.__name__} failed with exception: {e}")
            failed += 1
    
    print("\n" + "=" * 60)
    print(f"Lightning LoRA Test Results: {passed} passed, {failed} failed")
    print("=" * 60)
    
    if failed == 0:
        print("🎉 All Lightning LoRA tests passed!")
        print("\nKey Lightning Features Verified:")
        print("✅ Lightning LoRA configured as always-loaded base")
        print("✅ Lightning LoRA loaded and fused on startup")
        print("✅ Inference preserves Lightning LoRA state")
        print("✅ Multi-LoRA combination supported")
        print("✅ UI indicates Lightning always active")
        print("✅ Proper loading sequence implemented")
        return True
    else:
        print("⚠️ Some Lightning LoRA tests failed.")
        return False

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)