snapstream_gui.py

# -*- coding: utf-8 -*-
"""Unified GUI to download from YouTube or Twitch using yt-dlp."""

from __future__ import annotations

import sys
import threading
from pathlib import Path
from typing import Dict, List, Optional

import tkinter as tk
from tkinter import filedialog, messagebox, ttk
from tkinter.scrolledtext import ScrolledText

LEGAL_NOTE = (
    "Use this only for videos or streams you are entitled to download, "
    "respecting local laws and each platform's Terms of Service."
)

try:
    from yt_dlp import YoutubeDL
except Exception:  # pragma: no cover - external dependency
    print("Error: 'yt-dlp' not found. Install with: pip install yt-dlp", file=sys.stderr)
    raise


def _human_size(nbytes: Optional[int]) -> str:
    if not nbytes or nbytes <= 0:
        return "-"
    units = ["B", "KB", "MB", "GB", "TB", "PB"]
    size = float(nbytes)
    for unit in units:
        if size < 1024.0:
            return f"{size:.1f}{unit}"
        size /= 1024.0
    return f"{size:.1f}EB"


def _human_duration(seconds: Optional[float]) -> str:
    if seconds is None:
        return "-"
    seconds = int(seconds)
    hours, rem = divmod(seconds, 3600)
    minutes, secs = divmod(rem, 60)
    if hours:
        return f"{hours:02d}:{minutes:02d}:{secs:02d}"
    return f"{minutes:02d}:{secs:02d}"


