#!/usr/bin/env python3
"""
Quick test script to verify the edge detection 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',
            'generate_sample_image',
            'apply_sobel_filter',
            'apply_prewitt_filter',
            'apply_roberts_filter',
            'apply_laplacian_filter',
            'apply_canny_edge_detector',
            'compute_gradient_direction',
            'create_gradient_visualization',
            '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_edge_detection():
    """Test basic edge detection functions."""
    print("\n⚙️  Testing edge detection...")
    
    try:
        import app
        import numpy as np
        
        # Create a simple test image
        test_img = np.random.randint(0, 256, (128, 128), dtype=np.uint8)
        
        # Test Sobel
        edges, gx, gy = app.apply_sobel_filter(test_img, ksize=3)
        assert edges.shape == test_img.shape, "Sobel output shape mismatch"
        print(f"  ✅ Sobel filter works")
        
        # Test Prewitt
        edges, gx, gy = app.apply_prewitt_filter(test_img)
        assert edges.shape == test_img.shape, "Prewitt output shape mismatch"
        print(f"  ✅ Prewitt filter works")
        
        # Test Roberts
        edges, gx, gy = app.apply_roberts_filter(test_img)
        print(f"  ✅ Roberts filter works")
        
        # Test Laplacian
        edges = app.apply_laplacian_filter(test_img, ksize=3)
        assert edges.shape == test_img.shape, "Laplacian output shape mismatch"
        print(f"  ✅ Laplacian filter works")
        
        # Test Canny
        edges = app.apply_canny_edge_detector(test_img, 50, 150)
        assert edges.shape == test_img.shape, "Canny output shape mismatch"
        print(f"  ✅ Canny edge detector works")
        
        # Test gradient visualization
        viz = app.create_gradient_visualization(gx, gy)
        print(f"  ✅ Gradient visualization works")
        
        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("🔍 Edge Detection 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: Edge detection
    if passed:
        passed, failed = test_edge_detection()
        if not passed:
            print(f"\n❌ Edge detection 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())