app.py

import gradio as gr
import yt_dlp
import os
import re
import time
import json
import subprocess
from datetime import datetime
# Add imports for the other engines
from pytube import YouTube as PytubeYouTube
try:
    from pytubefix import YouTube as PytubefixYouTube
except ImportError:
    PytubefixYouTube = None

def time_to_seconds(time_str):
    if not time_str:
        return 0
    
    if re.match(r'^\d+$', time_str):
        return int(time_str)
    elif re.match(r'^\d+:\d+$', time_str):
        minutes, seconds = time_str.split(':')
        return int(minutes) * 60 + int(seconds)
    elif re.match(r'^\d+:\d+:\d+$', time_str):
        hours, minutes, seconds = time_str.split(':')
        return int(hours) * 3600 + int(minutes) * 60 + int(seconds)
    return 0

def seconds_to_time_str(seconds):
    hours, remainder = divmod(seconds, 3600)
    minutes, seconds = divmod(remainder, 60)
    return f"{hours:02d}:{minutes:02d}:{seconds:02d}"

# ====== YT-DLP Implementation ======

def get_video_formats(info):
    """Extract available video formats from yt-dlp info"""
    formats = []
    for f in info.get('formats', []):
        if f.get('resolution') != 'audio only' and f.get('vcodec') != 'none':
            height = f.get('height')
            if height:
                formats.append({
                    'format_id': f.get('format_id'),
                    'resolution': f"{height}p",
                    'fps': f.get('fps'),
                    'vcodec': f.get('vcodec'),
                    'acodec': f.get('acodec'),
                    'ext': f.get('ext'),
                    'filesize': f.get('filesize'),
                    'has_audio': f.get('acodec') != 'none'
                })
    
    return formats

def fetch_video_info_ytdlp(video_url, cookies_file=None):
    try:
        ydl_opts = {
            'quiet': True,
            'no_warnings': True,
            'skip_download': True,
        }
        
        # Add cookies file if provided
        if cookies_file and os.path.exists(cookies_file):
            ydl_opts['cookiefile'] = cookies_file
        
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info = ydl.extract_info(video_url, download=False)
            
            # Extract basic information
            title = info.get('title', '')
            length = info.get('duration', 0)
            author = info.get('uploader', '')
            views = info.get('view_count', 0)
            description = info.get('description', '')
            thumbnail = info.get('thumbnail', '')
            
            duration_str = seconds_to_time_str(length)
            
            formats = get_video_formats(info)
            
            all_resolutions = set()
            for f in formats:
                if f.get('resolution'):
                    all_resolutions.add(f.get('resolution'))
            
            qualities = sorted(list(all_resolutions), 
                              key=lambda x: int(x.replace('p', '')))
            
            print(f"Available qualities: {qualities}")
            
            if not qualities:
                qualities = ["360p", "480p", "720p", "1080p"]
            
            default_quality = qualities[-1] if qualities else "1080p"
            
            return "", title, author, f"{views:,}", duration_str, description[:500] + "..." if len(description) > 500 else description, thumbnail, gr.update(visible=True), gr.update(choices=qualities, value=default_quality)
    except Exception as e:
        return f"错误: {str(e)}", "", "", "", "", "", None, gr.update(visible=False), gr.update(choices=["360p", "480p", "720p", "1080p"], value="1080p")

