import numpy as np
import cv2
import io
import base64
from PIL import Image
from app.config import config

def preprocess_image(image_bytes: bytes) -> np.ndarray:
    """Preprocess uploaded image for model inference"""
    # Convert bytes to PIL Image
    image = Image.open(io.BytesIO(image_bytes))

    # Convert to RGB if needed
    if image.mode != 'RGB':
        image = image.convert('RGB')

    # Convert to numpy array
    image_np = np.array(image)

    # Resize to model input size from config
    image_resized = cv2.resize(
        image_np, 
        (config.IMG_WIDTH, config.IMG_HEIGHT), 
        interpolation=cv2.INTER_LINEAR
    )

    return image_resized

def create_overlay_image(original: np.ndarray, mask: np.ndarray) -> np.ndarray:
    """Create overlay visualization with red highlights"""
    overlay = original.copy()
    # Apply red color to predicted stroke regions
    overlay[mask > 127] = [255, 64, 64]  # Red color for stroke regions
    return overlay

def numpy_to_base64(image_np: np.ndarray) -> str:
    """Convert numpy array to base64 string"""
    if len(image_np.shape) == 2:  # Grayscale
        image_pil = Image.fromarray(image_np, mode='L')
    else:  # RGB
        image_pil = Image.fromarray(image_np.astype(np.uint8))

    buffer = io.BytesIO()
    image_pil.save(buffer, format='PNG')
    img_str = base64.b64encode(buffer.getvalue()).decode()
    return f"data:image/png;base64,{img_str}"

def images_to_base64_urls(original_image: np.ndarray, prob_map: np.ndarray, overlay_image: np.ndarray) -> tuple:
    """Convert all images to base64 URLs for response when external storage not available"""
    original_url = numpy_to_base64(original_image)
    prediction_url = numpy_to_base64((prob_map * 255).astype(np.uint8))
    overlay_url = numpy_to_base64(overlay_image)
    
    return original_url, prediction_url, overlay_url