import gradio as gr
import yt_dlp
import os
from openai import OpenAI

client = OpenAI(
    api_key=os.getenv("OPENAI_API_KEY"),
)

opnerouter_client =  OpenAI(
  base_url="https://openrouter.ai/api/v1",
  api_key=os.getenv("OPENROUTER_API_KEY"),
)

ydl_opts = {
    'outtmpl': 'demo.m4a',
    'format': 'm4a/bestaudio/best',
    'postprocessors': [{  # Extract audio using ffmpeg
        'key': 'FFmpegExtractAudio',
        'preferredcodec': 'm4a',
    }],
    # 'proxy': 'socks5://192.168.2.18:20170',
}

text = ''

def download_audio(url):
    if os.path.exists('demo.m4a'):
        os.remove('demo.m4a')
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        code = ydl.download([url])
    assert code == 0, "Failed to download audio"

    with open("demo.m4a", "rb") as f:
        transcription = client.audio.transcriptions.create(
            model="whisper-1", 
            file=f, 
            response_format="text"
        )
    global text
    text = transcription
    return transcription

def summarize_text():
    prompt = f"Please summarize the following article in 5 sentences or less use the language of the audio: [{text}]"
    print(prompt)
    completion = opnerouter_client.chat.completions.create(
        model="google/gemini-pro",
        messages=[
            {"role": "user", "content": prompt}
        ]
    )
    return completion.choices[0].message.content

with gr.Blocks() as demo:
    with gr.Column():
        name = gr.Textbox(label="Enter your youtube url")
        button_download = gr.Button("Download")

    with gr.Row():
        output = gr.TextArea(label="Output")

        with gr.Column():
            button_summary = gr.Button("Summary")
            summary = gr.TextArea(label="Summary")

    button_download.click(
        download_audio,
        inputs=[name],
        outputs=[output],
    )

    button_summary.click(
        summarize_text,
        inputs=[],
        outputs=[summary],
    )

demo.launch()