import pathlib
from typing import List, Optional

import numpy as np
from PIL import Image
from tqdm.contrib import itertools as tqdm_itertools

from ..config import Paths
from .atlas import SpritesheetData, fetch_data, fetch_expression_sheets, fetch_config

def compose(input_id: int, filters: Optional[List[str]] = None):
    Paths.IMAGES.mkdir(exist_ok=True)
    Paths.OUTPUT.mkdir(exist_ok=True)

    if input_id < 10000:
        chara_ids = fetch_data(input_id)
        savefolder = Paths.OUTPUT / str(input_id)
    else:
        print(f"Processing manually uploaded charaId {input_id}")
        savefolder = Paths.OUTPUT / "manual"
        chara_ids = [input_id]

    if not savefolder.is_dir():
        savefolder.mkdir(parents=True, exist_ok=True)

    if filters is not None:
        chara_ids = [ v for v in chara_ids if v in filters ]

    for char_id in chara_ids:
        expfolder = fetch_expression_sheets(savefolder.stem, char_id)
        config = fetch_config(char_id)
        process_sprite(expfolder, config, savefolder)

    print(f"Files have been saved at: {savefolder.absolute()}")


def calculate_counts(width: int, height: int, facesize: tuple[int, int]):
    return height // facesize[1], width // facesize[0]

def gen_main_sprite(folder: pathlib.Path):
    image = Image.open(folder / "0.png")
    width, height = image.size
    return image.crop((0, 0, width, height - 256))

def process_sprite(images_folder: pathlib.Path, configdata: SpritesheetData, outputfolder: pathlib.Path):
    main_sprite = gen_main_sprite(images_folder)
    save_sprite(main_sprite, outputfolder, f"{images_folder.stem}")

    for i in images_folder.iterdir():
        initial_row, index = 0, int(i.stem)
        expressions = Image.open(i)

        rowcount, colcount = calculate_counts(*expressions.size, configdata["facesize"])

        if i.name == "0.png" and 256 < configdata["facesize"][1]:
            continue
        elif i.name == "0.png":
            initial_row = rowcount - 1

        for x, y in tqdm_itertools.product(range(initial_row, rowcount), range(0, colcount), ascii="-="):
            img = generate_sprite(main_sprite, expressions, x, y, configdata)
            if img is not None:
                save_sprite(img, outputfolder, f"{images_folder.stem}", (x, y, colcount, index))


def generate_sprite(main_sprite: Image.Image, expressions: Image.Image, row: int, col: int, configdata: SpritesheetData) -> Image.Image | None:
    position, facesize = configdata["position"], configdata["facesize"]
    roi = (
        col * facesize[0],
        row * facesize[1],
        (col + 1) * facesize[0] - 1,
        (row + 1) * facesize[1] - 1
    )

    expression = expressions.crop(roi)

    if is_empty(expression):
        return None

    composition = main_sprite.copy()
    composition.paste(expression, position, expression)
    return composition


def save_sprite(image: Image.Image, outputfolder: pathlib.Path, name: str, info: tuple | None = None):
    savefolder = outputfolder / name
    if not savefolder.is_dir():
        savefolder.mkdir()

    postfix = "0"

    if info is not None:
        (row, col, column_count, file_idx) = info

        if file_idx == 0 and column_count == 4:
            postfix = str(col + 1)
        elif file_idx == 0:
            raise ValueError("Should not have any faces")
        elif column_count == 4:
            postfix = str((column_count * row + col + 1) + pow(column_count, 2) * (file_idx - 1) + column_count)
        elif column_count == 1:
            postfix = str((file_idx - 1) * 16 + 1 if file_idx >= 2 else file_idx * 4 + 1)
        elif column_count < 4:
            postfix = str((column_count * row + col + 1) + pow(column_count, 2) * (file_idx - 1))
        else:
            raise ValueError("Unaccounted case")

    outfile = savefolder / f"{postfix}.png"

    with open(outfile, 'wb') as file:
        image.save(file)


def is_empty(img: Image.Image):
    data = np.asarray(img.crop((96, 96, 160, 160)).convert('LA'))
    np.reshape(data, (-1, 1))
    _, count_unique = np.unique(data, return_counts=True)
    if count_unique.size < 10:
        return True
    return False