import io
import logging

import cv2
import numpy as np
from PIL import Image
from cv2 import cvtColor, COLOR_BGR2RGBA

from app.core.config import AIDA_CLOTHING
from app.service.utils.generate_uuid import generate_uuid
from app.service.utils.oss_client import oss_upload_image
from ..builder import PIPELINES
from ...utils.conversion_image import rgb_to_rgba
from ...utils.upload_image import upload_png_mask


@PIPELINES.register_module()
class Split(object):
    """
    Split image into front and back layer according to the segmentation result
    """

    # @ClassCallRunTime
    # KNet
    def __call__(self, result):
        try:

            if result['name'] in ('outwear', 'dress', 'blouse', 'skirt', 'trousers', 'tops', 'bottoms'):
                front_mask = result['front_mask']
                back_mask = result['back_mask']
                rgba_image = rgb_to_rgba(result['final_image'], front_mask + back_mask)
                new_size = (int(rgba_image.shape[1] * result["scale"] * result["resize_scale"][0]), int(rgba_image.shape[0] * result["scale"] * result["resize_scale"][1]))
                rgba_image = cv2.resize(rgba_image, new_size)
                result_front_image = np.zeros_like(rgba_image)
                front_mask = cv2.resize(front_mask, new_size)
                result_front_image[front_mask != 0] = rgba_image[front_mask != 0]
                result_front_image_pil = Image.fromarray(cvtColor(result_front_image, COLOR_BGR2RGBA))
                result['front_image'], result["front_image_url"], _ = upload_png_mask(result_front_image_pil, f'{generate_uuid()}', mask=None)

                height, width = front_mask.shape
                mask_image = np.zeros((height, width, 3))
                mask_image[front_mask != 0] = [0, 0, 255]

                if result["name"] in ('blouse', 'dress', 'outwear', 'tops'):
                    result_back_image = np.zeros_like(rgba_image)
                    back_mask = cv2.resize(back_mask, new_size)
                    result_back_image[back_mask != 0] = rgba_image[back_mask != 0]
                    result_back_image_pil = Image.fromarray(cvtColor(result_back_image, COLOR_BGR2RGBA))
                    result['back_image'], result["back_image_url"], _ = upload_png_mask(result_back_image_pil, f'{generate_uuid()}', mask=None)
                    mask_image[back_mask != 0] = [0, 255, 0]

                    rbga_mask = rgb_to_rgba(mask_image, front_mask + back_mask)
                    mask_pil = Image.fromarray(cvtColor(rbga_mask.astype(np.uint8), COLOR_BGR2RGBA))
                    image_data = io.BytesIO()
                    mask_pil.save(image_data, format='PNG')
                    image_data.seek(0)
                    image_bytes = image_data.read()
                    req = oss_upload_image(bucket=AIDA_CLOTHING, object_name=f"mask/mask_{generate_uuid()}.png", image_bytes=image_bytes)
                    result['mask_url'] = req.bucket_name + "/" + req.object_name
                else:
                    rbga_mask = rgb_to_rgba(mask_image, front_mask)
                    mask_pil = Image.fromarray(cvtColor(rbga_mask.astype(np.uint8), COLOR_BGR2RGBA))
                    image_data = io.BytesIO()
                    mask_pil.save(image_data, format='PNG')
                    image_data.seek(0)
                    image_bytes = image_data.read()
                    req = oss_upload_image(bucket=AIDA_CLOTHING, object_name=f"mask/mask_{generate_uuid()}.png", image_bytes=image_bytes)
                    result['mask_url'] = req.bucket_name + "/" + req.object_name
                    result['back_image'] = None
                    result["back_image_url"] = None
                    # result["back_mask_url"] = None
                    # result['back_mask_image'] = None
            # 创建中间图层
            result_pattern_image_rgba = rgb_to_rgba(result['pattern_image'], result['mask'])
            result_pattern_image_pil = Image.fromarray(cvtColor(result_pattern_image_rgba, COLOR_BGR2RGBA))
            result['pattern_image'], result['pattern_image_url'], _ = upload_png_mask(result_pattern_image_pil, f'{generate_uuid()}')
            return result
        except Exception as e:
            logging.warning(f"split runtime exception : {e} image_id : {result['image_id']}")