import subprocess
import random
import asyncio
import websockets
import json
import audiocraft
from audiocraft.models import MusicGen
from audiocraft.utils.notebook import display_audio
import torch
import numpy as np
import base64
import soundfile as sf
import os

model = MusicGen.get_pretrained('medium', device="cpu")

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": "audio",
            "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 tensor_to_wav(tensor: torch.Tensor, sample_rate: int, job_id: int) -> str:
    numpy_array = tensor.numpy()
    numpy_array = numpy_array.astype('float32')

    tmp_filename = '/tmp/'+str(job_id)+'.wav'
    sf.write(tmp_filename, numpy_array.T, sample_rate)
    with open(tmp_filename, 'rb') as wav_file:
        base64_string = base64.b64encode(wav_file.read()).decode()

    base64_string = "data:audio/wav;base64," + base64_string
    os.remove(tmp_filename)

    return base64_string


# Websocket handler function, this is executed everytime new user connects
async def handler(websocket):
    # Wait for the client to send a message
    message = await websocket.recv()

    # Parse the message as JSON
    message_data = json.loads(message)

    # Create messenger instance
    messenger = Messenger(websocket)

    # Check if the message contains a request to start the worker
    if message_data.get("started"):
        global next_job_id, model

        # Assign & increment the current job ID
        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 = message_data.get("prompt")
        settings = message_data.get("settings")
        top_k = int(settings['top_k'])
        top_p = float(settings['top_p'])
        temperature = float(settings['temperature'])
        duration = float(settings['duration'])
        cfg_coef = float(settings['cfg_coef'])

        await messenger.progress(33)
        model.set_generation_params(
            use_sampling=True,
            top_k=top_k,
            top_p=top_p,
            temperature=temperature,
            duration=duration,
            cfg_coef=cfg_coef,
        )

        await messenger.progress(66)
        output = model.generate([prompt], progress=True)
        samples = output.detach().cpu()
        output_encoded = tensor_to_wav(samples[0], 32000, current_job_id)

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

        # Send the result back to the client
        await messenger.success([output_encoded])


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)