import bpy
import math
import numpy as np

from pathlib import Path


def alpha_to_color(pixels_data, color):
    """Convert Alpha to WhiteBG"""
    new_pixels_data = []
    for i in pixels_data:
        height, width, array_d = i.shape
        mask = i[:, :, 3:]
        background = np.array([color[0], color[1], color[2], 1], dtype=np.float32)
        background = np.tile(background, (height * width))
        background = np.reshape(background, (height, width, 4))
        new_pixels_data.append(i * mask + background * (1 - mask))
        # print(new_pixels_data)#Dbg
    return new_pixels_data


def create_array(height, width):
    return np.zeros((height * width * 4), dtype=np.float32)


def read_pixels_data(img, source_height, source_width):
    img_w, img_h = img.size

    if img_w != source_width:
        scale = abs(img_w / source_width)
        img.scale(int(img_w / scale), int(img_h / scale))
        img_w, img_h = img.size

    array = create_array(img_h, img_w)
    img.pixels.foreach_get(array)
    array = array.reshape(img_h, img_w, 4)

    if array.shape[0] != source_height:
        # print('ARRAY SHAPE', array.shape[:], source_height)
        missing_height = int(abs(source_height - img_h) / 2)
        empty_array = create_array(missing_height, source_width)
        empty_array = empty_array.reshape(missing_height, source_width, 4)
        array = np.vstack((empty_array, array, empty_array))

    return array.reshape(source_height, source_width, 4)


def create_final(output_name, pixels_data, final_height, final_width):
    # print('output_name: ', output_name)

    new_img = bpy.data.images.get(output_name)
    if new_img:
        bpy.data.images.remove(new_img)

    new_img = bpy.data.images.new(output_name, final_width, final_height)
    new_img.generated_color = (0, 0, 0, 0)

    # print('pixels_data: ', pixels_data)
    new_img.pixels.foreach_set(pixels_data)

    return new_img


def guess_input_format(img_list):
    for i in img_list:
        if i.size[0] == i.size[1]:
            return i.size


def format_files(files, catalog_data):
    img_dict = {}
    for k, v in catalog_data.items():
        if "/" not in k:
            continue
        img_dict[v["name"]] = [f for f in files if v["name"] in f]

    return img_dict


def mosaic_export(
    files,
    catalog_data,
    row=2,
    columns=2,
    auto_calculate=True,
    bg_color=(
        0.18,
        0.18,
        0.18,
    ),
    resize_output=100,
):

    img_dict = format_files(files, catalog_data)

    for cat, files_list in img_dict.items():

        if not files_list:
            continue

        for i in bpy.data.images:
            bpy.data.images.remove(i)

        img_list = []

        chars = Path(files_list[0]).parts[-4]
        output_dir = str(Path(files_list[0]).parent.parent)

        ext = "jpg"
        output_name = f"{chars}_{cat}.{ext}"

        for img in files_list:
            img_list.append(bpy.data.images.load(img, check_existing=True))

        for i in img_list:
            i.colorspace_settings.name = "Raw"

        if auto_calculate:
            rows = int(math.sqrt(len(img_list)))
            columns = math.ceil(len(img_list) / rows)

        if rows * columns < len(img_list):
            raise AttributeError("Grid too small for number of images")

        src_w, src_h = img_list[0].size
        final_w = src_w * columns
        final_h = src_h * rows

        img_pixels = [read_pixels_data(img, src_h, src_w) for img in img_list]

        # Check if there is enough "data" to create an horizontal stack
        ##It not, create empty array
        h_stack = []
        total_len = rows * columns
        if len(img_pixels) < total_len:
            for i in range(total_len - len(img_pixels)):
                img_pixels.append(create_array(src_h, src_w).reshape(src_h, src_w, 4))

        img_pixels = alpha_to_color(img_pixels, bg_color)
        for i in range(0, len(img_pixels), columns):
            h_stack.append(np.hstack(img_pixels[i : i + columns]))
        if rows > 1:
            combined_stack = np.vstack(h_stack[::-1])
        else:
            combined_stack = np.hstack((h_stack[:]))

        combined_img = create_final(
            output_name, combined_stack.flatten(), final_h, final_w
        )

        if resize_output != 100:
            w, h = combined_img.size
            combined_img.scale(w * (resize_output * 0.01), h * (resize_output * 0.01))

        combined_img.filepath_raw = "/".join([output_dir, output_name])
        combined_img.file_format = "JPEG"
        combined_img.save()

        print(
            f"""
        Image saved: {combined_img.filepath_raw}
        """
        )