import gradio as gr
from transformers import pipeline
import numpy as np

# --- Model Loading ---
MODEL_PATH = "./whisper_afrispeech_twi_model"

print("Loading ASR model...")
asr_pipeline = pipeline("automatic-speech-recognition", model=MODEL_PATH)
print("Model loaded successfully!")

def transcribe_audio(audio):
    """
    Transcribe audio using the fine-tuned Whisper model.
    
    Args:
        audio: Can be either:
            - A file path (string) when audio is uploaded
            - A tuple of (sample_rate, audio_data) when recorded
            - None if no audio provided
    
    Returns:
        Transcription text
    """
    if audio is None:
        return "No audio provided. Please record or upload audio."
    
    try:
        # Gradio automatically handles resampling to 16kHz for the pipeline
        # Just pass the audio directly to the pipeline
        result = asr_pipeline(audio)
        return result["text"]
    
    except Exception as e:
        return f"Error during transcription: {str(e)}"

# --- Create Gradio Interface ---
with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 🎙️ Twi ASR Transcription App
        Record or upload Twi speech and get the transcription using a fine-tuned Whisper model.
        """
    )
    
    with gr.Tab("🎤 Record Audio"):
        with gr.Row():
            with gr.Column():
                audio_input_mic = gr.Audio(
                    sources=["microphone"],
                    type="filepath",
                    label="Record Audio"
                )
                transcribe_btn_mic = gr.Button("Transcribe Recording", variant="primary")
            
            with gr.Column():
                output_text_mic = gr.Textbox(
                    label="Transcription",
                    placeholder="Your transcription will appear here...",
                    lines=5
                )
        
        transcribe_btn_mic.click(
            fn=transcribe_audio,
            inputs=audio_input_mic,
            outputs=output_text_mic
        )
    
    with gr.Tab("📤 Upload Audio"):
        with gr.Row():
            with gr.Column():
                audio_input_upload = gr.Audio(
                    sources=["upload"],
                    type="filepath",
                    label="Upload Audio File"
                )
                transcribe_btn_upload = gr.Button("Transcribe Upload", variant="primary")
            
            with gr.Column():
                output_text_upload = gr.Textbox(
                    label="Transcription",
                    placeholder="Your transcription will appear here...",
                    lines=5
                )
        
        transcribe_btn_upload.click(
            fn=transcribe_audio,
            inputs=audio_input_upload,
            outputs=output_text_upload
        )
    
    gr.Markdown(
        """
        ---
        Powered by Fine-tuned Whisper on AfriSpeech-Twi 🇬🇭
        """
    )

# --- Launch the app ---
if __name__ == "__main__":
    demo.launch()  # share=True creates a public link for testing