def download_video_segment_ytdlp(video_url, start_time, end_time, quality, fps, progress=gr.Progress(), cookies_file=None):
    try:
        start_seconds = time_to_seconds(start_time)
        end_seconds = time_to_seconds(end_time)
        
        if end_seconds <= start_seconds:
            return "结束时间必须大于开始时间"
        
        # Extract video information
        progress(0.1, "正在获取视频信息...")
        ydl_opts = {
            'quiet': True,
            'no_warnings': True,
            'skip_download': True,
        }
        
        # Add cookies file if provided
        if cookies_file and os.path.exists(cookies_file):
            ydl_opts['cookiefile'] = cookies_file
        
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info = ydl.extract_info(video_url, download=False)
            title = info.get('title', 'video')
            safe_title = re.sub(r'[^\w\-_]', '_', title)
            formats = get_video_formats(info)
        
        resolution = quality
        fps_value = int(fps.replace("fps", ""))
        
        progress(0.3, f"正在选择视频流 (质量: {quality}, 帧率: {fps})...")
        
        selected_format = None
        for f in formats:
            if f['resolution'] == quality and f['fps'] and f['fps'] >= fps_value and f['has_audio']:
                selected_format = f
                break
        
        if not selected_format:
            progress(0.3, "正在选择最佳视频和音频流...")
            
            timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
            temp_video = f"temp_video_{timestamp}.mp4"
            temp_audio = f"temp_audio_{timestamp}.mp4"
            temp_output = f"temp_combined_{timestamp}.mp4"
            output_filename = f"{safe_title}_{timestamp}_{start_time.replace(':', '')}_{end_time.replace(':', '')}.mp4"
            
            progress(0.4, "正在下载视频...")
            video_ydl_opts = {
                'format': f'bestvideo[height<={quality.replace("p", "")}]',
                'outtmpl': temp_video,
                'quiet': True
            }
            # Add cookies file if provided
            if cookies_file and os.path.exists(cookies_file):
                video_ydl_opts['cookiefile'] = cookies_file
                
            with yt_dlp.YoutubeDL(video_ydl_opts) as ydl:
                ydl.download([video_url])
            
            progress(0.5, "正在下载音频...")
            audio_ydl_opts = {
                'format': 'bestaudio',
                'outtmpl': temp_audio,
                'quiet': True
            }
            with yt_dlp.YoutubeDL(audio_ydl_opts) as ydl:
                ydl.download([video_url])
            
            progress(0.6, "正在合并视频和音频...")
            output_file_temp = os.path.join(os.getcwd(), temp_output)
            cmd_combine = [
                'ffmpeg', '-y',
                '-i', temp_video,
                '-i', temp_audio,
                '-c', 'copy',
                output_file_temp
            ]
            subprocess.run(cmd_combine, check=True)
            
            progress(0.7, "正在剪辑视频片段...")
            output_file = os.path.join(os.getcwd(), output_filename)
            cmd_cut = [
                'ffmpeg', '-y',
                '-i', output_file_temp,
                '-ss', str(start_seconds),
                '-to', str(end_seconds),
                '-c:v', 'copy',
                '-c:a', 'copy',
                output_file
            ]
            subprocess.run(cmd_cut, check=True)
            
            os.remove(temp_video)
            os.remove(temp_audio)
            os.remove(output_file_temp)
            
            progress(1.0, "已完成!")
            return output_file
        
        format_id = selected_format['format_id']
        
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        temp_filename = f"temp_{timestamp}.mp4"
        output_filename = f"{safe_title}_{timestamp}_{start_time.replace(':', '')}_{end_time.replace(':', '')}.mp4"
        
        progress(0.4, "正在下载视频...")
        ydl_opts = {
            'format': format_id,
            'outtmpl': temp_filename,
            'quiet': True
        }
        # Add cookies file if provided
        if cookies_file and os.path.exists(cookies_file):
            ydl_opts['cookiefile'] = cookies_file
            
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            ydl.download([video_url])
        
        progress(0.7, "正在剪辑视频片段...")
        output_file = os.path.join(os.getcwd(), output_filename)
        cmd = [
            'ffmpeg', '-y',
            '-i', temp_filename,
            '-ss', str(start_seconds),
            '-to', str(end_seconds),
            '-c:v', 'copy',
            '-c:a', 'copy',
            output_file
        ]
        
        subprocess.run(cmd, check=True)
        
        os.remove(temp_filename)
        
        progress(1.0, "已完成!")
        return output_file
    
    except Exception as e:
        return f"错误: {str(e)}"

# ====== PyTube Implementation ======

