from io import BytesIO
from typing import Tuple
import wave
import gradio as gr
import numpy as np
from pydub.audio_segment import AudioSegment
import requests
from os.path import exists
from stt import Model

import torch
from transformers import AutoModelForCTC, Wav2Vec2Processor

import torchaudio
from speechbrain.pretrained import EncoderClassifier

# initialize language ID model
lang_classifier = EncoderClassifier.from_hparams(
    source="speechbrain/lang-id-commonlanguage_ecapa", 
    savedir="pretrained_models/lang-id-commonlanguage_ecapa"
)


# download STT model
model_info = {
    "mixteco": ("https://coqui.gateway.scarf.sh/mixtec/jemeyer/v1.0.0/model.tflite", "mixtec.tflite"),
    "chatino": ("https://coqui.gateway.scarf.sh/chatino/bozden/v1.0.0/model.tflite", "chatino.tflite"),
    "totonaco": ("https://coqui.gateway.scarf.sh/totonac/bozden/v1.0.0/model.tflite", "totonac.tflite"),
    "español": ("jonatasgrosman/wav2vec2-large-xlsr-53-spanish", "spanish_xlsr"),
    "inglés": ("facebook/wav2vec2-large-robust-ft-swbd-300h", "english_xlsr"),
}


def client(audio_data: np.array, sample_rate: int, default_lang: str):
    output_audio = _convert_audio(audio_data, sample_rate)
    waveform, _ = torchaudio.load(output_audio)
    out_prob, score, index, text_lab = lang_classifier.classify_batch(waveform)

    output_audio.seek(0)
    fin = wave.open(output_audio, 'rb')
    audio = np.frombuffer(fin.readframes(fin.getnframes()), np.int16)

    fin.close()

    if text_lab == 'Spanish':
        processor, model = STT_MODELS['español']
        inputs = processor(waveform)
        logits = model(inputs.input_values, attention_mask=inputs.attention_mask).logits
        result = processor.decode(torch.argmax(logits, dim=-1).cpu().tolist())

    else:
        ds = STT_MODELS[default_lang]
        result = ds.stt(audio)

    return f"{text_lab}: {result}"


def load_models(language):

    model_path, file_name = model_info.get("language", ("", ""))

    if model_path.startswith('http'):
        if not exists(file_name):
            print(f"Downloading {model_path}")
            r = requests.get(model_path, allow_redirects=True)
            with open(file_name, 'wb') as file:
                file.write(r.content)
        else:
            print(f"Found {file_name}. Skipping download...")
        return Model(file_name)

    processor = Wav2Vec2Processor.from_pretrained(model_path)
    model = AutoModelForCTC.from_pretrained(model_path)
    return processor, model



def stt(default_lang: str, audio: Tuple[int, np.array]):
    sample_rate, audio = audio
    use_scorer = False

    recognized_result = client(audio, sample_rate, default_lang)

    return recognized_result


def _convert_audio(audio_data: np.array, sample_rate: int):
    source_audio = BytesIO()
    source_audio.write(audio_data)
    source_audio.seek(0)
    output_audio = BytesIO()
    wav_file = AudioSegment.from_raw(
        source_audio,
        channels=1,
        sample_width=2,
        frame_rate=sample_rate
    )
    wav_file.set_frame_rate(16000).set_channels(1).export(output_audio, "wav", codec="pcm_s16le")
    output_audio.seek(0)
    return output_audio


iface = gr.Interface(
    fn=stt,
    inputs=[
        gr.inputs.Radio(choices=("chatino", "mixteco", "totonaco"), default="mixteco", label="Lengua principal"),
        gr.inputs.Audio(type="numpy", label="Audio", optional=False),
    ],
    outputs=gr.outputs.Textbox(label="Output"),
    title="Coqui STT Yoloxochitl Mixtec",
    theme="huggingface",
    description="Prueba de dictado a texto para el mixteco de Yoloxochitl,"
                " usando [el modelo entrenado por Josh Meyer](https://coqui.ai/mixtec/jemeyer/v1.0.0/)"
                " con [los datos recopilados por Rey Castillo y sus colaboradores](https://www.openslr.org/89)."
                " Esta prueba es basada en la de [Ukraniano](https://huggingface.co/spaces/robinhad/ukrainian-stt)."
                " \n\n"
                "Speech-to-text demo for Yoloxochitl Mixtec,"
                " using [the model trained by Josh Meyer](https://coqui.ai/mixtec/jemeyer/v1.0.0/)"
                " on [the corpus compiled by Rey Castillo and collaborators](https://www.openslr.org/89)."
                " This demo is based on the [Ukrainian STT demo](https://huggingface.co/spaces/robinhad/ukrainian-stt).",
)

STT_MODELS = {lang: load_models(lang) for lang in ("inglés", "español")}

iface.launch()