#!/usr/bin/env python3
"""
VIGNETTE CORRECTION NUITKA COMPATIBILITY TEST

Tests the multi-location flatFields search logic that we implemented
for Nuitka compatibility in mip/Vignette_Correction.py
"""

import sys
import os
import pathlib

def test_vignette_correction():
    print("🧪 VIGNETTE CORRECTION NUITKA COMPATIBILITY TEST")
    print("=" * 55)
    
    # Test 1: Import Test
    print("\n🔍 TEST 1: Import Test")
    try:
        from mip.Vignette_Correction import ApplyVig
        print("✅ Vignette_Correction.ApplyVig imported successfully")
    except ImportError as e:
        print(f"❌ Import failed: {e}")
        return False
    
    # Test 2: flatFields Directory Detection
    print("\n🔍 TEST 2: flatFields Directory Detection Test")
    
    # Check if flatFields exists in the expected location
    current_dir = pathlib.Path(os.getcwd())
    flatfields_path = current_dir / "flatFields"
    
    if flatfields_path.exists():
        print(f"✅ flatFields directory found at: {flatfields_path}")
        
        # Check subdirectories
        s3w_path = flatfields_path / "S3W"
        s3n_path = flatfields_path / "S3N"
        
        if s3w_path.exists():
            print(f"✅ S3W subdirectory found: {s3w_path}")
            # List files in S3W
            s3w_files = list(s3w_path.glob("*.tif"))
            print(f"✅ S3W contains {len(s3w_files)} .tif files")
            for file in s3w_files:
                print(f"    - {file.name}")
        else:
            print(f"❌ S3W subdirectory missing: {s3w_path}")
            
        if s3n_path.exists():
            print(f"✅ S3N subdirectory found: {s3n_path}")
            # List files in S3N
            s3n_files = list(s3n_path.glob("*.tif"))
            print(f"✅ S3N contains {len(s3n_files)} .tif files")
            for file in s3n_files:
                print(f"    - {file.name}")
        else:
            print(f"❌ S3N subdirectory missing: {s3n_path}")
            
    else:
        print(f"❌ flatFields directory not found at: {flatfields_path}")
        return False
    
    # Test 3: Simulate Nuitka Path Detection Logic
    print("\n🔍 TEST 3: Nuitka Path Detection Logic Test")
    
    # Test our multi-location search logic
    print("  Testing the multi-location search paths...")
    
    possible_paths = []
    
    # For PyInstaller (has _MEIPASS) - won't exist in development
    if hasattr(sys, '_MEIPASS'):
        possible_paths.append(pathlib.Path(sys._MEIPASS, "flatFields"))
        print("  ✅ sys._MEIPASS detected (PyInstaller mode)")
    else:
        print("  ✅ No sys._MEIPASS (development mode - expected)")
    
    # For Nuitka onefile (extracts to temp directory near executable)
    executable_dir = pathlib.Path(sys.executable).parent
    possible_paths.extend([
        executable_dir / "flatFields",  # Next to executable
        executable_dir.parent / "flatFields",  # One level up
        pathlib.Path(os.getcwd()) / "flatFields",  # Current working directory
        # Nuitka may extract to a subdirectory
        executable_dir / "backend_server.dist" / "flatFields",
    ])
    
    print(f"  Testing {len(possible_paths)} possible paths:")
    
    flatfields_base = None
    for i, path in enumerate(possible_paths, 1):
        exists = path.exists() and path.is_dir()
        print(f"    {i}. {path} - {'✅ EXISTS' if exists else '❌ not found'}")
        if exists and flatfields_base is None:
            flatfields_base = path
    
    if flatfields_base:
        print(f"  ✅ flatFields would be found at: {flatfields_base}")
    else:
        print(f"  ❌ flatFields would NOT be found in any location!")
        return False
    
    # Test 4: Mock Image Object Test
    print("\n🔍 TEST 4: Mock Image Object Test")
    
    # Create a mock image object to test the logic
    class MockImage:
        def __init__(self, camera_model):
            self.camera_model = camera_model
    
    # Test different camera models
    test_models = ['Survey3W', 'Survey3N', 'Unknown', None]
    
    for model in test_models:
        print(f"  Testing camera model: {model}")
        mock_image = MockImage(model)
        
        # Simulate the path selection logic from the actual function
        if getattr(sys, 'frozen', False):
            # This won't trigger in development, but we can test the logic
            print(f"    (In Nuitka mode, would use multi-location search)")
        else:
            # Development mode logic
            base_path = pathlib.Path(__file__).parent.parent
            if mock_image.camera_model == 'Survey3W':
                corr_path = str(pathlib.Path(base_path, "flatFields", "S3W"))
                camera_model = 'Survey3W'
            elif mock_image.camera_model == 'Survey3N':
                corr_path = str(pathlib.Path(base_path, "flatFields", "S3N"))
                camera_model = 'Survey3N'
            else:
                # Handle Unknown or other camera models - default to Survey3W
                corr_path = str(pathlib.Path(base_path, "flatFields", "S3W"))
                camera_model = 'Survey3W'
                
            print(f"    ✅ Would use path: {corr_path}")
            print(f"    ✅ Camera model: {camera_model}")
    
    print("\n🎉 ALL VIGNETTE CORRECTION TESTS PASSED!")
    print("✅ Nuitka flatFields compatibility is working correctly")
    return True

if __name__ == "__main__":
    success = test_vignette_correction()
    if success:
        print("\n🚀 Vignette Correction is ready for Nuitka compilation!")
    else:
        print("\n❌ Vignette Correction needs fixes before compilation!")
    
    input("\nPress Enter to continue...")