def fetch_video_info_pytube(video_url):
    try:
        yt = PytubeYouTube(video_url)
        
        title = yt.title
        length = yt.length
        author = yt.author
        views = yt.views
        description = yt.description
        thumbnail = yt.thumbnail_url
        
        duration_str = seconds_to_time_str(length)
        
        progressive_streams = yt.streams.filter(progressive=True)
        adaptive_streams = yt.streams.filter(adaptive=True, only_video=True)
        
        all_resolutions = set()
        for s in progressive_streams:
            if s.resolution:
                all_resolutions.add(s.resolution)
        for s in adaptive_streams:
            if s.resolution:
                all_resolutions.add(s.resolution)
        
        qualities = sorted(list(all_resolutions), 
                          key=lambda x: int(x.replace('p', '')))
        
        print(f"Available qualities: {qualities}")
        
        if not qualities:
            qualities = ["360p", "480p", "720p", "1080p"]
        
        default_quality = qualities[-1] if qualities else "1080p"
        
        return "", title, author, f"{views:,}", duration_str, description[:500] + "..." if len(description) > 500 else description, thumbnail, gr.update(visible=True), gr.update(choices=qualities, value=default_quality)
    except Exception as e:
        return f"错误: {str(e)}", "", "", "", "", "", None, gr.update(visible=False), gr.update(choices=["360p", "480p", "720p", "1080p"], value="1080p")

def download_video_segment_pytube(video_url, start_time, end_time, quality, fps, progress=gr.Progress()):
    try:
        start_seconds = time_to_seconds(start_time)
        end_seconds = time_to_seconds(end_time)
        
        if end_seconds <= start_seconds:
            return "结束时间必须大于开始时间"
        
        progress(0.1, "正在获取视频信息...")
        yt = PytubeYouTube(video_url)
        
        title = yt.title
        safe_title = re.sub(r'[^\w\-_]', '_', title)
        
        resolution = quality
        fps_value = int(fps.replace("fps", ""))
        
        progress(0.3, f"正在选择视频流 (质量: {quality}, 帧率: {fps})...")
        
        streams = yt.streams.filter(progressive=True, res=quality)
        
        if not streams:
            print(f"No progressive stream found for {quality}, trying adaptive...")
            video_streams = yt.streams.filter(adaptive=True, only_video=True, res=quality)
            
            if video_streams:
                video_stream = video_streams.first()
                audio_stream = yt.streams.filter(only_audio=True).first()
                
                timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
                temp_video = f"temp_video_{timestamp}.mp4"
                temp_audio = f"temp_audio_{timestamp}.mp4"
                temp_output = f"temp_combined_{timestamp}.mp4"
                output_filename = f"{safe_title}_{timestamp}_{start_time.replace(':', '')}_{end_time.replace(':', '')}.mp4"
                
                progress(0.4, "正在下载视频流...")
                video_path = video_stream.download(filename=temp_video)
                
                progress(0.5, "正在下载音频流...")
                audio_path = audio_stream.download(filename=temp_audio)
                
                progress(0.6, "正在合并视频和音频...")
                
                output_file_temp = os.path.join(os.getcwd(), temp_output)
                cmd_combine = [
                    'ffmpeg', '-y',
                    '-i', video_path,
                    '-i', audio_path,
                    '-c', 'copy',
                    output_file_temp
                ]
                subprocess.run(cmd_combine, check=True)
                
                progress(0.7, "正在剪辑视频片段...")
                output_file = os.path.join(os.getcwd(), output_filename)
                cmd_cut = [
                    'ffmpeg', '-y',
                    '-i', output_file_temp,
                    '-ss', str(start_seconds),
                    '-to', str(end_seconds),
                    '-c:v', 'copy',
                    '-c:a', 'copy',
                    output_file
                ]
                subprocess.run(cmd_cut, check=True)
                
                os.remove(video_path)
                os.remove(audio_path)
                os.remove(output_file_temp)
                
                progress(1.0, "已完成!")
                return output_file
            else:
                streams = yt.streams.filter(progressive=True).order_by('resolution').desc()
        
        if not streams:
            return "无法找到合适的视频流"
        
        video_stream = streams.first()
        print(f"Selected stream: {video_stream.resolution} {video_stream.fps}fps")
        
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        output_filename = f"{safe_title}_{timestamp}_{start_time.replace(':', '')}_{end_time.replace(':', '')}.mp4"
        
        progress(0.4, "正在下载视频...")
        download_path = video_stream.download(filename=f"temp_{output_filename}")
        
        progress(0.7, "正在剪辑视频片段...")
        output_file = os.path.join(os.getcwd(), output_filename)
        
        cmd = [
            'ffmpeg', '-y',
            '-i', download_path,
            '-ss', str(start_seconds),
            '-to', str(end_seconds),
            '-c:v', 'copy',
            '-c:a', 'copy',
            output_file
        ]
        
        subprocess.run(cmd, check=True)
        
        os.remove(download_path)
        
        progress(1.0, "已完成!")
        return output_file
    
    except Exception as e:
        return f"错误: {str(e)}"

