import os

import gradio as gr
import spaces
from transformers import pipeline

MODEL_ID = "NCAIR1/NigerianAccentedEnglish"
HF_TOKEN = os.environ.get("HF_TOKEN") or None
API_SECRET = os.environ.get("API_SECRET", "")

# Loaded once on CPU at startup. Only moved to the GPU inside the
# @spaces.GPU-decorated function, for the duration of that call — that's how
# ZeroGPU's shared/ephemeral GPU allocation model works.
pipe = pipeline("automatic-speech-recognition", model=MODEL_ID, token=HF_TOKEN)


@spaces.GPU(duration=45)
def transcribe(audio_path, secret):
    if API_SECRET and secret != API_SECRET:
        raise gr.Error("Unauthorized")
    if audio_path is None:
        raise gr.Error("No audio data received.")

    pipe.model.to("cuda")
    try:
        result = pipe(
            audio_path,
            generate_kwargs={
                "language": "en",
                "task": "transcribe",
                "num_beams": 5,
                "no_repeat_ngram_size": 3,
            },
        )
    finally:
        pipe.model.to("cpu")

    text = result.get("text", "") if isinstance(result, dict) else ""
    return text.strip()


demo = gr.Interface(
    fn=transcribe,
    inputs=[
        gr.Audio(type="filepath", label="Audio chunk"),
        gr.Textbox(label="Secret", type="password"),
    ],
    outputs=gr.Textbox(label="Transcript"),
    api_name="transcribe",
)

if __name__ == "__main__":
    demo.queue().launch()