import cv2
import numpy as np
import base64
from PIL import Image
from payload_model import PayloadModel

def convert_base64_to_cv2(base64_string: str):
    return cv2.imdecode(np.frombuffer(base64.b64decode(base64_string), np.uint8), cv2.IMREAD_COLOR)

def convert_cv2_to_pil(image: np.ndarray):
    return Image.fromarray(image).convert('RGB')

def convert_base64_to_pil(base64_string: str):
    return convert_cv2_to_pil(convert_base64_to_cv2(base64_string))

def get_images_using_bbox(payload: PayloadModel):
    images = []
    # Forcing that only a single image is received
    cv2_image = convert_base64_to_cv2(payload.input_data[0])
    print(f"Bbox: {payload.bbox}")
    images_bboxes = payload.bbox
    image_bboxes = images_bboxes[0]
    for _, bbox in enumerate(image_bboxes):
        x1, y1, x2, y2 = bbox
        image = cv2_image[y1:y2, x1:x2]
        pil_image = convert_cv2_to_pil(image)
        images.append(pil_image)
    return images

def get_whole_image(payload: PayloadModel):
    images = []
    # Forcing that only a single image is received
    pil_image = convert_base64_to_pil(payload.input_data[0])
    images.append(pil_image)
    return images