# ====== PyTubeFix Implementation ======

def fetch_video_info_pytubefix(video_url, po_token="", visitor_data=""):
    try:
        if PytubefixYouTube is None:
            return "错误: pytubefix 未安装", "", "", "", "", "", None, gr.update(visible=False), gr.update(choices=["360p", "480p", "720p", "1080p"], value="1080p")
        
        # Create YouTube object with token verification if tokens are provided
        if po_token and visitor_data:
            yt = PytubefixYouTube(
                video_url, 
                use_po_token=True, 
                po_token_verifier=lambda: (visitor_data, po_token)
            )
        else:
            yt = PytubefixYouTube(video_url)
        
        title = yt.title
        length = yt.length
        author = yt.author
        views = yt.views
        description = yt.description
        thumbnail = yt.thumbnail_url
        
        duration_str = seconds_to_time_str(length)
        
        progressive_streams = yt.streams.filter(progressive=True)
        adaptive_streams = yt.streams.filter(adaptive=True, only_video=True)
        
        all_resolutions = set()
        for s in progressive_streams:
            if s.resolution:
                all_resolutions.add(s.resolution)
        for s in adaptive_streams:
            if s.resolution:
                all_resolutions.add(s.resolution)
        
        qualities = sorted(list(all_resolutions), 
                          key=lambda x: int(x.replace('p', '')))
        
        print(f"Available qualities: {qualities}")
        
        if not qualities:
            qualities = ["360p", "480p", "720p", "1080p"]
        
        default_quality = qualities[-1] if qualities else "1080p"
        
        return "", title, author, f"{views:,}", duration_str, description[:500] + "..." if len(description) > 500 else description, thumbnail, gr.update(visible=True), gr.update(choices=qualities, value=default_quality)
    except Exception as e:
        return f"错误: {str(e)}", "", "", "", "", "", None, gr.update(visible=False), gr.update(choices=["360p", "480p", "720p", "1080p"], value="1080p")

