from io import BytesIO
from typing import List, Union

import torch
from diffusers import (
    ControlNetModel,
    StableDiffusionControlNetInpaintPipeline,
    StableDiffusionInpaintPipeline,
    UniPCMultistepScheduler,
)
from PIL import Image, ImageFilter, ImageOps

import internals.util.image as ImageUtil
from internals.data.result import Result
from internals.pipelines.controlnets import ControlNet
from internals.pipelines.remove_background import RemoveBackgroundV2
from internals.pipelines.upscaler import Upscaler
from internals.util.commons import download_image


class ReplaceBackground:
    def load(self, upscaler: Upscaler, remove_background: RemoveBackgroundV2):
        controlnet = ControlNetModel.from_pretrained(
            "lllyasviel/control_v11p_sd15_lineart", torch_dtype=torch.float16
        ).to("cuda")
        pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
            "runwayml/stable-diffusion-inpainting",
            controlnet=controlnet,
            torch_dtype=torch.float16,
        )
        pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
        pipe.to("cuda")

        upscaler.load()

        self.pipe = pipe
        self.upscaler = upscaler
        self.remove_background = remove_background

    def replace(
        self,
        image: Union[str, Image.Image],
        width: int,
        height: int,
        product_scale_width: float,
        prompt: Union[str, List[str]],
        negative_prompt: Union[str, List[str]],
        resize_dimension: int,
        seed: int,
        steps: int,
    ):
        if type(image) is str:
            image = download_image(image)

        torch.manual_seed(seed)
        torch.cuda.manual_seed(seed)

        image = image.convert("RGB")
        image = self.remove_background.remove(image)

        width = int(width)
        height = int(height)

        n_width = int(width * product_scale_width)
        n_height = int(n_width * height // width)

        print(width, height, n_width, n_height)

        image = ImageUtil.padd_image(image, n_width, n_height)

        f_image = Image.new("RGBA", (width, height), (0, 0, 0, 0))
        f_image.paste(image, ((width - n_width) // 2, (height - n_height) // 2))
        image = f_image

        mask = image.copy()
        pixdata = mask.load()

        w, h = mask.size
        for y in range(h):
            for x in range(w):
                item = pixdata[x, y]
                if item[3] == 0:
                    pixdata[x, y] = (255, 255, 255, 255)
                else:
                    pixdata[x, y] = (0, 0, 0, 255)

        mask = mask.convert("RGB")

        condition_image = ControlNet.linearart_condition_image(image)

        result = self.pipe.__call__(
            prompt=prompt,
            negative_prompt=negative_prompt,
            image=image,
            mask_image=mask,
            control_image=condition_image,
            guidance_scale=9,
            strength=1,
            height=height,
            width=width,
        )
        result = Result.from_result(result)

        images, has_nsfw = result

        if not has_nsfw:
            for i in range(len(images)):
                images[i].paste(image, (0, 0), image)
                out_bytes = self.upscaler.upscale(images[i], resize_dimension)
                images[i] = Image.open(BytesIO(out_bytes)).convert("RGB")

        return (images, has_nsfw)