import numpy as np
from PIL import Image
import io
from typing import Tuple, Dict, Union
import asyncio
from concurrent.futures import ThreadPoolExecutor

# Optimized thread pool for CPU-intensive image operations
image_executor = ThreadPoolExecutor(max_workers=2, thread_name_prefix="ImageProcessor")

# Common model input sizes
MODEL_SIZES = {
    'xceptionnet': (299, 299),
    'mesonet': (256, 256)
}


def _preprocess_image_core(image_data: bytes, target_size: Tuple[int, int]) -> np.ndarray:
    """
    Core image preprocessing function (synchronous).
    
    Args:
        image_data: Raw image bytes
        target_size: Target size as (width, height)
    
    Returns:
        Preprocessed numpy array ready for model prediction
    """
    # Read and validate image
    image = Image.open(io.BytesIO(image_data))
    
    # Convert to RGB if needed
    if image.mode != 'RGB':
        image = image.convert('RGB')
    
    # Resize and normalize
    image = image.resize(target_size, Image.Resampling.LANCZOS)
    img_array = np.array(image, dtype=np.float32) / 255.0
    
    # Add batch dimension
    return np.expand_dims(img_array, axis=0)


async def preprocess_image(image_data: bytes, target_size: Tuple[int, int]) -> np.ndarray:
    """
    Async image preprocessing for FastAPI.
    
    Args:
        image_data: Raw image bytes from uploaded file
        target_size: Target size as (width, height)
    
    Returns:
        Preprocessed numpy array ready for model prediction
    """
    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(image_executor, _preprocess_image_core, image_data, target_size)


async def preprocess_for_models(image_data: bytes) -> Dict[str, np.ndarray]:
    """
    Preprocess image for all models simultaneously (optimized).
    
    Args:
        image_data: Raw image bytes from uploaded file
    
    Returns:
        Dictionary with preprocessed arrays for each model
    """
    def _process_all_sizes(data: bytes) -> Dict[str, np.ndarray]:
        # Load image once
        image = Image.open(io.BytesIO(data))
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        results = {}
        for model_name, size in MODEL_SIZES.items():
            # Resize and normalize
            resized = image.resize(size, Image.Resampling.LANCZOS)
            img_array = np.array(resized, dtype=np.float32) / 255.0
            results[model_name] = np.expand_dims(img_array, axis=0)
        
        return results
    
    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(image_executor, _process_all_sizes, image_data)


def validate_image(image_data: bytes) -> bool:
    """
    Validate if the provided bytes represent a valid image.
    
    Args:
        image_data: Raw image bytes
    
    Returns:
        True if valid image, False otherwise
    """
    try:
        with Image.open(io.BytesIO(image_data)) as img:
            img.verify()
        return True
    except Exception:
        return False


# Legacy compatibility function (for any remaining Flask code)
def preprocess_image_sync(file_or_bytes: Union[bytes, object], target_size: Tuple[int, int]) -> np.ndarray:
    """Legacy synchronous preprocessing function for backward compatibility."""
    if isinstance(file_or_bytes, bytes):
        return _preprocess_image_core(file_or_bytes, target_size)
    else:
        # Assume it's a file-like object
        image = Image.open(file_or_bytes)
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        image = image.resize(target_size, Image.Resampling.LANCZOS)
        img_array = np.array(image, dtype=np.float32) / 255.0
        return np.expand_dims(img_array, axis=0)