import gradio as gr


from transformers import pipeline 

from datasets import load_dataset
import soundfile as sf
import torch
import os
import io 
import base64
import numpy as np
from pydub import AudioSegment

os.environ['TRANSFORMERS_CACHE'] = '.cache'

print ("----- setting up pipeline -----")


synthesiser = pipeline("text-to-speech", "microsoft/speecht5_tts")

print ("----- setting up dataset -----")


embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
speaker_embedding = torch.tensor(embeddings_dataset[7306]["xvector"]).unsqueeze(0)
# You can replace this embedding with your own as well.

print ("----- synthetizing audio -----")

#speech = synthesiser("Hello, my dog is cooler than you!", forward_params={"speaker_embeddings": speaker_embedding})

#sf.write("speech.wav", speech["audio"], samplerate=speech["sampling_rate"])

def greet(name):
    return "Hello " + name + "!!"


def synthesise_audio(text, forward_params=None):
    if len(text) > 100:
        raise ValueError("Error: El texto es demasiado largo. Por favor, limita tu entrada a 100 caracteres.")

    speech = synthesiser(text, forward_params={"speaker_embeddings": speaker_embedding})
    # sf.write("speech.wav", speech["audio"], samplerate=speech["sampling_rate"])
    # return "speech.wav"
    
    # Convert numpy array to audio
    #with io.BytesIO() as f:
    #    sf.write(f, speech["audio"], samplerate=speech["sampling_rate"], format='wav')
    #    audio = f.getvalue()

    # Convert numpy array to audio
    audio = np.int16(speech["audio"] * 32767)
    audio_segment = AudioSegment(audio, sample_width=2, frame_rate=speech["sampling_rate"], channels=1)



    #return speech["audio"]

    return audio
    
  
    


#demo = gr.Interface(fn=greet, inputs="text", outputs="text", description="----- TTS Testing -----")

input_text = gr.Textbox(lines=10, label="Enter text here")


demo = gr.Interface(fn=synthesise_audio,
                    inputs=input_text,
                    #outputs="audio",
                    outputs = gr.Audio(type="numpy"),
                    description="----- manuai Text To Speech generator -----",
                    allow_flagging = False)

demo.launch(debug = True)