import gradio as gr
import os
import time
import sys
import subprocess

# Clone and install faster-whisper from GitHub
subprocess.run(["git", "clone", "https://github.com/SYSTRAN/faster-whisper.git"], check=True)
subprocess.run(["pip", "install", "-e", "./faster-whisper"], check=True)

# Add the faster-whisper directory to the Python path
sys.path.append("./faster-whisper")

from faster_whisper import WhisperModel
from faster_whisper.transcribe import BatchedInferencePipeline

def transcribe_audio(audio_path, batch_size):
    # Initialize the model
    model = WhisperModel("cstr/whisper-large-v3-turbo-int8_float32", device="auto", compute_type="int8")
    batched_model = BatchedInferencePipeline(model=model)

    # Benchmark transcription time
    start_time = time.time()
    segments, info = batched_model.transcribe(audio_path, batch_size=batch_size)
    end_time = time.time()

    # Generate transcription
    transcription = ""
    for segment in segments:
        transcription += f"[{segment.start:.2f}s -> {segment.end:.2f}s] {segment.text}\n"

    # Calculate metrics
    transcription_time = end_time - start_time
    real_time_factor = info.duration / transcription_time
    audio_file_size = os.path.getsize(audio_path) / (1024 * 1024)  # Size in MB

    # Prepare output
    output = f"Transcription:\n\n{transcription}\n"
    output += f"\nLanguage: {info.language}, Probability: {info.language_probability:.2f}\n"
    output += f"Duration: {info.duration:.2f}s, Duration after VAD: {info.duration_after_vad:.2f}s\n"
    output += f"Transcription time: {transcription_time:.2f} seconds\n"
    output += f"Real-time factor: {real_time_factor:.2f}x\n"
    output += f"Audio file size: {audio_file_size:.2f} MB"

    return output

# Gradio interface
iface = gr.Interface(
    fn=transcribe_audio,
    inputs=[
        gr.Audio(type="filepath", label="Upload Audio File"),
        gr.Slider(minimum=1, maximum=32, step=1, value=16, label="Batch Size")
    ],
    outputs=gr.Textbox(label="Transcription and Metrics"),
    title="Faster Whisper Transcription",
    description="Upload an audio file to transcribe using Faster Whisper v3 turbo int8. Adjust the batch size for performance tuning.",
    examples=[["path/to/example/audio.mp3", 16]],
)

iface.launch()