import cv2
import numpy as np
from typing import Dict, List
from pycocotools import mask as coco_mask
from pytoshop import layers
import pytoshop
from pytoshop.enums import BlendMode


def generate_random_color():
    return np.random.randint(0, 256), np.random.randint(0, 256), np.random.randint(0, 256)


def create_base_layer(image):
    rgba_image = cv2.cvtColor(image, cv2.COLOR_RGB2RGBA)
    return [rgba_image]


def create_mask_layers(image, masks):
    layer_list = []

    for result in masks:
        rle = result['segmentation']
        mask = coco_mask.decode(rle).astype(np.uint8)
        rgba_image = cv2.cvtColor(image, cv2.COLOR_RGB2RGBA)
        rgba_image[..., 3] = cv2.bitwise_and(rgba_image[..., 3], rgba_image[..., 3], mask=mask)

        layer_list.append(rgba_image)

    return layer_list


def create_mask_gallery(image, masks):
    mask_array_list = []
    label_list = []

    for index, result in enumerate(masks):
        rle = result['segmentation']
        mask = coco_mask.decode(rle).astype(np.uint8)

        rgba_image = cv2.cvtColor(image, cv2.COLOR_RGB2RGBA)
        rgba_image[..., 3] = cv2.bitwise_and(rgba_image[..., 3], rgba_image[..., 3], mask=mask)

        mask_array_list.append(rgba_image)
        label_list.append(f'Part {index}')

    return [[img, label] for img, label in zip(mask_array_list, label_list)]


def create_mask_combined_images(image, masks):
    final_result = np.zeros_like(image)

    for result in masks:
        rle = result['segmentation']
        mask = coco_mask.decode(rle).astype(np.uint8)

        color = generate_random_color()
        colored_mask = np.zeros_like(image)
        colored_mask[mask == 1] = color

        final_result = cv2.addWeighted(final_result, 1, colored_mask, 0.5, 0)

    combined_image = cv2.addWeighted(image, 1, final_result, 0.5, 0)
    return [combined_image, "masked"]


def insert_psd_layer(psd, image_data, layer_name, blending_mode):
    channel_data = [layers.ChannelImageData(image=image_data[:, :, i], compression=1) for i in range(4)]

    layer_record = layers.LayerRecord(
        channels={-1: channel_data[3], 0: channel_data[0], 1: channel_data[1], 2: channel_data[2]},
        top=0, bottom=image_data.shape[0], left=0, right=image_data.shape[1],
        blend_mode=blending_mode,
        name=layer_name,
        opacity=255,
    )
    psd.layer_and_mask_info.layer_info.layer_records.append(layer_record)
    return psd


def save_psd(input_image_data, layer_data, layer_names, blending_modes, output_path):
    psd_file = pytoshop.core.PsdFile(num_channels=3, height=input_image_data.shape[0], width=input_image_data.shape[1])
    psd_file.layer_and_mask_info.layer_info.layer_records.clear()

    for index, layer in enumerate(layer_data):
        psd_file = insert_psd_layer(psd_file, layer, layer_names[index], blending_modes[index])

    with open(output_path, 'wb') as output_file:
        psd_file.write(output_file)


def save_psd_with_masks(
    image: np.ndarray,
    masks: Dict,
    output_path: str
):
    original_layer = create_base_layer(image)
    mask_layers = create_mask_layers(image, masks)
    names = [f'Part {i}' for i in range(len(mask_layers))]
    modes = [BlendMode.normal] * (len(mask_layers)+1)
    save_psd(image, original_layer+mask_layers, ['Original_Image']+names, modes, output_path)