import gradio as gr
import os
import time
import sys
import subprocess
import tempfile
import requests
from urllib.parse import urlparse

# 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)
subprocess.run(["pip", "install", "yt-dlp"], 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
import yt_dlp

def download_audio(url):
    parsed_url = urlparse(url)
    if parsed_url.netloc == 'www.youtube.com' or parsed_url.netloc == 'youtu.be':
        # YouTube video
        ydl_opts = {
            'format': 'bestaudio/best',
            'postprocessors': [{
                'key': 'FFmpegExtractAudio',
                'preferredcodec': 'mp3',
                'preferredquality': '192',
            }],
            'outtmpl': '%(id)s.%(ext)s',
        }
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info = ydl.extract_info(url, download=True)
            return f"{info['id']}.mp3"
    else:
        # Direct MP3 URL
        response = requests.get(url)
        if response.status_code == 200:
            with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as temp_file:
                temp_file.write(response.content)
                return temp_file.name
        else:
            raise Exception(f"Failed to download audio from {url}")

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

    # Handle input source
    if isinstance(input_source, str) and (input_source.startswith('http://') or input_source.startswith('https://')):
        # It's a URL, download the audio
        audio_path = download_audio(input_source)
    else:
        # It's a local file path
        audio_path = input_source

    # 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"

    # Clean up downloaded file if it was a URL
    if isinstance(input_source, str) and (input_source.startswith('http://') or input_source.startswith('https://')):
        os.remove(audio_path)

    return output

# Gradio interface
iface = gr.Interface(
    fn=transcribe_audio,
    inputs=[
        gr.inputs.Textbox(label="Audio Source (Upload, MP3 URL, or YouTube URL)"),
        gr.Slider(minimum=1, maximum=32, step=1, value=16, label="Batch Size")
    ],
    outputs=gr.Textbox(label="Transcription and Metrics"),
    title="Faster Whisper v3 turbo int8 transcription",
    description="Enter an audio file path, MP3 URL, or YouTube URL to transcribe using Faster Whisper v3 turbo (int8). Adjust the batch size for performance tuning.",
    examples=[
        ["https://www.youtube.com/watch?v=dQw4w9WgXcQ", 16],
        ["https://example.com/path/to/audio.mp3", 16],
        ["path/to/local/audio.mp3", 16]
    ],
)

iface.launch()