import os
import numpy as np
import torch
import torch.nn.functional as F
from torchvision.transforms.functional import normalize
from PIL import Image, ImageOps, ImageSequence
from typing import List
from pathlib import Path
from huggingface_hub import snapshot_download, hf_hub_download


def tensor_to_pil(images: torch.Tensor | List[torch.Tensor]) -> List[Image.Image]:
    if not isinstance(images, list):
        images = [images]
    imgs = []
    for image in images:
        i = 255.0 * image.cpu().numpy()
        img = Image.fromarray(np.clip(np.squeeze(i), 0, 255).astype(np.uint8))
        imgs.append(img)
    return imgs


def pad_image(input_image, background_color=(0, 0, 0)):
    w, h = input_image.size
    pad_w = (64 - w % 64) % 64
    pad_h = (64 - h % 64) % 64

    new_size = (w + pad_w, h + pad_h)
    im_padded = Image.new(input_image.mode, new_size, background_color)
    im_padded.paste(input_image, (pad_w // 2, pad_h // 2))

    if im_padded.size[0] == im_padded.size[1]:
        return im_padded
    elif im_padded.size[0] > im_padded.size[1]:
        new_size = (im_padded.size[0], im_padded.size[0])
        new_image = Image.new(im_padded.mode, new_size, background_color)
        new_image.paste(im_padded, (0, (new_size[1] - im_padded.size[1]) // 2))
        return new_image
    else:
        new_size = (im_padded.size[1], im_padded.size[1])
        new_image = Image.new(im_padded.mode, new_size, background_color)
        new_image.paste(im_padded, ((new_size[0] - im_padded.size[0]) // 2, 0))
        return new_image


def pil_to_tensor(image: Image.Image) -> tuple[torch.Tensor, torch.Tensor]:
    output_images = []
    output_masks = []
    for i in ImageSequence.Iterator(image):
        i = ImageOps.exif_transpose(i)
        if i.mode == "I":
            i = i.point(lambda i: i * (1 / 255))
        image = i.convert("RGB")
        image = np.array(image).astype(np.float32) / 255.0
        image = torch.from_numpy(image)[None,]
        if "A" in i.getbands():
            mask = np.array(i.getchannel("A")).astype(np.float32) / 255.0
            mask = 1.0 - torch.from_numpy(mask)
        else:
            mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
        output_images.append(image)
        output_masks.append(mask.unsqueeze(0))

    if len(output_images) > 1:
        output_image = torch.cat(output_images, dim=0)
        output_mask = torch.cat(output_masks, dim=0)
    else:
        output_image = output_images[0]
        output_mask = output_masks[0]

    return (output_image, output_mask)


def preprocess_image(im: np.ndarray, model_input_size: list) -> torch.Tensor:
    if len(im.shape) < 3:
        im = im[:, :, np.newaxis]
    # orig_im_size=im.shape[0:2]
    im_tensor = torch.tensor(im, dtype=torch.float32).permute(2, 0, 1)
    im_tensor = F.interpolate(
        torch.unsqueeze(im_tensor, 0), size=model_input_size, mode="bilinear"
    ).type(torch.uint8)
    image = torch.divide(im_tensor, 255.0)
    image = normalize(image, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
    return image


def postprocess_image(result: torch.Tensor, im_size: list) -> np.ndarray:
    result = torch.squeeze(F.interpolate(result, size=im_size, mode="bilinear"), 0)
    ma = torch.max(result)
    mi = torch.min(result)
    result = (result - mi) / (ma - mi)
    im_array = (result * 255).permute(1, 2, 0).cpu().data.numpy().astype(np.uint8)
    im_array = np.squeeze(im_array)
    return im_array


def downloadModels():
    MODEL_PATH = snapshot_download(
        repo_id="RunDiffusion/Juggernaut-XL-v6", allow_patterns="*.safetensors"
    )
    LAYERS_PATH = snapshot_download(
        repo_id="LayerDiffusion/layerdiffusion-v1", allow_patterns="*.safetensors"
    )
    for file in Path(LAYERS_PATH).glob("*.safetensors"):
        target_path = Path(f"./ComfyUI/models/layer_model/{file.name}")
        if not target_path.exists():
            os.symlink(file, target_path)
    for model in Path(MODEL_PATH).glob("*.safetensors"):
        model_target_path = Path(f"./ComfyUI/models/checkpoints/{model.name}")
        if not model_target_path.exists():
            os.symlink(model, model_target_path)


examples = [
    [
        "A very cute monster cat on a glass bottle",
        "ugly distorted image, low quality, text, bad, not good ,watermark",
        None,
        False,
        None,
        1231231,
        5,
    ],
    [
        "A picture from above captures a beautiful, small toucan bird flying in the sky.",
        "ugly distorted image, low quality, text, bad, not good ,watermark",
        "./examples/bg.png",
        False,
        "SDXL, Background",
        1234144,
        8,
    ],
    [
        "a photo a men surrounded by a crowd of people in a circle",
        "ugly distorted image, low quality, text, bad, not good ,watermark",
        "./examples/lecun.png",
        True,
        "SDXL, Foreground",
        123123,
        10,
    ],
    [
        "An image of a galaxy",
        "ugly distorted image, low quality, text, bad, not good ,watermark",
        "./examples/julien.png",
        True,
        "SDXL, Foreground",
        123123,
        10,
    ],
    [
        "a men jumping on swiming pool full of people",
        "ugly distorted image, low quality, text, bad, not good ,watermark",
        "./examples/old_jump.png",
        False,
        "SDXL, Foreground",
        5350795678007195000,
        10,
    ],
    [
        "a cute cat flying over Manhattan time square",
        "ugly distorted image, low quality, text, bad, not good ,watermark",
        "./examples/cat.png",
        True,
        "SDXL, Foreground",
        123123,
        10,
    ],
]