import sys
import subprocess
import pkg_resources

# Função para verificar e instalar pacotes faltantes
def install_missing_packages():
    required = {'torch', 'torchaudio', 'transformers', 'gradio', 'numpy', 'accelerate', 'soundfile'}
    installed = {pkg.key for pkg in pkg_resources.working_set}
    missing = required - installed

    if missing:
        python = sys.executable
        subprocess.check_call([python, '-m', 'pip', 'install', *missing], stdout=subprocess.DEVNULL)
        
# Instalar pacotes faltantes
install_missing_packages()

import spaces
import gradio as gr
import torch
from string import punctuation
import re

from parler_tts import ParlerTTSForConditionalGeneration
from transformers import AutoTokenizer, AutoFeatureExtractor, set_seed
from helpers.text_normalizer import BrazilianPortugueseTextNormalizer

device = "cuda:0" if torch.cuda.is_available() else "cpu"

repo_id = "GustavoNeneve/parler-ttsbr"  # Seu repositório no Hugging Face

model = ParlerTTSForConditionalGeneration.from_pretrained(repo_id).to(device)
tokenizer = AutoTokenizer.from_pretrained(repo_id)
feature_extractor = AutoFeatureExtractor.from_pretrained(repo_id)

SAMPLE_RATE = feature_extractor.sampling_rate
SEED = 42

number_normalizer = BrazilianPortugueseTextNormalizer()

default_text = "Olá! Esta é a primeira versão do modelo em português brasileiro."
default_description = "Francisca fala de forma clara e profissional com sotaque brasileiro neutro."
examples = [
    [
        "Olá, como posso ajudar você hoje?",
        "Francisca fala de forma acolhedora e profissional com sotaque brasileiro neutro."
    ],
    [
        "Esta é uma mensagem importante para todos os usuários.",
        "Francisca fala com ênfase e clareza em ritmo moderado.",
    ],
    [
        "Seja bem-vindo ao nosso sistema automatizado.",
        "Francisca fala de forma amigável em um ritmo natural.",
    ]
]


def preprocess(text):
    text = number_normalizer.normalize(text)
    if text[-1] not in punctuation:
        text = f"{text}."
    return text


@spaces.GPU
def gen_tts(text, description):
    inputs = tokenizer(description, return_tensors="pt").to(device)
    prompt = tokenizer(preprocess(text), return_tensors="pt").to(device)

    set_seed(SEED)
    generation = model.generate(input_ids=inputs.input_ids, prompt_input_ids=prompt.input_ids)
    audio_arr = generation.cpu().numpy().squeeze()

    return SAMPLE_RATE, audio_arr


css = """
        #share-btn-container {
            display: flex;
            padding-left: 0.5rem !important;
            padding-right: 0.5rem !important;
            background-color: #000000;
            justify-content: center;
            align-items: center;
            border-radius: 9999px !important; 
            width: 13rem;
            margin-top: 10px;
            margin-left: auto;
            flex: unset !important;
        }
        #share-btn {
            all: initial;
            color: #ffffff;
            font-weight: 600;
            cursor: pointer;
            font-family: 'IBM Plex Sans', sans-serif;
            margin-left: 0.5rem !important;
            padding-top: 0.25rem !important;
            padding-bottom: 0.25rem !important;
            right:0;
        }
        #share-btn * {
            all: unset !important;
        }
        #share-btn-container div:nth-child(-n+2){
            width: auto !important;
            min-height: 0px !important;
        }
        #share-btn-container .wrap {
            display: none !important;
        }
"""
with gr.Blocks(css=css) as block:
    gr.HTML(
        """
            <div style="text-align: center; max-width: 700px; margin: 0 auto;">
              <div style="display: inline-flex; align-items: center; gap: 0.8rem; font-size: 1.75rem;">
                <h1 style="font-weight: 900; margin-bottom: 7px; line-height: normal;">
                  Parler-TTS: Português Brasileiro 🇧🇷
                </h1>
              </div>
            </div>
        """
    )
    gr.HTML(
        """
        <p>Este é um modelo de text-to-speech (TTS) adaptado para português brasileiro.
        Gere áudio de alta qualidade especificando a voz e o estilo desejado através de uma descrição em texto.</p>
        
        <p>Dicas para boas gerações:
        <ul>
            <li>Use a voz "Francisca" para resultados consistentes</li>
            <li>Você pode especificar diferentes emoções como: "profissional", "acolhedora", "enfática"</li>
            <li>Use pontuação para controlar as pausas na fala</li>
        </ul>
        </p>
        """
    )
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(label="Texto de Entrada", lines=2, value=default_text, elem_id="input_text")
            description = gr.Textbox(label="Descrição", lines=2, value=default_description, elem_id="input_description")
            run_button = gr.Button("Gerar Áudio", variant="primary")
        with gr.Column():
            audio_out = gr.Audio(label="Geração Parler-TTS", type="numpy", elem_id="audio_out")

    inputs = [input_text, description]
    outputs = [audio_out]
    gr.Examples(examples=examples, fn=gen_tts, inputs=inputs, outputs=outputs, cache_examples=True)
    run_button.click(fn=gen_tts, inputs=inputs, outputs=outputs, queue=True)
    gr.HTML(
        """
        <p>Para melhorar a prosódia e a naturalidade da fala, estamos ampliando a quantidade de dados de treinamento.
        A versão v1 do modelo será treinada com esses dados, além de otimizações de inferência que melhorarão a latência.
        Se você quiser saber mais sobre como este modelo foi treinado e até mesmo ajustá-lo, confira o 
        <a href="https://github.com/huggingface/parler-tts"> repositório Parler-TTS</a> no GitHub. O código e os checkpoints associados ao Parler-TTS estão licenciados sob <a href='https://github.com/huggingface/parler-tts?tab=Apache-2.0-1-ov-file#readme'> Apache 2.0</a>.</p>
        """
    )

block.queue()
block.launch(share=True)