#!/usr/bin/env python3
"""
Quick test script to verify the 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'),
        ('huggingface_hub', 'HuggingFace Hub'),
    ]
    
    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_sample_image',
            'downsample_image',
            'quantize_image',
            'apply_sampling_and_quantization',
            'calculate_file_size',
            'compress_image_jpeg',
            'compress_image_png',
            'main_loop',
        ]
        
        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}")
        return False, [str(e)]


def test_image_generation():
    """Test that we can generate a sample image."""
    print("\n🖼️  Testing image generation...")
    
    try:
        import app
        import numpy as np
        
        img = app.generate_sample_image(size=256)
        
        # Check image properties
        assert img.shape == (256, 256, 3), f"Wrong shape: {img.shape}"
        assert img.dtype == np.uint8, f"Wrong dtype: {img.dtype}"
        assert img.min() >= 0 and img.max() <= 255, "Invalid pixel values"
        
        print(f"  ✅ Generated {img.shape[1]}×{img.shape[0]} image")
        return True, []
    
    except Exception as e:
        print(f"  ❌ Error: {e}")
        return False, [str(e)]


def test_image_processing():
    """Test basic image processing functions."""
    print("\n⚙️  Testing image processing...")
    
    try:
        import app
        import numpy as np
        
        # Create a simple test image
        test_img = np.random.randint(0, 256, (128, 128, 3), dtype=np.uint8)
        
        # Test downsampling
        downsampled = app.downsample_image(test_img, sampling_rate=2)
        print(f"  ✅ Downsampling works")
        
        # Test quantization
        quantized = app.quantize_image(test_img, bits_per_pixel=4)
        print(f"  ✅ Quantization works")
        
        # Test combined
        processed = app.apply_sampling_and_quantization(test_img, 2, 4)
        print(f"  ✅ Combined processing works")
        
        # Test file size calculation
        size = app.calculate_file_size(test_img.shape, 2, 4)
        assert size > 0, "File size should be positive"
        print(f"  ✅ File size calculation works ({size} bytes)")
        
        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("🎨 Image Sampling & Quantization Demo - 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: Image generation
    if passed:
        passed, failed = test_image_generation()
        if not passed:
            print(f"\n❌ Image generation test failed.")
            all_passed = False
    
    # Test 4: Image processing
    if passed:
        passed, failed = test_image_processing()
        if not passed:
            print(f"\n❌ Image processing 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: ./run_simple.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())