import shutil
import tempfile
from pathlib import Path

import spaces
import gradio as gr
from huggingface_hub import hf_hub_download
from transformers import AutoProcessor, CohereAsrForConditionalGeneration
from transformers.audio_utils import load_audio

MODEL_ID = "CohereLabs/cohere-transcribe-03-2026"

processor = AutoProcessor.from_pretrained(MODEL_ID)
model = CohereAsrForConditionalGeneration.from_pretrained(MODEL_ID, device_map="auto")

LANGUAGES = {
    "English": "en",
    "French": "fr",
    "German": "de",
    "Italian": "it",
    "Spanish": "es",
    "Portuguese": "pt",
    "Greek": "el",
    "Dutch": "nl",
    "Polish": "pl",
    "Chinese (Mandarin)": "zh",
    "Japanese": "ja",
    "Korean": "ko",
    "Vietnamese": "vi",
    "Arabic": "ar",
}

_cached_audio = hf_hub_download(
    repo_id=MODEL_ID, filename="demo/voxpopuli_test_en_demo.wav"
)
EXAMPLE_AUDIO = str(Path(tempfile.gettempdir()) / Path(_cached_audio).name)
shutil.copy2(_cached_audio, EXAMPLE_AUDIO)


@spaces.GPU
def transcribe(audio_path: str | None, language: str) -> str:
    if audio_path is None:
        return ""
    audio = load_audio(audio_path, sampling_rate=16000)
    lang_code = LANGUAGES[language]
    inputs = processor(
        audio, sampling_rate=16000, return_tensors="pt", language=lang_code
    )
    inputs.to(model.device, dtype=model.dtype)
    audio_chunk_index = inputs.get("audio_chunk_index")
    outputs = model.generate(**inputs, max_new_tokens=1024)
    text = processor.decode(
        outputs,
        skip_special_tokens=True,
        audio_chunk_index=audio_chunk_index,
        language=lang_code,
    )
    if isinstance(text, list):
        text = text[0]
    return text.strip()


with gr.Blocks() as demo:
    gr.Markdown("# Cohere Transcribe")
    gr.Markdown(
        "Transcribe audio using"
        " [CohereLabs/cohere-transcribe-03-2026](https://huggingface.co/CohereLabs/cohere-transcribe-03-2026)."
    )
    with gr.Row():
        with gr.Column():
            audio_input = gr.Audio(label="Audio", type="filepath")
            language = gr.Dropdown(
                label="Language",
                choices=list(LANGUAGES.keys()),
                value="English",
            )
            btn = gr.Button("Transcribe")
        with gr.Column():
            output = gr.Textbox(label="Transcription", lines=10)

    btn.click(fn=transcribe, inputs=[audio_input, language], outputs=output)

    gr.Examples(
        examples=[[EXAMPLE_AUDIO, "English"]],
        inputs=[audio_input, language],
    )

if __name__ == "__main__":
    demo.launch(max_file_size="25MB")