import base64
import io
import json
import os
import random
import re
from typing import Any, Dict, List

import inflect
import numpy as np
import torch
import torchaudio
import uroman as ur
from audiotokenizer import AudioTokenizerForLocal
from outetts.wav_tokenizer.decoder import WavTokenizer
from transformers import AutoModelForCausalLM, AutoTokenizer


class EndpointHandler:
    def __init__(self, path=""):
        # Preload all the elements you are going to need at inference.
        # pseudo:
        # self.model= load_model(path)
        hf_path = "aybdee/igbo-SpeechSynthesis"
        wav_tokenizer_config_path = "/data/wavtokenizer_mediumdata_frame75_3s_nq1_code4096_dim512_kmeans200_attn.yaml"
        wav_tokenizer_model_path = "/data/wavtokenizer_large_speech_320_24k.ckpt"
        audio_tokenizer = AudioTokenizerForLocal(
            hf_path, wav_tokenizer_model_path, wav_tokenizer_config_path
        )
        self.model = AutoModelForCausalLM.from_pretrained(
            hf_path, torch_dtype="auto"
        ).to(audio_tokenizer.device)

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
         data args:
              inputs (:obj: `str` | `PIL.Image` | `np.array`)
              kwargs
        Return:
              A :obj:`list` | `dict`: will be serialized and returned
        """
        prompt = audio_tokenizer.create_prompt(data, "igbo", "igbo_male2")
        input_ids = audio_tokenizer.tokenize_prompt(prompt)
        output = self.model.generate(
            input_ids=input_ids,
            temperature=0.1,
            repetition_penalty=1.1,
            num_beams=4,
            max_length=8192,
        )
        codes = audio_tokenizer.get_codes(output)

        # converts the codes to audio
        audio = audio_tokenizer.get_audio(codes)
        buffer = io.BytesIO()
        torchaudio.save(buffer, audio, sample_rate=sample_rate, format="wav")
        buffer.seek(0)  # go to beginning

        # Base64 encode the buffer content
        audio_bytes = buffer.read()
        audio_base64 = base64.b64encode(audio_bytes).decode("utf-8")
        return {"audio": audio_base64}