import albumentations as A
import io
import numpy as np
from PIL import Image, UnidentifiedImageError
from fastapi import HTTPException

val_test_transform = A.Compose([
    A.Resize(224, 224),
    A.Normalize(),  
])

async def pre_process_image(image_data: bytes):
    try:
        pil_image = Image.open(io.BytesIO(image_data)).convert("RGB")
        original_width, original_height = pil_image.size
        np_image = np.array(pil_image)
        augmented = val_test_transform(image=np_image)
        img_array = augmented["image"]
        img_array = np.expand_dims(img_array, axis=0).astype(np.float32)
        processed_image_dimensions = img_array.shape[1:]
        return img_array, original_width, original_height, processed_image_dimensions
    except UnidentifiedImageError as e:
        raise HTTPException(status_code=400, detail=f"Cannot identify image file. Please upload a valid image. Error: {e}")
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Image processing failed: {e}")