# app.py
from pyharp import *

import torch
import whisper
import numpy as np
import tempfile
import soundfile as sf
import gradio as gr

# -----------------------------
# Load Whisper model
# -----------------------------
MODEL_NAME = "base"  # options: tiny, base, small, medium, large
device = "cuda" if torch.cuda.is_available() else "cpu"

print(f"Loading Whisper model ({MODEL_NAME}) on {device}...")
model = whisper.load_model(MODEL_NAME, device=device)

# -----------------------------
# Helper: transcribe waveform
# -----------------------------
def transcribe_waveform(waveform: np.ndarray, sample_rate: int) -> str:
    # Convert stereo → mono
    if waveform.ndim == 2:
        waveform = np.mean(waveform, axis=0)

    # Whisper expects audio file OR 16k waveform
    # Easiest + safest: write temp WAV file
    with tempfile.NamedTemporaryFile(suffix=".wav") as tmp:
        sf.write(tmp.name, waveform, sample_rate)

        result = model.transcribe(tmp.name)

    return result["text"].strip()

# -----------------------------
# PyHARP process function
# -----------------------------
def process_fn(audio_path, params=None):
    result = model.transcribe(audio_path)
    return {"transcript": result["text"].strip()}

# -----------------------------
# Build PyHARP endpoint
# -----------------------------
import gradio as gr

if __name__ == "__main__":
    with gr.Blocks() as demo:

        model_card = ModelCard(
            name="Whisper Speech-to-Text",
            description="OpenAI Whisper STT with PyHARP",
            author="Aydin",
            tags=["speech", "transcription"]
        )

        input_components = [
            gr.Audio(type="filepath", label="Input Audio")
        ]

        output_components = [
            gr.Textbox(label="Transcription")
        ]

        build_endpoint(
            model_card,
            input_components,
            output_components,
            process_fn
        )

    demo.queue()
    demo.launch()