import gradio as gr
import os
from pytube import YouTube

def cmd(url):
    # 创建 YouTube 对象
    yt = YouTube(url)

    # 获取视频标题
    title = yt.title.replace('.','')
    print(f'视频标题：{title}')

    # 获取视频时长
    duration = yt.length
    print(f'视频时长：{duration} 秒')

    # 过滤视频流，只保留渐进式下载和 MP4 格式
    streams = yt.streams.filter(progressive=True, file_extension='mp4')

    # 按分辨率降序排列视频流
    streams = streams.order_by('resolution').desc()

    # 选择第一个视频流并下载
    stream = streams.first()

    # 初始化循环变量
    retry_count = 0

    while retry_count < 3:
        try:
            # 获取视频大小
            file_size = stream.filesize
            print(f'视频大小：{file_size} 字节')

            # 创建下载目录
            download_dir = os.path.join(os.getcwd(), 'downloads')
            if not os.path.exists(download_dir):
                os.makedirs(download_dir)

            # 显示下载进度条
            stream.download(download_dir, filename=f'{title}.MP4')

            # 打印下载完成信息
            print('视频下载完成！')

            # 下载成功，退出循环
            #break
            return download_dir+f'/{title}.MP4'
        except Exception as e:
            print(f'视频下载失败：{e}')

        # 下载失败，增加重试次数
        retry_count += 1

    # 如果下载失败，打印提示信息
    if retry_count == 3:
        print('视频下载失败，请重试。')
    

# 创建 Gradio 应用程序
app = gr.Interface(fn=cmd, inputs=gr.Textbox(label="YouTube URL",lines=1), outputs=[gr.Video(label="Download the output Video")])

# 启动 Gradio 应用程序
app.launch()