import gradio as gr
import openai
import yt_dlp
import os
openai.api_key = os.environ['OPENAI_API_KEY']

def asr(url):
    # download audio
    # Options for youtube-dl
    ydl_opts = {
       'format': 'bestaudio/best',
        'outtmpl': 'audio_downloaded.%(ext)s',
        'no_continue': True,
    }

    # Create a youtube-dl object
    ydl = yt_dlp.YoutubeDL(ydl_opts)

    # Download the video
    info_dict = ydl.extract_info(url, download=True)
    if info_dict is not None:
        audio_file_name = "audio_downloaded.{}".format(info_dict["ext"])
    else:
        return "下载音频发生错误，请确认链接再试一次。", "Error downloading the audio. Check the URL and try again."
    audio_file= open(audio_file_name, "rb")
    try:
        transcript = openai.Audio.transcribe("whisper-1", audio_file)
    except:
        return "视频过大（超过25mb）了，无法处理。", "The audio file is too big (25mb)."
    output = openai.ChatCompletion.create(
    model="gpt-3.5-turbo",
    messages=[
            {"role": "user", "content": "Transcript: {transcript}. \n Translate the video conversation transcript into fluent Chinese. Chinese: ".format(transcript=transcript["text"])},
        ]
    )

    # delete the video
    os.system("rm {}".format(audio_file_name))
    return output['choices'][0]['message']['content'], transcript["text"]

title = """
轻声细译"""
# Create an instruction input component
instruction = """
<div style="border: 2px solid #000; padding: 10px; border-radius: 5px;">
一键输入视频链接，轻松实现中文翻译，畅享视频无障碍沟通 <span style="color: grey;">-- powered by OpenAI Whisper & ChatGPT.</span>.<br>

1.将视频链接（支持Twitter、YouTube）复制粘贴至输入框，点击提交（Submit）即可;
2.为保证翻译质量，目前仅支持处理时长不超过5分钟的短视频。
</div>"""
# Create a text input component
text_input = gr.inputs.Textbox()

demo = gr.Interface(fn=asr, 
                    inputs=gr.inputs.Textbox(label="粘贴视频链接"),
                    outputs=[
                        gr.outputs.Textbox(label="中文"),
                        gr.outputs.Textbox(label="英文")
                    ],
                    title=title,
                    description=instruction,theme='huggingface',)

demo.launch()