def download_video_segment_pytubefix(video_url, start_time, end_time, quality, fps, progress=gr.Progress(), po_token="", visitor_data=""):
    try:
        if PytubefixYouTube is None:
            return "错误: pytubefix 未安装"
        
        start_seconds = time_to_seconds(start_time)
        end_seconds = time_to_seconds(end_time)
        
        if end_seconds <= start_seconds:
            return "结束时间必须大于开始时间"
        
        progress(0.1, "正在获取视频信息...")
        
        # Create YouTube object with token verification if tokens are provided
        if po_token and visitor_data:
            yt = PytubefixYouTube(
                video_url, 
                use_po_token=True, 
                po_token_verifier=lambda: (visitor_data, po_token)
            )
        else:
            yt = PytubefixYouTube(video_url)
        
        title = yt.title
        safe_title = re.sub(r'[^\w\-_]', '_', title)
        
        resolution = quality
        fps_value = int(fps.replace("fps", ""))
        progress(0.3, f"正在选择视频流 (质量: {quality}, 帧率: {fps})...")
        
        streams = yt.streams.filter(progressive=True, res=quality)
        
        if not streams:
            print(f"No progressive stream found for {quality}, trying adaptive...")
            video_streams = yt.streams.filter(adaptive=True, only_video=True, res=quality)
            
            if video_streams:
                video_stream = video_streams.first()
                audio_stream = yt.streams.filter(only_audio=True).first()
                
                timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
                temp_video = f"temp_video_{timestamp}.mp4"
                temp_audio = f"temp_audio_{timestamp}.mp4"
                temp_output = f"temp_combined_{timestamp}.mp4"
                output_filename = f"{safe_title}_{timestamp}_{start_time.replace(':', '')}_{end_time.replace(':', '')}.mp4"
                
                progress(0.4, "正在下载视频流...")
                video_path = video_stream.download(filename=temp_video)
                
                progress(0.5, "正在下载音频流...")
                audio_path = audio_stream.download(filename=temp_audio)
                
                progress(0.6, "正在合并视频和音频...")
                
                output_file_temp = os.path.join(os.getcwd(), temp_output)
                cmd_combine = [
                    'ffmpeg', '-y',
                    '-i', video_path,
                    '-i', audio_path,
                    '-c', 'copy',
                    output_file_temp
                ]
                subprocess.run(cmd_combine, check=True)
                
                progress(0.7, "正在剪辑视频片段...")
                output_file = os.path.join(os.getcwd(), output_filename)
                cmd_cut = [
                    'ffmpeg', '-y',
                    '-i', output_file_temp,
                    '-ss', str(start_seconds),
                    '-to', str(end_seconds),
                    '-c:v', 'copy',
                    '-c:a', 'copy',
                    output_file
                ]
                subprocess.run(cmd_cut, check=True)
                
                os.remove(video_path)
                os.remove(audio_path)
                os.remove(output_file_temp)
                
                progress(1.0, "已完成!")
                return output_file
            else:
                streams = yt.streams.filter(progressive=True).order_by('resolution').desc()
        
        if not streams:
            return "无法找到合适的视频流"
        
        video_stream = streams.first()
        print(f"Selected stream: {video_stream.resolution} {video_stream.fps}fps")
        
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        output_filename = f"{safe_title}_{timestamp}_{start_time.replace(':', '')}_{end_time.replace(':', '')}.mp4"
        
        progress(0.4, "正在下载视频...")
        download_path = video_stream.download(filename=f"temp_{output_filename}")
        
        progress(0.7, "正在剪辑视频片段...")
        output_file = os.path.join(os.getcwd(), output_filename)
        
        cmd = [
            'ffmpeg', '-y',
            '-i', download_path,
            '-ss', str(start_seconds),
            '-to', str(end_seconds),
            '-c:v', 'copy',
            '-c:a', 'copy',
            output_file
        ]
        
        subprocess.run(cmd, check=True)
        
        os.remove(download_path)
        
        progress(1.0, "已完成!")
        return output_file
    
    except Exception as e:
        return f"错误: {str(e)}"

# ====== Unified Interface Functions ======

def fetch_video_info(video_url, engine, po_token="", visitor_data="", cookies_file=None):
    if engine == "yt_dlp":
        return fetch_video_info_ytdlp(video_url, cookies_file)
    elif engine == "pytube":
        return fetch_video_info_pytube(video_url)
    elif engine == "pytubefix":
        return fetch_video_info_pytubefix(video_url, po_token, visitor_data)
    else:
        return f"错误: 未知引擎 {engine}", "", "", "", "", "", None, gr.update(visible=False), gr.update(choices=["360p", "480p", "720p", "1080p"], value="1080p")

