from PIL import Image
from io import BytesIO
import numpy as np
import base64
import torch
import torchvision.transforms.functional as F
from S2I import Sketch2ImagePipeline



class Sketch2ImageController():
    def __init__(self, gr):
        super().__init__()
        self.gr = gr
        self.style_list = [
            {"name": "Comic",
             "prompt": "comic {prompt} . graphic illustration, comic art, graphic novel art, vibrant, highly detailed"},
            {"name": "Cinematic",
             "prompt": "cinematic still {prompt} . emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy"},
            {"name": "3D Model",
             "prompt": "professional 3d model {prompt} . octane render, highly detailed, volumetric, dramatic lighting"},
            {"name": "Anime",
             "prompt": "anime artwork {prompt} . anime style, key visual, vibrant, studio anime,  highly detailed"},
            {"name": "Digital Art",
             "prompt": "concept art {prompt} . digital artwork, illustrative, painterly, matte painting, highly detailed"},
            {"name": "Photographic",
             "prompt": "cinematic photo {prompt} . 35mm photograph, film, bokeh, professional, 4k, highly detailed"},
            {"name": "Pixel art", "prompt": "pixel-art {prompt} . low-res, blocky, pixel art style, 8-bit graphics"},
            {"name": "Fantasy art",
             "prompt": "ethereal fantasy concept art of  {prompt} . magnificent, celestial, ethereal, painterly, epic, majestic, magical, fantasy art, cover art, dreamy"},
            {"name": "Neonpunk",
             "prompt": "neonpunk style {prompt} . cyberpunk, vaporwave, neon, vibes, vibrant, stunningly beautiful, crisp, detailed, sleek, ultramodern, magenta highlights, dark purple shadows, high contrast, cinematic, ultra detailed, intricate, professional"},
            {"name": "Manga",
             "prompt": "manga style {prompt} . vibrant, high-energy, detailed, iconic, Japanese comic style"},
        ]

        self.styles = {k["name"]: k["prompt"] for k in self.style_list}
        self.STYLE_NAMES = list(self.styles.keys())
        self.DEFAULT_STYLE_NAME = "Fantasy art"
        self.MAX_SEED = np.iinfo(np.int32).max

        # Initialize the model once here
        self.pipe = None
        self.zero_options = None
    def load_pipeline(self, zero_options):
        if self.pipe is None or zero_options != self.zero_options:
            self.pipe = Sketch2ImagePipeline()
            self.zero_options = zero_options

    def update_canvas(self, use_line, use_eraser):
        brush_size = 20 if use_eraser else 4
        _color = "#ffffff" if use_eraser else "#000000"
        return self.gr.update(brush_radius=brush_size, brush_color=_color, interactive=True)

    def upload_sketch(self, file):
        _img = Image.open(file.name).convert("L")
        return self.gr.update(value=_img, source="upload", interactive=True)

    @staticmethod
    def pil_image_to_data_uri(img, format="PNG"):
        buffered = BytesIO()
        img.save(buffered, format=format)
        img_str = base64.b64encode(buffered.getvalue()).decode()
        return f"data:image/{format.lower()};base64,{img_str}"

    def artwork(self, options, image, prompt, prompt_template, style_name, seed, val_r, faster, model_name, type_flag):
        self.load_pipeline(zero_options=options)

        prompt = prompt_template.replace("{prompt}", prompt)

        if type_flag == 'live-sketch':
            img = Image.fromarray(np.array(image["composite"])[:, :, -1])
        elif type_flag == 'upload':
            img = image["composite"]

        img = img.convert("RGB")
        img = img.resize((512, 512))

        image_t = F.to_tensor(img) > 0.5
        c_t = image_t.unsqueeze(0).cuda().float()

        torch.manual_seed(seed)
        _, _, H, W = c_t.shape
        noise = torch.randn((1, 4, H // 8, W // 8), device=c_t.device)

        with torch.no_grad():
            output_image = self.pipe.generate(c_t, prompt, r=val_r, noise_map=noise, half_model=faster, model_name=model_name)

        output_pil = F.to_pil_image(output_image[0].cpu() * 0.5 + 0.5)

        if type_flag == 'live-sketch':
            input_uri = self.pil_image_to_data_uri(Image.fromarray(255 - np.array(img)))
        else:
            input_uri = self.pil_image_to_data_uri(img)

        return output_pil, self.gr.update(link=input_uri)