import gradio as gr
import os
from transformers import pipeline
import torch

print("Loading Amharic Whisper model...")
print("This may take a few minutes on first run...")

# Load the model with proper settings for Spaces
pipe = pipeline(
    "automatic-speech-recognition",
    model="chappM/whisper-amharic-small-v2",
    device=-1,  # Use CPU
    torch_dtype=torch.float32
)

print("Model ready!")

def transcribe_audio(audio_file):
    """Transcribe audio to Amharic text"""
    
    if audio_file is None:
        return "⚠️ Please record or upload an audio file first."
    
    try:
        # Get the file path
        if isinstance(audio_file, str):
            audio_path = audio_file
        else:
            audio_path = audio_file
        
        print(f"Processing: {audio_path}")
        
        # Check if file exists and has content
        if not os.path.exists(audio_path):
            return "❌ Audio file not found. Please try again."
        
        # Get file size
        file_size = os.path.getsize(audio_path)
        if file_size < 1000:  # Less than 1KB
            return "❌ Audio file is too small. Please record a longer clip."
        
        # Transcribe with Amharic language forced
        result = pipe(
            audio_path,
            generate_kwargs={"language": "am", "task": "transcribe"},
            return_timestamps=True
        )
        
        transcription = result["text"].strip()
        
        if transcription == "":
            return "❌ No speech detected. Please speak clearly in Amharic."
        
        # Return with formatting
        return f"📝 **Transcription:**\n\n{transcription}"
        
    except Exception as e:
        return f"❌ Error: {str(e)}\n\nPlease make sure your audio is clear and try again."

# Create the interface
with gr.Blocks(title="🎙️ Amharic Speech-to-Text", theme="soft") as demo:
    gr.Markdown("""
    # 🎙️ Amharic Speech-to-Text
    
    **Record or upload Amharic speech and get accurate transcriptions.**
    
    This model is fine-tuned specifically for Amharic language using the Whisper small architecture.
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            audio_input = gr.Audio(
                type="filepath",
                label="🎤 Record or Upload Amharic Speech",
                sources=["microphone", "upload"]
            )
            submit_btn = gr.Button("🚀 Transcribe", variant="primary", size="lg")
            clear_btn = gr.Button("🗑️ Clear", variant="stop")
            
        with gr.Column(scale=1):
            output_text = gr.Textbox(
                label="📝 Transcription",
                lines=10,
                placeholder="Your Amharic transcription will appear here...",
                interactive=False
            )
    
    with gr.Row():
        gr.Markdown("""
        ### 💡 Tips:
        - Speak clearly and at a normal pace
        - Short phrases (5-15 seconds) work best
        - Try to minimize background noise
        - Supported formats: WAV, MP3, M4A, FLAC, OGG
        """)
    
    # Event handlers
    submit_btn.click(
        fn=transcribe_audio,
        inputs=[audio_input],
        outputs=[output_text]
    )
    
    clear_btn.click(
        fn=lambda: "",
        outputs=[output_text]
    )
    
    # Also transcribe when audio is uploaded/recorded (optional)
    audio_input.change(
        fn=transcribe_audio,
        inputs=[audio_input],
        outputs=[output_text]
    )

# This is CRITICAL for Hugging Face Spaces
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", 7860))
    )