import os
import subprocess

# Ensure required dependencies are installed
try:
    import whisper
except ImportError:
    subprocess.run(["pip", "install", "openai-whisper"])
    import whisper

import gradio as gr
import torch
import librosa
from pydub import AudioSegment
from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC

# Load Models
models = {
    "Odia (AI4Bharat IndicWav2Vec)": {
        "processor": Wav2Vec2Processor.from_pretrained("ai4bharat/indicwav2vec-odia"),
        "model": Wav2Vec2ForCTC.from_pretrained("ai4bharat/indicwav2vec-odia").to("cuda" if torch.cuda.is_available() else "cpu"),
    },
    "Hindi (AI4Bharat IndicWav2Vec)": {
        "processor": Wav2Vec2Processor.from_pretrained("ai4bharat/indicwav2vec-hindi"),
        "model": Wav2Vec2ForCTC.from_pretrained("ai4bharat/indicwav2vec-hindi").to("cuda" if torch.cuda.is_available() else "cpu"),
    },
    "English (Facebook Wav2Vec2-960h)": {
        "processor": Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-large-960h"),
        "model": Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-large-960h").to("cuda" if torch.cuda.is_available() else "cpu"),
    }
}

# Function to convert audio to WAV (if needed)
def convert_to_wav(audio_file):
    file_ext = os.path.splitext(audio_file)[-1].lower()
    if file_ext != ".wav":
        temp_wav = "converted_temp.wav"
        sound = AudioSegment.from_file(audio_file, format=file_ext[1:])
        sound = sound.set_channels(1).set_frame_rate(16000)  # Convert to mono 16kHz
        sound.export(temp_wav, format="wav")
        return temp_wav
    return audio_file

# Transcription Function with Time Snippets
def transcribe(audio, language):
    # Convert audio to WAV if needed
    audio_path = convert_to_wav(audio)
    
    # Load and process audio
    audio, sr = librosa.load(audio_path, sr=16000)
    duration = librosa.get_duration(y=audio, sr=sr)
    
    # Select appropriate model
    selected_model = models[f"{language.capitalize()} (AI4Bharat IndicWav2Vec)" if language != "english" else "English (Facebook Wav2Vec2-960h)"]
    processor = selected_model["processor"]
    model = selected_model["model"]
    
    inputs = processor(audio, sampling_rate=16000, return_tensors="pt", padding=True)
    with torch.no_grad():
        logits = model(**inputs).logits
    
    predicted_ids = torch.argmax(logits, dim=-1)
    transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
    
    # Clean up temporary file if created
    if audio_path == "converted_temp.wav":
        os.remove(audio_path)
    
    return transcription

# Gradio Interface
app = gr.Interface(
    fn=transcribe,
    inputs=[
        gr.Audio(type="filepath", label="Upload Audio File"),
        gr.Dropdown(choices=["odia", "hindi", "english"], label="Select Language", value="odia"),
    ],
    outputs=gr.Textbox(label="Transcription"),
    title="Multilingual ASR Web App",
    description="Upload an audio file and select a language (Odia, Hindi, or English) to generate transcription.",
)

# Launch the App
app.launch()