import os
import shutil
import subprocess
from io import BytesIO
from tempfile import NamedTemporaryFile

from PIL import Image, ImageOps

from config import IMAGE_STRIP_HEIGHT, IMAGE_QUALITY, IMAGE_CACHE, IMAGE_MAX_WIDTH, THUMBNAIL_WIDTH, THUMBNAIL_HEIGHT, \
    THUMBNAIL_QUALITY

Image.MAX_IMAGE_PIXELS = 1000000000
from glob import glob

def _copy(src, dst):
    try:
        os.link(src, dst)
    except:
        shutil.copy(src, dst)


def rechunk(im):

    height = im.height
    width = im.width
    last_chunk = (height // IMAGE_STRIP_HEIGHT) - 1

    for i, cursor in enumerate(range(0, height, IMAGE_STRIP_HEIGHT)):

        if i == last_chunk:
            chunk_height = IMAGE_STRIP_HEIGHT + (height % IMAGE_STRIP_HEIGHT)
        else:
            chunk_height = min(IMAGE_STRIP_HEIGHT, height)

        yield im.crop((0, cursor, width, cursor + chunk_height))

        if i == last_chunk:
            break

def webp_convert(input_file, output_file, quality=IMAGE_QUALITY):

    subprocess.run([
        "ffmpeg", "-y", "-loglevel", "error", "-hide_banner", "-nostats", "-i", input_file,
        "-c:v", "libwebp", "-preset", "drawing", "-quality", str(quality), "-compression_level", "6",
        output_file
    ], text=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)

def copy_from_cache(cache_folder, output_folder):

    for filename in glob(os.path.join(cache_folder, "*.webp")):
        _copy(filename, os.path.join(output_folder, os.path.basename(filename)))

def is_cached(cache_folder):
    if os.path.exists(os.path.join(cache_folder, "ok")):
        return True

    return False

def convert_cover(cache_key, filename, output_folder):
    cache_folder = os.path.join(IMAGE_CACHE, cache_key)

    os.makedirs(output_folder, exist_ok=True)
    output_filename = os.path.join(output_folder, "cover.webp")
    if os.path.exists(output_filename):
        os.remove(output_filename)

    if is_cached(cache_folder):
        copy_from_cache(cache_folder, output_folder)
        return

    shutil.rmtree(cache_folder, ignore_errors=True)
    os.makedirs(cache_folder)

    im = open_image(filename)

    if im.width > IMAGE_MAX_WIDTH:
        aspect_ratio = im.height / im.width
        new_height = int(IMAGE_MAX_WIDTH * aspect_ratio)

        im = im.resize((IMAGE_MAX_WIDTH, new_height), resample=Image.Resampling.LANCZOS)

    if im.mode != "RGB":
        im = im.convert("RGB")

    try:
        with NamedTemporaryFile(delete=False, suffix=".png") as f:
            im.save(f)

        webp_convert(f.name, os.path.join(cache_folder, f"cover.webp"), quality=90)

    finally:
        os.remove(f.name)

    with open(os.path.join(cache_folder, "ok"), "w"):
        pass

    copy_from_cache(cache_folder, output_folder)

def convert_rgb(im):
    if im.mode == "RGBA":
        background = Image.new('RGBA', im.size, color=(0, 0, 0))
        alpha_composite = Image.alpha_composite(background, im)
        im = alpha_composite.convert("RGB")

    if im.mode != "RGB":
        im = im.convert("RGB")

    return im

def convert_rgba(im):
    if im.mode != "RGBA":
        im = im.convert("RGBA")

    return im


def convert_chapter(cache_key, images, output_folder):

    # TODO: for non-long strips, stack horizontally (?)

    cache_folder = os.path.join(IMAGE_CACHE, cache_key)

    shutil.rmtree(output_folder, ignore_errors=True)
    os.makedirs(output_folder)

    if is_cached(cache_folder):
        copy_from_cache(cache_folder, output_folder)
        return

    shutil.rmtree(cache_folder, ignore_errors=True)
    os.makedirs(cache_folder)

    resized = []

    # Resize images
    for filename in images:
        im = open_image(filename)

        im = convert_rgb(im)

        if im.width > IMAGE_MAX_WIDTH:

            aspect_ratio = im.height / im.width
            new_height = int(IMAGE_MAX_WIDTH * aspect_ratio)

            if new_height == 0:
                continue

            im = im.resize((IMAGE_MAX_WIDTH, new_height), resample=Image.Resampling.LANCZOS)

        resized.append(im)

    # Prepare canvases
    canvases = [
        [resized[0]]
    ]
    current_canvas = canvases[0]
    last_image_width = resized[0].width
    for i in range(1, len(resized)):

        if last_image_width != resized[i].width:
            canvases.append([])
            current_canvas = canvases[-1]

        current_canvas.append(resized[i])
        last_image_width = resized[i].width

    # Stack
    stacked_images = []
    for canvas in canvases:
        width = canvas[0].width
        height = sum(im.height for im in canvas)

        new_image = Image.new("RGB", (width, height), color=(0, 0, 0, 0))

        cursor = 0
        for img in canvas:
            new_image.paste(img, (0, cursor))
            cursor += img.height

        stacked_images.append(new_image)

    # Rechunk
    chunks = []
    for im in stacked_images:
        for im_chunk in rechunk(im):
            chunks.append(im_chunk)

    # Convert
    for i, chunk in enumerate(chunks):

        try:
            with NamedTemporaryFile(delete=False, suffix=".png") as f:
                chunk.save(f)

            webp_convert(f.name, os.path.join(cache_folder, f"{i:03d}.webp"))

        finally:
            os.remove(f.name)

    with open(os.path.join(cache_folder, "ok"), "w"):
        pass

    copy_from_cache(cache_folder, output_folder)

def open_image(filename):

    with open(filename, "rb") as f:
        data = f.read()

    io = BytesIO(data)
    io.name = filename
    return Image.open(io)

def create_thumbnail_sheet(images, output_filename):
    images = [open_image(image["cover"]) for image in images]

    width = len(images) * THUMBNAIL_WIDTH
    height = THUMBNAIL_HEIGHT

    new_image = Image.new("RGBA", (width, height), color=(0, 0, 0, 0))

    for i, im in enumerate(images):
        im = convert_rgba(im)

        im = ImageOps.pad(im, (THUMBNAIL_WIDTH, THUMBNAIL_HEIGHT), color=(0, 0, 0, 0), method=Image.Resampling.LANCZOS)
        new_image.paste(im, (i * THUMBNAIL_WIDTH, 0))

    new_image.save(output_filename, quality=THUMBNAIL_QUALITY)