import os
import re
import uuid
import glob
import shutil
import zipfile
import subprocess

import gradio as gr

def backup_tiktok(url):
    """
    Downloads TikTok video(s) from a provided URL using yt-dlp.
    If multiple videos are downloaded (e.g., from a TikTok profile),
    return a zip of all videos. Otherwise, return the single video file.
    """

    # Validate URL
    if "tiktok.com" not in url.lower():
        return "Please provide a valid TikTok Account or Video URL (must contain 'tiktok.com')."

    # Create a unique temporary folder
    download_id = str(uuid.uuid4())
    download_folder = f"downloads_{download_id}"
    os.makedirs(download_folder, exist_ok=True)

    # Run yt-dlp command
    cmd = ["yt-dlp", url, "-o", f"{download_folder}/%(title)s.%(ext)s"]
    try:
        subprocess.run(cmd, check=True)
    except subprocess.CalledProcessError as e:
        return f"Error downloading: {str(e)}"

    # Gather all downloaded files
    video_files = glob.glob(f"{download_folder}/*")

    if len(video_files) == 0:
        return "No videos found."

    # If only one video, return it directly
    if len(video_files) == 1:
        return video_files[0]

    # Otherwise, zip all files and return the zip
    zip_name = f"tiktok_backup_{download_id}.zip"
    with zipfile.ZipFile(zip_name, "w", zipfile.ZIP_DEFLATED) as zipf:
        for file_path in video_files:
            arcname = os.path.basename(file_path)
            zipf.write(file_path, arcname=arcname)

    return zip_name


with gr.Blocks() as demo:
    gr.Markdown("# TikTok Account/Video Backup\n\nEnter a TikTok Account or Video URL to download the video(s).")
    
    with gr.Row():
        url_input = gr.Textbox(label="TikTok URL", placeholder="https://www.tiktok.com/@username/video/...")
        download_button = gr.Button("Download")
        
    output_file = gr.File(label="Downloaded File or ZIP")

    download_button.click(fn=backup_tiktok, inputs=url_input, outputs=output_file)

demo.launch()