import gradio as gr
import yt_dlp
import os
import zipfile

def download_video(url, quality, format_choice, download_playlist, cookies_file):
    try:
        ydl_opts = {
            'outtmpl': 'downloads/%(title)s.%(ext)s',
            'noplaylist': not download_playlist,
        }
        
        if cookies_file is not None:
            cookie_path = "temp_cookies.txt"
            with open(cookie_path, "wb") as f:
                f.write(cookies_file)
            ydl_opts['cookiefile'] = cookie_path
        
        if quality == "Best":
            ydl_opts['format'] = 'bestvideo+bestaudio/best'
        elif quality == "Medium":
            ydl_opts['format'] = 'bestvideo[height<=720]+bestaudio/best[height<=720]'
        else:
            ydl_opts['format'] = 'bestvideo[height<=480]+bestaudio/best[height<=480]'
        
        if format_choice == "MP4":
            ydl_opts['merge_output_format'] = 'mp4'
        elif format_choice == "MKV":
            ydl_opts['merge_output_format'] = 'mkv'
        else:
            ydl_opts['format'] = 'bestaudio'
            ydl_opts['postprocessors'] = [{
                'key': 'FFmpegExtractAudio',
                'preferredcodec': 'mp3',
                'preferredquality': '192',
            }]
        
        if not os.path.exists('downloads'):
            os.makedirs('downloads')
            
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info = ydl.extract_info(url, download=True)
            filename = ydl.prepare_filename(info)
        
        if cookies_file is not None and os.path.exists(cookie_path):
            os.remove(cookie_path)
        
        # Nếu là playlist, nén toàn bộ thư mục downloads thành ZIP
        if download_playlist and os.path.isdir('downloads'):
            zip_path = "downloads/playlist.zip"
            with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
                for root, _, files in os.walk('downloads'):
                    for file in files:
                        if file != "playlist.zip":  # Tránh nén chính file ZIP
                            zipf.write(os.path.join(root, file), file)
            return "Playlist đã tải thành công dưới dạng ZIP", zip_path
        
        return f"Video đã tải thành công: {os.path.basename(filename)}", filename
    
    except Exception as e:
        return f"Lỗi: {str(e)}", None

interface = gr.Interface(
    fn=download_video,
    inputs=[
        gr.Textbox(label="Dán URL video vào đây", placeholder="https://..."),
        gr.Dropdown(choices=["Best", "Medium", "Low"], label="Chất lượng video", value="Best"),
        gr.Dropdown(choices=["MP4", "MKV", "Audio Only (MP3)"], label="Định dạng đầu ra", value="MP4"),
        gr.Checkbox(label="Tải toàn bộ playlist (nếu có)", value=False),
        gr.File(label="Upload file cookie (nếu cần)", type="binary")
    ],
    outputs=[
        gr.Textbox(label="Kết quả"),
        gr.File(label="Tải file về máy")
    ],
    title="Trình tải video trên Hugging Face Spaces",
    description="Dán URL video và upload file cookie nếu cần. Tải file hoặc playlist (ZIP) về máy."
)

interface.launch()