main.py

import os
import uuid
import logging
from pathlib import Path
import subprocess

import gradio as gr
import yt_dlp

# --- Logging Setup ---
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s"
)

logging.getLogger("urllib3").setLevel(logging.WARNING)
logging.getLogger("httpx").setLevel(logging.WARNING)

logger = logging.getLogger(__name__)

# --- Load YouTube cookies from HF Secrets ---
if "YOUTUBE_COOKIES" in os.environ:
    with open("/tmp/cookies.txt", "w", encoding="utf-8") as f:
        f.write(os.environ["YOUTUBE_COOKIES"])


def download_media(url: str, format_choice: str, progress=gr.Progress()):
    if not url:
        raise gr.Error("Please enter a valid YouTube URL.")

    logger.info(f"New download request. URL: {url}, Format: {format_choice}")

    # Unique session folder
    session_id = str(uuid.uuid4())
    out_dir = Path("/tmp/downloads") / session_id
    out_dir.mkdir(parents=True, exist_ok=True)

    logger.info(f"Created temporary directory: {out_dir}")

    def progress_hook(d):
        if d["status"] == "downloading":
            total = d.get("total_bytes") or d.get("total_bytes_estimate", 0)
            downloaded = d.get("downloaded_bytes", 0)

            if total > 0:
                progress(downloaded / total, desc="Downloading...")

        elif d["status"] == "finished":
            progress(1.0, desc="Processing with ffmpeg...")

    try:
        # --- Step 1: Extract Video Info ---
        info_opts = {
            "quiet": False,
            "verbose": True,
            "no_warnings": False,
            "cookiefile": "/tmp/cookies.txt"
            if os.path.exists("/tmp/cookies.txt")
            else None,
            "extractor_retries": 10,
            "skip_download": True,
            "noplaylist": True,
            "extractor_args": {
                "youtube": {
                    "player_skip": ["webpage", "configs"],
                }
            },
        }

        with yt_dlp.YoutubeDL(info_opts) as ydl:
            logger.info("Extracting video information...")
            info_dict = ydl.extract_info(
                                url,
                                download=False,
                                process=False
                            )

            formats = info_dict.get("formats", [])

            logger.info("----- AVAILABLE FORMATS -----")

            for f in formats:
                logger.info(
                    f"ID: {f.get('format_id')} | "
                    f"EXT: {f.get('ext')} | "
                    f"RES: {f.get('resolution')} | "
                    f"VCODEC: {f.get('vcodec')} | "
                    f"ACODEC: {f.get('acodec')}"
                )

            logger.info("-----------------------------")

        # --- Step 2: Download Options ---
        download_opts = {
            "outtmpl": str(out_dir / "%(title)s.%(ext)s"),
            "progress_hooks": [progress_hook],
            "quiet": True,
            "no_warnings": True,
            "cookiefile": "/tmp/cookies.txt"
            if os.path.exists("/tmp/cookies.txt")
            else None,
            "extractor_retries": 5,
            "extractor_args": {
                "youtube": {
                    "player_skip": ["webpage", "configs"],
                }
            },
        }

        # --- Audio Download ---
        if format_choice == "Audio (MP3)":
            download_opts.update({
                "format": "bestaudio/best",
                "postprocessors": [{
                    "key": "FFmpegExtractAudio",
                    "preferredcodec": "mp3",
                    "preferredquality": "192",
                }],
            })

        # --- Video Download ---
        else:
            download_opts.update({
                "format": "bestvideo+bestaudio/best",
                "merge_output_format": "mp4",
            })

        logger.info(
            f"Downloading using selector: {download_opts['format']}"
        )

        with yt_dlp.YoutubeDL(download_opts) as ydl:
            ydl.download([url])

        logger.info("Download completed successfully.")

        downloaded_files = list(out_dir.glob("*"))

        if downloaded_files:
            logger.info(f"Downloaded file: {downloaded_files[0]}")
            return str(downloaded_files[0])

        raise gr.Error("Download failed: No file generated.")

    except Exception as e:
        logger.error(
            f"yt-dlp exception: {str(e)}",
            exc_info=True
        )

        raise gr.Error(
            f"Download failed.\n\nError:\n{str(e)}"
        )


# --- Theme ---
custom_theme = gr.themes.Monochrome(
    font=[
        gr.themes.GoogleFont("Inter"),
        "ui-sans-serif",
        "system-ui",
        "sans-serif",
    ]
)

# --- UI ---
with gr.Blocks(
    title="YouTube Downloader",
    theme=custom_theme
) as app:

    gr.Markdown("# YouTube Cloud Downloader")

    gr.Markdown(
        "Download YouTube videos or audio directly."
    )

    url_input = gr.Textbox(
        label="YouTube URL",
        placeholder="https://www.youtube.com/watch?v=..."
    )

    format_input = gr.Radio(
        choices=[
            "Video (MP4)",
            "Audio (MP3)"
        ],
        label="Format",
        value="Video (MP4)"
    )

    download_btn = gr.Button(
        "Download",
        variant="primary",
        size="lg"
    )

    output_file = gr.File(
        label="Downloaded File"
    )

    download_btn.click(
        fn=download_media,
        inputs=[url_input, format_input],
        outputs=output_file
    )

# --- Launch ---
if __name__ == "__main__":
    app.launch(
        server_name="0.0.0.0",
        server_port=7860
    )