import os
import gradio as gr
import torch
import librosa
from utils import download_from_gdrive, save_as_docx

# Keeping your existing engine loading logic
engines = {"Whisper V3 Turbo": None, "Bhashini (AI4Bharat)": None, "Sarvam": None}

def get_audio_duration(file_path):
    """Returns duration in minutes and a formatted string."""
    try:
        duration_sec = librosa.get_duration(path=file_path)
        duration_min = duration_sec / 60
        return duration_min, f"{duration_min:.2f} min"
    except Exception:
        return 0, "Unknown duration"

def get_engine(model_choice):
    """Lazy loads the selected engine."""
    if engines[model_choice] is None:
        if model_choice == "Whisper V3 Turbo":
            from engines.whisper_engine import transcribe_whisper
            engines[model_choice] = transcribe_whisper
        elif model_choice == "Bhashini (AI4Bharat)":
            from engines.indic_engine import transcribe_indic
            engines[model_choice] = transcribe_indic
        elif model_choice == "Sarvam":
            from engines.sarvam_engine import transcribe_sarvam
            engines[model_choice] = transcribe_sarvam
    return engines[model_choice]

def process_audio(files, gdrive_link, model_choice, input_lang, progress=gr.Progress()):
    all_transcripts = []
    processed_files = []
    
    # 1. Handle Google Drive Link if provided
    if gdrive_link:
        progress(0, desc="Fetching from Google Drive...")
        path = download_from_gdrive(gdrive_link)
        if path:
            # Create a simple object with a .name attribute to match gr.File objects
            class MockFile:
                def __init__(self, name):
                    self.name = name
            processed_files.append(MockFile(path))

    # 2. Add uploaded files
    if files:
        processed_files.extend(files)

    if not processed_files:
        raise gr.Error("Please upload files or provide a G-Drive link.")

    # 3. Load Engine
    progress(0.1, desc=f"Loading {model_choice} weights...")
    transcribe_fn = get_engine(model_choice)
    
    # 4. Process each file
    for i, file in enumerate(processed_files):
        file_name = os.path.basename(file.name)
        _, duration_str = get_audio_duration(file.name)
        
        progress((i + 0.2) / len(processed_files), 
                 desc=f"Transcribing {file_name} ({duration_str})...")
        
        try:
            # Actually call the transcription function and store the result
            transcript_data = transcribe_fn(file.name)
            all_transcripts.append({"filename": file_name, "text": transcript_data})
        except Exception as e:
            print(f"Error processing {file_name}: {e}")
            all_transcripts.append({"filename": file_name, "text": f"Error: {str(e)}"})

    # 5. Create the Docx
    progress(0.9, desc="Generating Document...")
    output_path = save_as_docx(all_transcripts)
    return output_path

# Custom CSS for the UI
custom_css = """
.gradio-container { font-family: 'Noto Sans', sans-serif; }
#title-container { text-align: center; padding: 20px; }
"""

with gr.Blocks() as demo:
    gr.HTML("""
        <div id="title-container">
            <h1 style="font-size: 2.5em; margin-bottom: 0;">🎙️ Indic Transcribe Pro</h1>
            <p style="font-size: 1.2em; color: #666;">v1.6 - Fixed Empty Output Bug</p>
        </div>
    """)
    
    with gr.Row():
        with gr.Column():
            file_input = gr.File(label="Upload Audio Files", file_count="multiple")
            link_input = gr.Textbox(label="OR Paste Google Drive Link (Public)")
            
            lang_selector = gr.Dropdown(
                choices=["Auto-Detect", "Hindi", "English", "Punjabi", "Telugu"],
                value="Hindi",
                label="Primary Language (Forces script output)"
            )
            
            model_selector = gr.Radio(
                choices=["Whisper V3 Turbo", "Bhashini (AI4Bharat)", "Sarvam"],
                value="Whisper V3 Turbo",
                label="Select AI Model"
            )
            submit_btn = gr.Button("🚀 Start Transcription", variant="primary")
            
        with gr.Column():
            output_file = gr.File(label="Download Transcript (.docx)")

    submit_btn.click(
        fn=process_audio,
        inputs=[file_input, link_input, model_selector, lang_selector],
        outputs=[output_file]
    )

if __name__ == "__main__":
    # Theme and CSS moved here to launch() to satisfy Gradio 6.0
    demo.launch(
        theme=gr.themes.Soft(),
        css=custom_css
    )