import io
import base64
import time

import torch.nn.functional as tnnf
import comfy.utils
from server import PromptServer
from PIL import Image

def scale_latent(latent, max_pixels=1000000):
    h = latent.shape[2]
    w = latent.shape[3]
    pixels = h * 8 * w * 8
    if pixels <= max_pixels:
        return latent
    scale = (max_pixels / pixels) ** 0.5
    new_h = int(h * scale)
    new_w = int(w * scale)
    new_h = max(1,new_h)
    new_w = max(1,new_w)
    return tnnf.interpolate(latent, size = (new_h, new_w), mode = "bilinear",
    align_corners = False)

def tensor_to_base64(image):
    # Receive a vae image and change it to base64, then send it to front-end.
    image = image.cpu().clamp(0, 1)
    image = (image[0].numpy() * 255).astype("uint8")
    img = Image.fromarray(image)
    buffer = io.BytesIO()
    img.save(buffer, format = "PNG")
    return base64.b64encode(buffer.getvalue()).decode()

def create_preview_callback(model, vae, steps, preview_decode = "Scale", interval = 1, sync_preview = True):
    # Sampler's callback, it has two mode.
    # Sync: The HD preview will sync with latent's step, one step = one preview image.
    # Fast: The HD preview will be rendered once at intervals.
    pbar = comfy.utils.ProgressBar(steps)
    def inner_callback(step, x0, x, total_steps):
        if(steps % interval == 0 or steps == total_steps):
            latent = (model.model.latent_format.process_out(x0))
            if preview_decode == "Scale":
                latent = scale_latent(latent)
            image = vae.decode(latent)
            data = tensor_to_base64(image)
            PromptServer.instance.send_sync(
            "high_preview", {
                "image": data,
                "step": step + 1,
                "total":total_steps
                }
            )
            if sync_preview:
                # Give browser some time to flash image.
                time.sleep(0.15)
        # Not use default preview window, so there is None.
        pbar.update_absolute(step + 1, total_steps, None)
    return inner_callback