import gradio as gr
from transformers import pipeline
from librosa import resample
import numpy as np
import os
import sys
import glob
import torch
from huggingface_hub import snapshot_download
from fairseq_chunking import infer_fairseq_with_chunking

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
fairseq_cache = {}


def load_dictionary(dict_path):
    """Load the fairseq dictionary file"""
    dictionary = {}
    special_tokens = ['<s>', '<pad>', '</s>', '<unk>']
    for i, token in enumerate(special_tokens):
        dictionary[i] = token
    with open(dict_path, 'r', encoding='utf-8') as f:
        for line in f:
            parts = line.strip().split()
            if len(parts) >= 1:
                token = parts[0]
                dictionary[len(dictionary)] = token
    return dictionary


def decode_predictions(emissions, dictionary, verbose=False):
    predictions = emissions.argmax(dim=-1)
    vocab_size = emissions.shape[-1]
    blank_idx = 0

    tokens = []
    prev_token_id = None

    for pred in predictions[0]:
        token_id = pred.item()
        if token_id == blank_idx:
            prev_token_id = None
            continue
        if token_id == prev_token_id or token_id < 4:
            if token_id != prev_token_id and token_id < 4:
                prev_token_id = None
            continue
        token = dictionary.get(token_id, f'<unk_{token_id}>')
        tokens.append(token)
        prev_token_id = token_id

    transcription = ''.join(tokens).replace('|', ' ').strip()
    return transcription


def load_fairseq_model(model_id):
    """Download the fairseq model repo, import the encoders, and load the checkpoint"""
    if model_id in fairseq_cache:
        return fairseq_cache[model_id]

    local_path = snapshot_download(model_id)
    if local_path not in sys.path:
        sys.path.insert(0, local_path)

    original_load = torch.load
    torch.load = lambda *args, **kwargs: original_load(*args, **{**kwargs, "weights_only": False})
    try:
        import fairseq_extra_encoders
        import fairseq
        checkpoint_path = os.path.join(local_path, "fairseq_checkpoint.pt")
        models, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path])
    finally:
        torch.load = original_load

    model = models[0]
    model.eval()
    model = model.to(device)

    dict_path = os.path.join(local_path, "dict.ltr.txt")
    dictionary = load_dictionary(dict_path)

    fairseq_cache[model_id] = (model, dictionary, cfg.task.normalize)
    return fairseq_cache[model_id]


def transcribe(input_audio, model_id):
    sr, speech = input_audio
    # Convert to mono if stereo
    if speech.ndim > 1:
        speech = speech.mean(axis=1)
    # Convert to float32 if needed
    if speech.dtype != "float32":
        speech = speech.astype(np.float32)
    # Resample if sampling rate is not 16kHz
    if sr != 16000:
        speech = resample(speech, orig_sr=sr, target_sr=16000)
        sr = 16000

    if "ebranch" in model_id:
        model, dictionary, normalize_audio = load_fairseq_model(model_id)
        output = infer_fairseq_with_chunking(
            audio=speech,
            sampling_rate=sr,
            model=model,
            dictionary=dictionary,
            device=device,
            normalize_audio=normalize_audio,
            chunk_length_s=30.0,
            stride_length_s=(5.0, 5.0),
            model_downsample_ratio=320.0,
            decode_fn=lambda emissions, dict: decode_predictions(emissions, dict, verbose=False)
        )
    else:
        pipe = pipeline(
            "automatic-speech-recognition",
            model=model_id,
            device="cpu"
        )
        output = pipe(speech, chunk_length_s=30, stride_length_s=5)['text']
    return output


model_ids_list = [
 "GetmanY1/wav2vec2-large-sami-cont-pt-22k-finetuned",
 "GetmanY1/wav2vec2-large-ebranch-sami-18k-finetuned-experimental"
]
gradio_app = gr.Interface(
    fn=transcribe,
    inputs=[
        gr.Audio(sources=["upload","microphone"]),
        gr.Dropdown(
            label="Model",
            value="GetmanY1/wav2vec2-large-sami-cont-pt-22k-finetuned",
            choices=model_ids_list
        )
        ],
    outputs="text",
    title="Sámi Automatic Speech Recognition",
    description ="Choose a model from the list."
)
if __name__ == "__main__":
    gradio_app.launch(server_name="0.0.0.0", server_port=7860)