class SnapStreamGUI(tk.Tk):
    def __init__(self) -> None:
        super().__init__()
        self.title("SnapStream - YouTube + Twitch")
        self.geometry("1040x720")
        self.minsize(900, 640)

        self.url_var = tk.StringVar()
        self.output_dir_var = tk.StringVar(value=str(Path.cwd()))
        self.quality_var = tk.StringVar(value="Auto (best)")
        self.mode_var = tk.StringVar(value="Video MP4")
        self.audio_quality_var = tk.StringVar(value="192")
        self.status_var = tk.StringVar(value="Ready")
        self.title_var = tk.StringVar(value="No video loaded")

        self.formats_cache: List[Dict] = []
        self.downloading = False

        self._build_style()
        self._build_ui()

    # ---------------------------- UI helpers ----------------------------
    def _build_style(self) -> None:
        style = ttk.Style(self)
        try:
            style.theme_use("clam")
        except tk.TclError:
            pass
        bg = "#1f2024"
        fg = "#f5f6f7"
        accent = "#5dd39e"
        style.configure("Main.TFrame", background=bg)
        style.configure("Main.TLabel", background=bg, foreground=fg, font=("Segoe UI", 10))
        style.configure("Title.TLabel", background=bg, foreground=fg, font=("Segoe UI", 18, "bold"))
        style.configure("Card.TLabelframe", background=bg, foreground=fg)
        style.configure("Card.TLabelframe.Label", background=bg, foreground=fg, font=("Segoe UI", 11, "bold"))
        style.configure("Accent.TButton", font=("Segoe UI", 10, "bold"), padding=6)
        style.map("Accent.TButton", background=[("!disabled", accent)])
        style.configure("TButton", padding=5)
        style.configure("Treeview", background="#232429", foreground=fg, fieldbackground="#232429", rowheight=26)
        style.configure("Treeview.Heading", font=("Segoe UI", 10, "bold"))
        style.map("Treeview", background=[("selected", "#2f80ed")])

    def _build_ui(self) -> None:
        main = ttk.Frame(self, padding=16, style="Main.TFrame")
        main.pack(fill=tk.BOTH, expand=True)

        header = ttk.Frame(main, style="Main.TFrame")
        header.pack(fill=tk.X, pady=(0, 14))
        ttk.Label(header, text="SnapStream", style="Title.TLabel").pack(side=tk.LEFT)
        ttk.Label(header, text="YouTube + Twitch in one place", style="Main.TLabel").pack(side=tk.LEFT, padx=12)
        ttk.Label(header, textvariable=self.status_var, style="Main.TLabel").pack(side=tk.RIGHT)

        top = ttk.LabelFrame(main, text="Source", padding=12, style="Card.TLabelframe")
        top.pack(fill=tk.X)

        url_row = ttk.Frame(top, style="Main.TFrame")
        url_row.pack(fill=tk.X, pady=4)
        ttk.Label(url_row, text="Video or stream URL:", style="Main.TLabel").pack(side=tk.LEFT)
        url_entry = ttk.Entry(url_row, textvariable=self.url_var, font=("Segoe UI", 10))
        url_entry.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=8)
        ttk.Button(url_row, text="Fetch formats", style="Accent.TButton", command=self.fetch_info).pack(side=tk.LEFT, padx=4)

        dest_row = ttk.Frame(top, style="Main.TFrame")
        dest_row.pack(fill=tk.X, pady=4)
        ttk.Label(dest_row, text="Save to:", style="Main.TLabel").pack(side=tk.LEFT)
        dest_entry = ttk.Entry(dest_row, textvariable=self.output_dir_var, width=50)
        dest_entry.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=8)
        ttk.Button(dest_row, text="Browse...", command=self.choose_output_dir).pack(side=tk.LEFT, padx=4)

        options = ttk.Frame(top, style="Main.TFrame")
        options.pack(fill=tk.X, pady=6)

        ttk.Label(options, text="Mode:", style="Main.TLabel").pack(side=tk.LEFT)
        mode_combo = ttk.Combobox(
            options,
            state="readonly",
            values=["Video MP4", "Audio only MP3"],
            textvariable=self.mode_var,
            width=16,
        )
        mode_combo.pack(side=tk.LEFT, padx=6)

        ttk.Label(options, text="Quality:", style="Main.TLabel").pack(side=tk.LEFT, padx=(14, 0))
        quality_combo = ttk.Combobox(
            options,
            state="readonly",
            values=["Auto (best)", "1080p", "720p", "480p", "360p"],
            textvariable=self.quality_var,
            width=14,
        )
        quality_combo.pack(side=tk.LEFT, padx=6)

        ttk.Label(options, text="Audio kbps:", style="Main.TLabel").pack(side=tk.LEFT, padx=(14, 0))
        audio_combo = ttk.Combobox(
            options,
            state="readonly",
            values=["320", "256", "192", "160", "128"],
            textvariable=self.audio_quality_var,
            width=8,
        )
        audio_combo.pack(side=tk.LEFT, padx=6)

        ttk.Label(options, textvariable=self.title_var, style="Main.TLabel").pack(side=tk.RIGHT)

        table_box = ttk.LabelFrame(
            main,
            text="Detected formats (pick one or use quick quality)",
            padding=12,
            style="Card.TLabelframe",
        )
        table_box.pack(fill=tk.BOTH, expand=True, pady=12)
        self.table = ttk.Treeview(
            table_box,
            columns=("id", "res", "fps", "vcodec", "acodec", "ext", "size", "note"),
            show="headings",
            selectmode="browse",
        )
        for col, width in [
            ("id", 80),
            ("res", 110),
            ("fps", 60),
            ("vcodec", 120),
            ("acodec", 120),
            ("ext", 60),
            ("size", 100),
            ("note", 200),
        ]:
            self.table.heading(col, text=col)
            self.table.column(col, width=width, anchor=tk.CENTER if col in {"fps", "ext"} else tk.W)
        self.table.pack(fill=tk.BOTH, expand=True)

        actions = ttk.Frame(main, style="Main.TFrame")
        actions.pack(fill=tk.X, pady=(4, 12))
        ttk.Button(actions, text="Download", style="Accent.TButton", command=self.start_download).pack(side=tk.LEFT)
        ttk.Button(actions, text="Clear list", command=self.clear_formats).pack(side=tk.LEFT, padx=6)
        ttk.Label(actions, text=LEGAL_NOTE, style="Main.TLabel", wraplength=520, justify=tk.LEFT).pack(side=tk.RIGHT)

        log_box = ttk.LabelFrame(main, text="Log", padding=10, style="Card.TLabelframe")
        log_box.pack(fill=tk.BOTH, expand=True)
        self.log_text = ScrolledText(
            log_box,
            height=10,
            wrap=tk.WORD,
            state=tk.DISABLED,
            background="#121318",
            foreground="#e9e9ea",
        )
        self.log_text.pack(fill=tk.BOTH, expand=True)

    # ---------------------------- Actions ----------------------------
    def choose_output_dir(self) -> None:
        folder = filedialog.askdirectory(title="Select destination folder")
        if folder:
            self.output_dir_var.set(folder)

    def set_status(self, message: str) -> None:
        self.status_var.set(message)

    def log(self, message: str) -> None:
        def _append() -> None:
            self.log_text.configure(state=tk.NORMAL)
            self.log_text.insert(tk.END, message + "\n")
            self.log_text.see(tk.END)
            self.log_text.configure(state=tk.DISABLED)

        self.after(0, _append)

    def clear_formats(self) -> None:
        for item in self.table.get_children():
            self.table.delete(item)
        self.formats_cache.clear()
        self.title_var.set("No video loaded")

    def fetch_info(self) -> None:
        url = self.url_var.get().strip()
        if not url:
            messagebox.showwarning("SnapStream", "Enter the URL first.")
            return
        self.set_status("Querying yt-dlp...")
        self.log(f"Fetching formats for: {url}")
        threading.Thread(target=self._fetch_info_worker, args=(url,), daemon=True).start()

    def _fetch_info_worker(self, url: str) -> None:
        opts = {"quiet": True, "no_warnings": True, "skip_download": True}
        try:
            with YoutubeDL(opts) as ydl:
                info = ydl.extract_info(url, download=False)
            self.after(0, lambda: self._update_formats(info))
            self.log("Ready. Pick a format or use quick quality.")
            self.set_status("Metadata loaded")
        except Exception as exc:
            self.log(f"Error while fetching formats: {exc}")
            self.set_status("Error")
            self.after(0, lambda: messagebox.showerror("SnapStream", f"Could not read the URL:\n{exc}"))

    def _update_formats(self, info: Dict) -> None:
        self.clear_formats()
        title = info.get("title") or "No title"
        duration = _human_duration(info.get("duration"))
        uploader = info.get("uploader") or "-"
        self.title_var.set(f"{title} · {duration} · {uploader}")

        formats = info.get("formats", []) or []
        self.formats_cache = []
        for fmt in formats:
            fmt_id = str(fmt.get("format_id", ""))
            res = fmt.get("resolution") or "-"
            height = fmt.get("height")
            width = fmt.get("width")
            if width and height:
                res = f"{width}x{height}"
            fps = fmt.get("fps") or "-"
            vcodec = fmt.get("vcodec") or "-"
            acodec = fmt.get("acodec") or "-"
            ext = fmt.get("ext") or "-"
            size = _human_size(fmt.get("filesize") or fmt.get("filesize_approx"))
            note = fmt.get("format_note") or "-"
            has_audio = acodec not in ("none", "-", None)

            self.formats_cache.append(
                {
                    "id": fmt_id,
                    "has_audio": has_audio,
                }
            )
            self.table.insert("", tk.END, values=(fmt_id, res, fps, vcodec, acodec, ext, size, note))

    def _build_format_string(self, selected_id: Optional[str], audio_mode: bool, quality: str) -> str:
        if audio_mode:
            return "bestaudio/best"

        if selected_id:
            has_audio = False
            for item in self.formats_cache:
                if item["id"] == selected_id:
                    has_audio = item["has_audio"]
                    break
            if has_audio:
                return selected_id
            return f"{selected_id}+bestaudio/best"

        quality_map = {
            "1080p": "bestvideo[height<=1080][ext=mp4]+bestaudio[ext=m4a]/best[height<=1080][ext=mp4]/best",
            "720p": "bestvideo[height<=720][ext=mp4]+bestaudio[ext=m4a]/best[height<=720][ext=mp4]/best",
            "480p": "bestvideo[height<=480][ext=mp4]+bestaudio[ext=m4a]/best[height<=480][ext=mp4]/best",
            "360p": "bestvideo[height<=360][ext=mp4]+bestaudio[ext=m4a]/best[height<=360][ext=mp4]/best",
        }
        return quality_map.get(quality, "bestvideo[ext=mp4]+bestaudio[ext=m4a]/best")

    def start_download(self) -> None:
        if self.downloading:
            messagebox.showinfo("SnapStream", "A download is already running.")
            return
        url = self.url_var.get().strip()
        if not url:
            messagebox.showwarning("SnapStream", "Enter the URL first.")
            return
        out_dir = Path(self.output_dir_var.get().strip() or ".")
        try:
            out_dir.mkdir(parents=True, exist_ok=True)
        except Exception as exc:
            messagebox.showerror("SnapStream", f"Cannot use that folder:\n{exc}")
            return

        audio_mode = self.mode_var.get() == "Audio only MP3"
        quality = self.quality_var.get()
        selected = self.table.selection()
        selected_id = None
        if selected:
            values = self.table.item(selected[0]).get("values", [])
            if values:
                selected_id = str(values[0])

        fmt = self._build_format_string(selected_id, audio_mode, quality)

        opts: Dict = {
            "outtmpl": str(out_dir / "%(title)s.%(ext)s"),
            "quiet": True,
            "no_warnings": True,
            "noplaylist": True,
            "restrictfilenames": False,
            "concurrent_fragment_downloads": 4,
            "format": fmt,
            "progress_hooks": [self._progress_hook],
        }

        if audio_mode:
            opts["postprocessors"] = [
                {
                    "key": "FFmpegExtractAudio",
                    "preferredcodec": "mp3",
                    "preferredquality": self.audio_quality_var.get(),
                }
            ]

        self.downloading = True
        self.set_status("Downloading...")
        self.log(f"Download started -> {url}")
        threading.Thread(target=self._download_worker, args=(opts, url), daemon=True).start()

    def _progress_hook(self, data: Dict) -> None:
        if data.get("status") == "downloading":
            downloaded = data.get("downloaded_bytes") or 0
            total = data.get("total_bytes") or data.get("total_bytes_estimate")
            percent = 0.0
            if total:
                percent = downloaded / total * 100
            speed = data.get("speed")
            eta = data.get("eta")
            parts = [f"{percent:5.1f}%"]
            if speed:
                parts.append(f"{_human_size(int(speed))}/s")
            if eta:
                parts.append(f"ETA {eta}s")
            self.set_status(" ".join(parts))
        elif data.get("status") == "finished":
            self.set_status("Processing with FFmpeg...")

    def _download_worker(self, opts: Dict, url: str) -> None:
        try:
            with YoutubeDL(opts) as ydl:
                ydl.download([url])
            self.log("Download completed.")
            self.set_status("Ready")
        except Exception as exc:
            self.log(f"Download error: {exc}")
            self.set_status("Error")
            self.after(0, lambda: messagebox.showerror("SnapStream", f"Could not download:\n{exc}"))
        finally:
            self.downloading = False


def main() -> None:
    app = SnapStreamGUI()
    app.mainloop()


if __name__ == "__main__":
    main()