import gradio as gr
from transformers import pipeline
import yt_dlp as youtube_dl
import os
os.environ['PATH'] = 'C:\\Program Files (x86)\\ffmpeg-7.0-essentials_build\\bin;' + os.environ['PATH']

# Define the Whisper pipeline
whisper = pipeline(model="openai/whisper-large-v3")

def download_youtube_audio(youtube_url):
    # Options for yt-dlp
    ydl_opts = {
    'format': 'bestaudio/best',
    'postprocessors': [{
        'key': 'FFmpegExtractAudio',
        'preferredcodec': 'mp3',
        'preferredquality': '192',
    }],
    'ffmpeg_location': 'C:\\Program Files (x86)\\ffmpeg-7.0-essentials_build\\bin',  # Add this line
    'outtmpl': 'downloaded_audio.%(ext)s',
    'noplaylist': True,
}
    # Use yt-dlp to download the audio
    with youtube_dl.YoutubeDL(ydl_opts) as ydl:
        info_dict = ydl.extract_info(youtube_url, download=True)
        audio_file = ydl.prepare_filename(info_dict)
        # Replace the extension with 'mp3' since we know we're extracting to mp3
        audio_file = os.path.splitext(audio_file)[0] + '.mp3'

    return audio_file

def transcribe(audio_file_path, task):
    if task == "transcribe":
        # Load the audio file and run it through the Whisper model
        transcription = whisper(audio_file_path)
        # Cleanup downloaded audio after transcription
        os.remove(audio_file_path)
        return transcription["text"]
    else:
        # Assuming you would have a translation function or model
        pass  # Replace with your translation logic

with gr.Blocks() as demo:
    with gr.Tab("YouTube"):
        with gr.Group():
            youtube_url = gr.Textbox(placeholder="Paste the URL to a YouTube video here")
            task = gr.Radio(["transcribe", "translate"], label="Task", value="transcribe")
            submit_button = gr.Button("Submit")
        output = gr.Textbox(label="output 1")

    submit_button.click(
        fn=lambda url, task: transcribe(download_youtube_audio(url), task),
        inputs=[youtube_url, task],
        outputs=output
    )

if __name__ == "__main__":
    demo.launch(share=True)