#!/usr/bin/env python3
"""
Quick test script to verify the colorspace demo setup
"""

import sys

def test_imports():
    """Test that all required packages can be imported."""
    print("🧪 Testing package imports...")
    
    packages = [
        ('numpy', 'NumPy'),
        ('cv2', 'OpenCV'),
        ('streamlit', 'Streamlit'),
        ('PIL', 'Pillow'),
        ('matplotlib', 'Matplotlib'),
        ('colorsys', 'colorsys'),
    ]
    
    failed = []
    for module, name in packages:
        try:
            __import__(module)
            print(f"  ✅ {name}")
        except ImportError as e:
            print(f"  ❌ {name}: {e}")
            failed.append(name)
    
    return len(failed) == 0, failed


def test_app_structure():
    """Test that app.py has the expected structure."""
    print("\n📁 Testing app structure...")
    
    try:
        import app
        
        # Check for key functions
        functions = [
            'load_image',
            'get_available_images',
            'tab_rgb',
            'tab_hsv',
            'tab_lab',
            'tab_cmyk',
            'tab_ycbcr',
            'tab_gamma_wb',
            'tab_colorblind',
            'simulate_colorblindness',
            'main',
        ]
        
        failed = []
        for func in functions:
            if hasattr(app, func):
                print(f"  ✅ Function '{func}' found")
            else:
                print(f"  ❌ Function '{func}' missing")
                failed.append(func)
        
        return len(failed) == 0, failed
    
    except Exception as e:
        print(f"  ❌ Error loading app: {e}")
        import traceback
        traceback.print_exc()
        return False, [str(e)]


def test_colorspace_conversions():
    """Test basic colorspace conversion functions."""
    print("\n🌈 Testing colorspace conversions...")
    
    try:
        import numpy as np
        import cv2
        
        # Create a simple RGB image
        test_img = np.array([[[255, 0, 0], [0, 255, 0], [0, 0, 255]]], dtype=np.uint8)
        
        # Test RGB to HSV
        hsv_img = cv2.cvtColor(test_img, cv2.COLOR_RGB2HSV)
        print(f"  ✅ RGB to HSV conversion works")
        
        # Test RGB to LAB
        lab_img = cv2.cvtColor(test_img, cv2.COLOR_RGB2LAB)
        print(f"  ✅ RGB to LAB conversion works")
        
        # Test RGB to YCrCb
        ycrcb_img = cv2.cvtColor(test_img, cv2.COLOR_RGB2YCrCb)
        print(f"  ✅ RGB to YCrCb conversion works")
        
        # Test gamma correction
        img_float = test_img.astype(np.float32) / 255.0
        gamma = 2.2
        img_gamma = np.power(img_float, gamma)
        img_gamma = (img_gamma * 255).astype(np.uint8)
        print(f"  ✅ Gamma correction works")
        
        return True, []
    
    except Exception as e:
        print(f"  ❌ Error: {e}")
        import traceback
        traceback.print_exc()
        return False, [str(e)]


def test_colorblindness_simulation():
    """Test color blindness simulation."""
    print("\n👁️  Testing color blindness simulation...")
    
    try:
        import app
        import numpy as np
        
        # Create a simple test image
        test_img = np.random.randint(0, 256, (32, 32, 3), dtype=np.uint8)
        
        # Test different types
        types = [
            "Normal Vision",
            "Protanopia (No Red)",
            "Deuteranopia (No Green)",
            "Tritanopia (No Blue)",
            "Monochromacy (Grayscale)"
        ]
        
        for cb_type in types:
            result = app.simulate_colorblindness(test_img, cb_type)
            assert result.shape == test_img.shape, f"Shape mismatch for {cb_type}"
            print(f"  ✅ {cb_type}")
        
        return True, []
    
    except Exception as e:
        print(f"  ❌ Error: {e}")
        import traceback
        traceback.print_exc()
        return False, [str(e)]


def test_image_loading():
    """Test that images can be loaded from the images folder."""
    print("\n🖼️  Testing image loading...")
    
    try:
        from pathlib import Path
        import app
        
        # Check if images folder exists
        images_dir = Path("images")
        if not images_dir.exists():
            print("  ⚠️  images/ folder not found (this is OK for fresh setup)")
            return True, []
        
        # Get available images
        images = app.get_available_images()
        if images:
            print(f"  ✅ Found {len(images)} images in images/ folder")
            
            # Try to load one
            img = app.load_image(images[0])
            if img is not None:
                print(f"  ✅ Successfully loaded '{images[0]}'")
            else:
                print(f"  ⚠️  Could not load '{images[0]}' (may be corrupted)")
        else:
            print("  ⚠️  No images found in images/ folder (this is OK for fresh setup)")
        
        return True, []
    
    except Exception as e:
        print(f"  ❌ Error: {e}")
        import traceback
        traceback.print_exc()
        return False, [str(e)]


def main():
    """Run all tests."""
    print("=" * 60)
    print("🎨 Colorspace Explorer - Test Suite")
    print("=" * 60)
    
    all_passed = True
    
    # Test 1: Imports
    passed, failed = test_imports()
    if not passed:
        print(f"\n❌ Import test failed. Missing packages: {', '.join(failed)}")
        print("\n💡 Run: pip install -r requirements.txt")
        all_passed = False
    
    # Test 2: App structure
    if passed:  # Only run if imports work
        passed, failed = test_app_structure()
        if not passed:
            print(f"\n❌ App structure test failed.")
            all_passed = False
    
    # Test 3: Colorspace conversions
    if passed:
        passed, failed = test_colorspace_conversions()
        if not passed:
            print(f"\n❌ Colorspace conversion test failed.")
            all_passed = False
    
    # Test 4: Color blindness simulation
    if passed:
        passed, failed = test_colorblindness_simulation()
        if not passed:
            print(f"\n❌ Color blindness simulation test failed.")
            all_passed = False
    
    # Test 5: Image loading
    if passed:
        passed, failed = test_image_loading()
        if not passed:
            print(f"\n❌ Image loading test failed.")
            all_passed = False
    
    # Summary
    print("\n" + "=" * 60)
    if all_passed:
        print("✅ All tests passed! Ready to run the demo.")
        print("\n🚀 Start with:")
        print("   ./run.sh")
        print("   Or: make run")
        print("   Or: streamlit run app.py")
        return 0
    else:
        print("❌ Some tests failed. Please fix the issues above.")
        return 1


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