import numpy as np
import torch
import gradio as gr
from transformers import WhisperForConditionalGeneration, WhisperProcessor
from transformers.pipelines.audio_utils import ffmpeg_read

MODEL_ID = "openai/whisper-small"
TARGET_SAMPLE_RATE = 16000

processor = WhisperProcessor.from_pretrained(MODEL_ID)
model = WhisperForConditionalGeneration.from_pretrained(MODEL_ID)
model.eval()


def _normalize_waveform(data: np.ndarray) -> np.ndarray:
    if data.ndim > 1:
        data = data.mean(axis=1)

    if np.issubdtype(data.dtype, np.integer):
        max_int = np.iinfo(data.dtype).max
        data = data.astype(np.float32) / float(max_int)
    else:
        data = data.astype(np.float32)
        peak = np.max(np.abs(data)) if data.size else 0.0
        if peak > 1.0:
            data = data / peak

    return data


def _extract_audio(audio_input):
    # Gradio may send audio as (sample_rate, np.ndarray), a filepath string, or a FileData-like dict.
    if isinstance(audio_input, tuple) and len(audio_input) == 2:
        sample_rate, data = audio_input
        return int(sample_rate), _normalize_waveform(data)

    path = None
    if isinstance(audio_input, str):
        path = audio_input
    elif isinstance(audio_input, dict):
        path = audio_input.get("path")

    if path:
        with open(path, "rb") as f:
            audio_bytes = f.read()
        data = ffmpeg_read(audio_bytes, TARGET_SAMPLE_RATE)
        return TARGET_SAMPLE_RATE, _normalize_waveform(data)

    raise ValueError("Unsupported audio input format. Please upload a valid audio file.")


def transcribe_audio(audio, task):
    if audio is None:
        return "Please upload or record audio first."

    sample_rate, data = _extract_audio(audio)

    inputs = processor(
        data,
        sampling_rate=sample_rate,
        return_tensors="pt",
    )

    forced_decoder_ids = processor.get_decoder_prompt_ids(task=task)

    with torch.inference_mode():
        predicted_ids = model.generate(
            inputs.input_features,
            forced_decoder_ids=forced_decoder_ids,
        )

    text = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0].strip()
    return text or "(No speech detected)"


with gr.Blocks(title="Whisper Small STT CPU") as demo:
    gr.Markdown(
        """
# Whisper Small STT (Free CPU)
Public speech-to-text using `openai/whisper-small`.

- `transcribe`: keep original language
- `translate`: translate speech to English
"""
    )

    audio_input = gr.Audio(
        label="Audio Input",
        type="numpy",
        sources=["upload", "microphone"],
    )
    task_input = gr.Dropdown(
        choices=["transcribe", "translate"],
        value="transcribe",
        label="Task",
    )
    run_btn = gr.Button("Convert Speech to Text", variant="primary")
    text_output = gr.Textbox(label="Transcript", lines=12)

    run_btn.click(
        fn=transcribe_audio,
        inputs=[audio_input, task_input],
        outputs=[text_output],
        api_name="transcribe",
    )

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