import torch
import SpoutGL

from itertools import islice, cycle, repeat
import array
from random import randint
import time
from OpenGL import GL

from multiprocessing import Queue
import numpy as np
from PIL import Image


TARGET_FPS = 30
SEND_WIDTH = 512
SEND_HEIGHT = 512

alpha_cache = np.full((512, 512, 1), 255, dtype=np.uint8)



def spout_buffer_to_tensor(buffer, width, height):
    # np_buffer = np.asarray(buffer, dtype=np.uint8)
    np_buffer=np.frombuffer(buffer, dtype=np.uint8)
    image_bgra = np_buffer.reshape((height, width, 4))

    image_rgb = image_bgra[..., [2, 1, 0]]
    image_float = image_rgb.astype(np.float32) / 255.0
    # image_normalized = (image_float * 2.0) - 1.0    
    tensor = torch.from_numpy(image_float).permute(2, 0, 1)

    del np_buffer  # Free memory
    del image_bgra  # Free memory
    del image_rgb  # Free memory
    del image_float  # Free memory


    return tensor.unsqueeze(0)


def get_spout_image(queue, wwidth: int, wheight: int) -> None:
    with SpoutGL.SpoutReceiver() as receiver:
        receiver.setReceiverName("Spout DX11 Sender")
        image_bgra = np.zeros((SEND_HEIGHT, SEND_WIDTH, 4), dtype=np.uint8)

        while True:
            result = receiver.receiveImage(image_bgra, GL.GL_RGBA, False, 0)
            # print("Receive result", result)

            if receiver.isUpdated():
                continue
                # width = receiver.getSenderWidth()
                # height = receiver.getSenderHeight()
                # image_bgra = array.array('B', [0] * (width * height * 4))  # Correctly reallocate buffer with updated size
                # print("Spout Receiver updated, Buffer size", width, height)

            # if buffer and result and not SpoutGL.helpers.isBufferEmpty(buffer):
            if SpoutGL.helpers.isBufferEmpty(image_bgra):
                continue
                # pixels=spout_buffer_to_tensor(buffer, width, height)
                # print("get_spout_image", pixels.shape)
            
            image_rgb_array= image_bgra[:, :, [2, 1, 0]]
            pixels=Image.fromarray(image_rgb_array, 'RGB')
            queue.put(pixels, block=False)
                        

            # Wait until the next frame is ready
            # Wait time is in milliseconds; note that 0 will return immediately
            # receiver.waitFrameSync("SpoutSender", 10000)




def randcolor():
    return randint(0, 255)


def tensor_to_spout_image(tensor):
    image = tensor.squeeze(0)
    if image.device.type != "cpu":
        image = image.cpu()
    image = image.permute(1, 2, 0).numpy()

    if image.min() < 0:
        image = (image + 1) / 2  # Scale from [-1, 1] to [0, 1]
    image = np.clip(image * 255, 0, 255).astype(np.uint8)

    # h, w, _ = image_np.shape
    # alpha = np.full((h, w, 1), 255, dtype=np.uint8)
    image_rgba = np.concatenate((image, alpha_cache), axis=-1)
    image_bgra = image_rgba[..., [2, 1, 0, 3]]

    del image  # Free memory

    return np.ascontiguousarray(image_bgra)  # Ensure the array is contiguous in memory

def send_spout_image(queue: Queue, width: int, height: int)->None:
    
    with SpoutGL.SpoutSender() as sender:
        sender.setSenderName("StreamDiffusion")
        
        while True:

            # Check if there are images in the queue
            if not queue.empty():
                output_image = queue.get(block=False)
                # pixels = tensor_to_spout_image(image)
               
                output_bgr_array = np.array(output_image, dtype=np.uint8)[:, :, ::-1]
                output_bgra_array = np.zeros((SEND_HEIGHT, SEND_WIDTH, 4), dtype=np.uint8)
                output_bgra_array[:, :, :3] = output_bgr_array
                output_bgra_array[:, :, 3] = 255
                buffer = output_bgra_array
               
                
                result = sender.sendImage(buffer, width, height, GL.GL_RGBA, False, 0)
                # print("Send result", result)
                
                # Indicate that a frame is ready to read
                sender.setFrameSync("StreamDiffusion")

                # Wait for next send attempt
                # time.sleep(1./TARGET_FPS)