import gradio as gr
from yt_dlp import YoutubeDL
import os
import tempfile
import shutil

# -----------------------
# Helper Functions
# -----------------------
def download_tiktok_thumbnails(username_url, cookies_file=None, max_videos=20):
    tmp_dir = tempfile.mkdtemp()
    ydl_opts = {
        "skip_download": True,
        "write_thumbnail": True,
        "outtmpl": os.path.join(tmp_dir, "%(id)s.%(ext)s"),
        "quiet": True,
    }
    if cookies_file:
        ydl_opts["cookiefile"] = cookies_file.name

    try:
        with YoutubeDL(ydl_opts) as ydl:
            ydl.download([username_url])
    except Exception as e:
        return f"❌ Error: {e}", None, None

    # Collect all thumbnails
    thumbs_all = [os.path.join(tmp_dir, f) for f in os.listdir(tmp_dir)
                  if f.lower().endswith((".jpg", ".webp", ".png"))]

    # Apply limiter
    thumbs = thumbs_all[:max_videos]

    if not thumbs:
        return "❌ No thumbnails found", None, None

    return f"✅ Fetched {len(thumbs)} thumbnails", thumbs, tmp_dir

def prepare_zip(tmp_dir):
    zip_path = os.path.join(tempfile.gettempdir(), "tiktok_thumbnails.zip")
    shutil.make_archive(zip_path.replace(".zip",""), 'zip', tmp_dir)
    return zip_path

def fetch_and_zip_tiktok(username_url, cookies_file, max_videos):
    status, thumbs, tmp_dir = download_tiktok_thumbnails(username_url, cookies_file, max_videos)
    zip_file = prepare_zip(tmp_dir) if thumbs else None
    return status, thumbs, zip_file

# -----------------------
# Gradio Interface
# -----------------------
with gr.Blocks() as demo:
    gr.Markdown("## 🎵 TikTok Account Thumbnails Downloader with Cookies")
    gr.Markdown(
        "Upload your `cookies.txt` (exported from your browser) first, then enter a TikTok account URL.\n"
        "Thumbnails from private or age-restricted videos may require cookies."
    )

    cookies_upload = gr.File(label="Upload cookies.txt (required for private content)", file_types=[".txt"])
    url_input = gr.Textbox(label="TikTok Account URL", placeholder="https://www.tiktok.com/@username")
    max_videos_slider = gr.Slider(minimum=1, maximum=50, step=1, value=20, label="Max Videos to Fetch")
    fetch_btn = gr.Button("📥 Fetch Thumbnails")

    status_output = gr.Textbox(label="Status")
    thumbs_gallery = gr.Gallery(label="Thumbnails Preview", elem_id="thumbs_gallery", columns=5, height="auto")
    download_btn = gr.File(label="Download All Thumbnails (ZIP)")

    inputs_list = [url_input, cookies_upload, max_videos_slider]
    outputs_list = [status_output, thumbs_gallery, download_btn]

    fetch_btn.click(fetch_and_zip_tiktok, inputs=inputs_list, outputs=outputs_list)

if __name__ == "__main__":
    demo.launch()