from diffusers import StableDiffusionPipeline
import PIL.Image
import asyncio
import websockets
import json
import base64
import torch
import io

model = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
next_job_id = 1


class Messenger:
    def __init__(self, websocket):
        self.websocket = websocket

    async def progress(self, progress):
        payload = {
            "progress": progress
        }
        payload_serialized = json.dumps(payload)
        await self.websocket.send(payload_serialized)

    async def success(self, result):
        payload = {
            "type": "image",
            "success": result
        }
        payload_serialized = json.dumps(payload)
        await self.websocket.send(payload_serialized)
        await self.websocket.close()

    async def error(self, error):
        payload = {
            "error": error
        }
        payload_serialized = json.dumps(payload)
        await self.websocket.send(payload_serialized)
        await self.websocket.close()


def to_base64(image: PIL.Image.Image) -> str:
    with io.BytesIO() as output:
        image.save(output, format='PNG')
        contents = output.getvalue()
    img_str = base64.b64encode(contents).decode()
    return f"data:image/png;base64,{img_str}"


# Websocket handler function, this is executed everytime new user connects
async def handler(websocket):
    message = await websocket.recv()
    message_data = json.loads(message)
    messenger = Messenger(websocket)

    if message_data.get("started"):
        global next_job_id, model

        current_job_id: int = next_job_id
        next_job_id += 1
        print(f"Started worker with job id: {current_job_id}")
        await messenger.progress(0)

        # Get the prompt, and model from the message
        prompt: str = message_data.get("prompt")
        settings = message_data.get("settings")

        await messenger.progress(25)
        pipe = model.to(settings["device"])

        await messenger.progress(50)

        generated_images_arr: list = pipe(prompt).images

        encoded_images_arr: list = [to_base64(img) for _, img in enumerate(generated_images_arr)]

        print(f"Sequence complete for job id: {current_job_id}")
        await messenger.progress(100)

        # Send the result back to the client
        await messenger.success(encoded_images_arr)


if __name__ == "__main__":
    start_server = websockets.serve(handler, "0.0.0.0", 1337)
    asyncio.get_event_loop().run_until_complete(start_server)
    print("Server started...")

    try:
        asyncio.get_event_loop().run_forever()
    except KeyboardInterrupt:
        print("Keyboard Interrupt: Exiting server")
    except:
        print("Unknown error")
    finally:
        exit(1)