src/modules/input_module.py

"""
Input Module - Image Upload and Validation
=========================================

Handles image file upload, format validation, and preprocessing
for the image deblurring system.
"""

import cv2
import numpy as np
from PIL import Image
import io
import streamlit as st
from typing import Optional, Tuple, Union
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class ImageValidator:
    """Validates and processes uploaded images"""
    
    SUPPORTED_FORMATS = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'}
    MAX_SIZE_MB = 50
    MIN_RESOLUTION = (100, 100)
    MAX_RESOLUTION = (8192, 8192)
    
    @classmethod
    def validate_format(cls, file) -> bool:
        """Validate if file format is supported"""
        try:
            if hasattr(file, 'name'):
                filename = file.name.lower()
                return any(filename.endswith(fmt) for fmt in cls.SUPPORTED_FORMATS)
            return False
        except Exception as e:
            logger.error(f"Format validation error: {e}")
            return False
    
    @classmethod
    def validate_size(cls, file) -> bool:
        """Validate file size"""
        try:
            if hasattr(file, 'size'):
                size_mb = file.size / (1024 * 1024)
                return size_mb <= cls.MAX_SIZE_MB
            return True
        except Exception as e:
            logger.error(f"Size validation error: {e}")
            return False
    
    @classmethod
    def validate_resolution(cls, image: np.ndarray) -> bool:
        """Validate image resolution"""
        try:
            height, width = image.shape[:2]
            
            # Check minimum resolution
            if width < cls.MIN_RESOLUTION[0] or height < cls.MIN_RESOLUTION[1]:
                return False
            
            # Check maximum resolution
            if width > cls.MAX_RESOLUTION[0] or height > cls.MAX_RESOLUTION[1]:
                return False
            
            return True
        except Exception as e:
            logger.error(f"Resolution validation error: {e}")
            return False

def load_image_from_upload(uploaded_file) -> Optional[np.ndarray]:
    """
    Load and validate image from Streamlit file upload
    
    Args:
        uploaded_file: Streamlit UploadedFile object
    
    Returns:
        np.ndarray: Image as OpenCV format (BGR) or None if invalid
    """
    try:
        # Validate format
        if not ImageValidator.validate_format(uploaded_file):
            st.error("❌ Unsupported file format. Please use: JPG, PNG, BMP, or TIFF")
            return None
        
        # Validate size
        if not ImageValidator.validate_size(uploaded_file):
            st.error(f"❌ File too large. Maximum size: {ImageValidator.MAX_SIZE_MB}MB")
            return None
        
        # Load image
        file_bytes = uploaded_file.getvalue()
        image = Image.open(io.BytesIO(file_bytes))
        
        # Convert to numpy array
        img_array = np.array(image)
        
        # Handle different formats
        if len(img_array.shape) == 3:
            if img_array.shape[2] == 4:  # RGBA
                img_array = cv2.cvtColor(img_array, cv2.COLOR_RGBA2BGR)
            elif img_array.shape[2] == 3:  # RGB
                img_array = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
        elif len(img_array.shape) == 2:  # Grayscale
            img_array = cv2.cvtColor(img_array, cv2.COLOR_GRAY2BGR)
        
        # Validate resolution
        if not ImageValidator.validate_resolution(img_array):
            min_res = ImageValidator.MIN_RESOLUTION
            max_res = ImageValidator.MAX_RESOLUTION
            st.error(f"❌ Invalid resolution. Must be {min_res[0]}x{min_res[1]} to {max_res[0]}x{max_res[1]}")
            return None
        
        logger.info(f"Successfully loaded image: {img_array.shape}")
        return img_array
        
    except Exception as e:
        logger.error(f"Error loading image: {e}")
        st.error(f"❌ Error loading image: {str(e)}")
        return None

def load_image_from_path(image_path: str) -> Optional[np.ndarray]:
    """
    Load image from file path
    
    Args:
        image_path: Path to image file
    
    Returns:
        np.ndarray: Image as OpenCV format (BGR) or None if error
    """
    try:
        image = cv2.imread(image_path)
        if image is None:
            logger.error(f"Could not load image from {image_path}")
            return None
        
        # Validate resolution
        if not ImageValidator.validate_resolution(image):
            logger.error(f"Invalid resolution for image: {image.shape}")
            return None
        
        logger.info(f"Loaded image from path: {image.shape}")
        return image
        
    except Exception as e:
        logger.error(f"Error loading image from path: {e}")
        return None

def preprocess_image(image: np.ndarray, max_size: Tuple[int, int] = (1024, 1024)) -> np.ndarray:
    """
    Preprocess image for processing (resize if needed, normalize)
    
    Args:
        image: Input image
        max_size: Maximum dimensions for processing
    
    Returns:
        np.ndarray: Preprocessed image
    """
    try:
        height, width = image.shape[:2]
        
        # Resize if too large
        if width > max_size[0] or height > max_size[1]:
            # Calculate scale factor
            scale = min(max_size[0] / width, max_size[1] / height)
            new_width = int(width * scale)
            new_height = int(height * scale)
            
            # Resize with high quality
            image = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_LANCZOS4)
            logger.info(f"Resized image to: {image.shape}")
        
        # Ensure image is in correct format
        image = image.astype(np.uint8)
        
        return image
        
    except Exception as e:
        logger.error(f"Error preprocessing image: {e}")
        return image

def validate_and_load_image(uploaded_file, preprocess: bool = True) -> Optional[np.ndarray]:
    """
    Complete image validation and loading pipeline
    
    Args:
        uploaded_file: Streamlit UploadedFile object
        preprocess: Whether to preprocess the image
    
    Returns:
        np.ndarray: Validated and preprocessed image or None
    """
    # Load image
    image = load_image_from_upload(uploaded_file)
    if image is None:
        return None
    
    # Preprocess if requested
    if preprocess:
        image = preprocess_image(image)
    
    return image

def get_image_info(image: np.ndarray) -> dict:
    """
    Get comprehensive image information
    
    Args:
        image: Input image
    
    Returns:
        dict: Image information
    """
    try:
        height, width = image.shape[:2]
        channels = image.shape[2] if len(image.shape) == 3 else 1
        
        return {
            'width': width,
            'height': height,
            'channels': channels,
            'total_pixels': width * height,
            'data_type': str(image.dtype),
            'memory_size_mb': image.nbytes / (1024 * 1024),
            'aspect_ratio': width / height
        }
    except Exception as e:
        logger.error(f"Error getting image info: {e}")
        return {}

# Example usage and testing
if __name__ == "__main__":
    print("Input Module - Image Upload and Validation")
    print("==========================================")
    
    # Test with sample image creation
    test_image = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
    
    # Test validation
    validator = ImageValidator()
    print(f"Resolution validation: {validator.validate_resolution(test_image)}")
    
    # Test preprocessing
    processed = preprocess_image(test_image)
    print(f"Original shape: {test_image.shape}")
    print(f"Processed shape: {processed.shape}")
    
    # Test image info
    info = get_image_info(test_image)
    print(f"Image info: {info}")