def download_video_segment(video_url, start_time, end_time, quality, fps, engine, po_token="", visitor_data="", cookies_file=None, progress=gr.Progress()):
    if engine == "yt_dlp":
        return download_video_segment_ytdlp(video_url, start_time, end_time, quality, fps, progress, cookies_file)
    elif engine == "pytube":
        return download_video_segment_pytube(video_url, start_time, end_time, quality, fps, progress)
    elif engine == "pytubefix":
        return download_video_segment_pytubefix(video_url, start_time, end_time, quality, fps, progress, po_token, visitor_data)
    else:
        return f"错误: 未知引擎 {engine}"

# ====== Gradio UI ======

with gr.Blocks(title="YouTube 视频剪辑器", theme="NoCrypt/miku") as demo:
    gr.Markdown("# YouTube 视频剪辑工具")
    gr.Markdown("输入 YouTube 视频链接，查看视频信息并下载指定时间段的视频片段")
    
    with gr.Accordion(label="下载引擎", open=False):
        engine = gr.Radio(
            label="选择用于下载视频的库",
            choices=["yt_dlp", "pytube", "pytubefix"],
            value="yt_dlp"
        )

        with gr.Group(visible=True) as ytdlp_group:
            cookies_file = gr.File(
                label="Cookies 文件 (可选 - Netscape 格式，用于访问受限视频)", 
                file_count="single", 
                type="filepath",  # Change from "file" to "filepath"
                file_types=[".txt"]
            )
            
        with gr.Group(visible=False) as pytubefix_group:
            po_token = gr.Textbox(label="poToken (optional)", placeholder="INNERTUBE_API_KEY value")
            visitor_data = gr.Textbox(label="visitorData (optional)", placeholder="CgtsZG1hLXVpLW...")
    
    with gr.Row():
        video_url = gr.Textbox(label="YouTube 视频链接", placeholder="https://www.youtube.com/watch?v=...")
        fetch_button = gr.Button("获取视频信息", variant="primary")
    
    error_message = gr.Textbox(label="状态", visible=True, interactive=False)
    
    with gr.Group(visible=False) as video_info_group:
        with gr.Row():
            with gr.Column(scale=2):
                thumbnail = gr.Image(label="缩略图")
            
            with gr.Column(scale=3):
                title = gr.Textbox(label="标题", interactive=False)
                author = gr.Textbox(label="作者", interactive=False)
                
                with gr.Row():
                    views = gr.Textbox(label="播放量", interactive=False)
                    duration = gr.Textbox(label="时长", interactive=False)
                
                description = gr.Textbox(label="描述", interactive=False, lines=5)
        
        gr.Markdown("### 选择要下载的视频片段")
        
        with gr.Row():
            start_time = gr.Textbox(label="开始时间 (HH:MM:SS)", placeholder="00:01:30")
            end_time = gr.Textbox(label="结束时间 (HH:MM:SS)", placeholder="00:02:30")
        
        with gr.Row():
            quality = gr.Dropdown(label="视频质量", choices=["360p", "480p", "720p", "1080p"], value="1080p")
            fps = gr.Radio(label="帧率", choices=["30fps", "60fps"], value="60fps")
        
        download_button = gr.Button("下载视频片段", variant="primary")
        
        output_file = gr.File(label="下载结果")
    
    def update_engine_groups(engine_value):
        return gr.update(visible=(engine_value == "pytubefix")), gr.update(visible=(engine_value == "yt_dlp"))

    engine.change(
        fn=update_engine_groups,
        inputs=[engine],
        outputs=[pytubefix_group, ytdlp_group]
    )

    fetch_button.click(
        fn=fetch_video_info,
        inputs=[video_url, engine, po_token, visitor_data, cookies_file],
        outputs=[error_message, title, author, views, duration, description, thumbnail, video_info_group, quality]
    )

    download_button.click(
        fn=download_video_segment,
        inputs=[video_url, start_time, end_time, quality, fps, engine, po_token, visitor_data, cookies_file],
        outputs=[output_file]
    )
if __name__ == "__main__":
    demo.launch()