mediafish.py

"""Mediafish: an interactive multi-site media downloader powered by yt-dlp."""

from __future__ import annotations

import argparse
import os
import shutil
import sys
import textwrap
import urllib.request
import zipfile
from pathlib import Path
from urllib.parse import urlparse

AUDIO_FORMATS = {"aac", "alac", "flac", "m4a", "mp3", "opus", "vorbis", "wav"}
VIDEO_FORMATS = {"avi", "flv", "gif", "mkv", "mov", "mp4", "webm"}
FORMATS = ("best", *sorted(VIDEO_FORMATS), *sorted(AUDIO_FORMATS))
TOOLS = Path(__file__).resolve().parent / ".tools"
COLOR = sys.stdout.isatty() and "NO_COLOR" not in os.environ
try:
    "┌─┐│└┘›".encode(sys.stdout.encoding or "utf-8")
    UNICODE = True
except UnicodeEncodeError:
    UNICODE = False

BANNER = (
    "                O  o",
    r"          _\_    o",
    r">('>   \\/  o\ .",
    r"       //\___=",
    "          ''",
)


class QuietLogger:
    debug = warning = error = lambda *args: None


def style(text: str, code: str) -> str:
    return f"\033[{code}m{text}\033[0m" if COLOR else text


def prepare_screen(content_height: int) -> None:
    if not sys.stdout.isatty():
        return
    rows = shutil.get_terminal_size((80, 24)).lines
    print("\033[2J\033[H" + "\n" * max((rows - content_height) // 2, 0), end="")


def header() -> None:
    width = shutil.get_terminal_size((80, 24)).columns
    margin = " " * max((width - max(map(len, BANNER))) // 2, 0)
    print()
    for line in BANNER:
        print(style(margin + line, "1;97"))
    print()
    print(style("mediafish · cast a link. land the media.".center(width), "1"))
    sites = "youtube · x · threads · tiktok · instagram · soundcloud · +1000 more"
    print(style(sites.center(width), "2"))
    print()


def panel(title: str, body: str) -> int:
    lines = body.splitlines() or [""]
    width = max(len(title) + 4, *(len(line) for line in lines))
    margin = max((shutil.get_terminal_size((80, 24)).columns - width - 2) // 2, 0)
    pad = " " * margin
    top_left, horizontal, top_right, vertical, bottom_left, bottom_right = ("┌", "─", "┐", "│", "└", "┘") if UNICODE else ("+", "-", "+", "|", "+", "+")
    print(pad + style(f"{top_left}{horizontal} ", "2") + style(title, "1") + style(" " + horizontal * (width - len(title) - 2) + top_right, "2"))
    for line in lines:
        print(f"{pad}{vertical} {line:<{width}} {vertical}")
    print(pad + style(bottom_left + horizontal * (width + 2) + bottom_right, "2"))
    return margin


def ensure_binary(name: str, url: str) -> bool:
    if shutil.which(name):
        return True
    executable = next(TOOLS.rglob(f"{name}.exe"), None) if TOOLS.exists() else None
    if executable:
        os.environ["PATH"] = str(executable.parent) + os.pathsep + os.environ.get("PATH", "")
        return True

    print(style(f"  First run: installing {name}...", "33"))
    archive = TOOLS / f"{name}.zip"
    try:
        TOOLS.mkdir(parents=True, exist_ok=True)
        urllib.request.urlretrieve(url, archive)
        with zipfile.ZipFile(archive) as bundle:
            if any(Path(member).is_absolute() or ".." in Path(member).parts for member in bundle.namelist()):
                raise ValueError("unsafe archive")
            bundle.extractall(TOOLS / name)
        archive.unlink()
    except (OSError, ValueError, zipfile.BadZipFile) as error:
        print(style(f"  Could not install {name}: {error}", "31"), file=sys.stderr)
        return False

    executable = next((TOOLS / name).rglob(f"{name}.exe"), None)
    if not executable:
        print(style(f"  Could not find {name}.exe after installation.", "31"), file=sys.stderr)
        return False
    os.environ["PATH"] = str(executable.parent) + os.pathsep + os.environ.get("PATH", "")
    return True


def bootstrap_media_tools() -> None:
    if sys.platform != "win32":
        return
    ensure_binary("ffmpeg", "https://www.gyan.dev/ffmpeg/builds/ffmpeg-release-essentials.zip")
    if not any(shutil.which(runtime) for runtime in ("deno", "node", "bun")):
        ensure_binary("deno", "https://github.com/denoland/deno/releases/latest/download/deno-x86_64-pc-windows-msvc.zip")


def parse_formats(values: list[str]) -> list[str]:
    formats = list(dict.fromkeys(part.strip().lower() for value in values for part in value.split(",") if part.strip()))
    invalid = set(formats) - set(FORMATS)
    if not formats or invalid:
        raise ValueError(f"Choose one or more of: {', '.join(FORMATS)}")
    return formats


def valid_url(value: str) -> str:
    parsed = urlparse(value)
    if parsed.scheme not in {"http", "https"} or not parsed.netloc:
        raise argparse.ArgumentTypeError(f"Invalid URL: {value}")
    return value


def options_for(media_format: str, output: Path, cookies_browser: str | None) -> dict:
    options: dict = {
        "outtmpl": str(output / "%(title).180B [%(id)s].%(ext)s"),
        "noplaylist": False,
        "windowsfilenames": sys.platform == "win32",
    }
    if cookies_browser:
        options["cookiesfrombrowser"] = (cookies_browser,)
    if media_format == "best":
        options["format"] = "best"
    elif media_format in AUDIO_FORMATS:
        options.update(
            format="bestaudio/best",
            postprocessors=[{"key": "FFmpegExtractAudio", "preferredcodec": media_format}],
        )
    else:
        options.update(
            format="bestvideo*+bestaudio/best",
            postprocessors=[{"key": "FFmpegVideoConvertor", "preferedformat": media_format}],
        )
    return options


def fetch_metadata(urls: list[str], cookies_browser: str | None = None) -> list[dict]:
    import yt_dlp

    options = {"quiet": True, "no_warnings": True, "skip_download": True, "extract_flat": True}
    if cookies_browser:
        options["cookiesfrombrowser"] = (cookies_browser,)
    results = []
    with yt_dlp.YoutubeDL(options) as client:
        for url in urls:
            try:
                info = client.extract_info(url, download=False) or {}
                info["_requested_url"] = url
                results.append(info)
            except yt_dlp.utils.DownloadError as error:
                results.append({"_requested_url": url, "_error": str(error)})
    return results


def metadata_summary(items: list[dict], formats: list[str], output: Path) -> str:
    def value(item: object) -> str:
        return textwrap.shorten(str(item), width=60, placeholder="…")

    def duration(seconds: object) -> str:
        total = int(seconds or 0)
        hours, remainder = divmod(total, 3600)
        minutes, seconds = divmod(remainder, 60)
        return f"{hours}:{minutes:02}:{seconds:02}" if hours else f"{minutes}:{seconds:02}"

    blocks = []
    for number, info in enumerate(items, 1):
        fields = [("URL", info["_requested_url"])]
        if info.get("_error"):
            fields.append(("Metadata", f"Unavailable — {info['_error']}"))
        else:
            upload_date = info.get("upload_date")
            entries = info.get("entries")
            fields = [
                ("Name", info.get("title") or info.get("fulltitle")),
                ("Source", info.get("extractor_key") or info.get("extractor")),
                ("Media ID", info.get("id")),
                ("Artist", info.get("artist")),
                ("Creator", info.get("uploader") or info.get("channel") or info.get("creator")),
                ("Duration", duration(info.get("duration")) if info.get("duration") is not None else None),
                ("Published", f"{upload_date[:4]}-{upload_date[4:6]}-{upload_date[6:]}" if upload_date and len(upload_date) == 8 else upload_date),
                ("Views", f"{info['view_count']:,}" if info.get("view_count") is not None else None),
                ("Playlist", f"{len(entries)} items" if entries is not None else None),
                ("Resolution", info.get("resolution")),
                ("Container", info.get("ext")),
                ("URL", info.get("webpage_url") or info["_requested_url"]),
            ]
        lines = [f"Source {number}"] + [f"{label:<11}{value(field)}" for label, field in fields if field not in {None, ""}]
        blocks.append("\n".join(lines))
    blocks.append(f"Output\n{'Formats':<11}{', '.join(formats)}\n{'Folder':<11}{value(output.expanduser())}")
    return "\n\n".join(blocks)


def progress_text(media_format: str, progress: dict) -> str:
    total = progress.get("total_bytes") or progress.get("total_bytes_estimate") or 0
    downloaded = progress.get("downloaded_bytes") or 0
    ratio = min(downloaded / total, 1) if total else 0
    filled = round(ratio * 16)
    bar = ("█" if UNICODE else "#") * filled + ("·" if UNICODE else "-") * (16 - filled)
    speed = progress.get("speed") or 0
    for unit in ("B", "KiB", "MiB", "GiB"):
        if speed < 1024 or unit == "GiB":
            speed_text = f"{speed:.1f} {unit}/s"
            break
        speed /= 1024
    eta = progress.get("eta")
    eta_text = f"{int(eta) // 60}:{int(eta) % 60:02}" if eta is not None else "--:--"
    name = textwrap.shorten(str(progress.get("info_dict", {}).get("title") or "media"), width=22, placeholder="…")
    return f"{media_format}  {name}  [{bar}] {ratio:>4.0%}  {speed_text}  ETA {eta_text}"


def show_status(message: str, code: str = "") -> None:
    width = shutil.get_terminal_size((80, 24)).columns
    print(style(message.center(width), code) if code else message.center(width))


def interactive_args(parser: argparse.ArgumentParser) -> argparse.Namespace:
    width = shutil.get_terminal_size((80, 24)).columns
    prepare_screen(len(BANNER) + 9)
    header()
    margin = panel("Paste a link", "One link per line. Press Enter on an empty line when done.")
    urls = []
    while True:
        value = input(" " * (margin + 2) + style(f"{'›' if UNICODE else '>'} ", "1;97")).strip()
        if not value:
            break
        try:
            urls.append(valid_url(value))
        except argparse.ArgumentTypeError as error:
            print(" " * (margin + 2) + style(str(error), "31"))
    if not urls:
        parser.error("Enter at least one URL")

    menu = "\n".join(f"{number:>2}. {name}" for number, name in enumerate(FORMATS, 1))
    margin = panel("Output formats", menu + "\n\nChoose one or more numbers, separated by commas. [1]")
    selection = input(" " * (margin + 2) + style(f"{'›' if UNICODE else '>'} ", "1;97")).strip() or "1"
    try:
        formats = [FORMATS[int(item.strip()) - 1] for item in selection.split(",")]
        if not all(1 <= int(item.strip()) <= len(FORMATS) for item in selection.split(",")):
            raise ValueError
    except (ValueError, IndexError):
        parser.error(f"Choose format numbers from 1 to {len(FORMATS)}")

    arrow = "›" if UNICODE else ">"
    margin = panel("Download folder", "Where should the files be saved? [downloads]")
    output = input(" " * (margin + 2) + style(f"{arrow} ", "1;97")).strip() or "downloads"
    print(" " * (margin + 2) + style("Reading source metadata…", "2"))
    summary = metadata_summary(fetch_metadata(urls), parse_formats(formats), Path(output))
    prepare_screen(len(summary.splitlines()) + 4)
    margin = panel("Review download", summary)
    if input(" " * (margin + 2) + style("Start download? ", "1") + f"[Y/n] {arrow} ").strip().lower() not in {"", "y", "yes"}:
        print(" " * (margin + 2) + style("Cancelled.", "33"))
        raise SystemExit(0)
    return argparse.Namespace(urls=urls, formats=formats, output=Path(output), cookies_from_browser=None)


def download(urls: list[str], formats: list[str], output: Path, cookies_browser: str | None = None) -> int:
    if any(item != "best" for item in formats) and not shutil.which("ffmpeg"):
        print("Error: FFmpeg is required for conversion. Install it and ensure ffmpeg is on PATH.", file=sys.stderr)
        return 2

    try:
        import yt_dlp
    except ImportError:
        print("Error: yt-dlp is not installed. Run: python -m pip install -e .", file=sys.stderr)
        return 2

    output.mkdir(parents=True, exist_ok=True)
    failed = False
    prepare_screen(len(formats) + 4)
    show_status("Download progress", "1")
    print()
    for number, media_format in enumerate(formats, 1):
        width = shutil.get_terminal_size((80, 24)).columns

        def progress_hook(progress: dict) -> None:
            if not sys.stdout.isatty():
                return
            if progress.get("status") == "downloading":
                line = progress_text(media_format, progress)
            elif progress.get("status") == "finished":
                line = f"{media_format}  download complete · processing…"
            else:
                return
            print("\r" + line.center(width)[:width], end="", flush=True)

        show_status(f"{number}/{len(formats)}  Preparing {media_format}…", "2")
        try:
            options = options_for(media_format, output, cookies_browser)
            options.update(quiet=True, no_warnings=True, noprogress=True, logger=QuietLogger(), progress_hooks=[progress_hook])
            with yt_dlp.YoutubeDL(options) as client:
                result = client.download(urls)
            if sys.stdout.isatty():
                print("\r" + " " * width + "\r", end="")
            if result:
                failed = True
                show_status(f"{'✗' if UNICODE else 'X'}  {media_format} failed", "31")
            else:
                show_status(f"{'✓' if UNICODE else 'OK'}  {media_format} complete", "32")
        except yt_dlp.utils.DownloadError as error:
            if sys.stdout.isatty():
                print("\r" + " " * width + "\r", end="")
            show_status(f"{'✗' if UNICODE else 'X'}  {media_format} failed", "31")
            show_status(textwrap.shorten(str(error).removeprefix("ERROR: "), width=68, placeholder="…"), "31")
            failed = True
    return int(failed)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Download and convert media from sites supported by yt-dlp.")
    parser.add_argument("urls", nargs="*", type=valid_url, help="one or more media URLs")
    parser.add_argument("-f", "--format", dest="formats", action="append", default=[], help="output format; repeat or comma-separate")
    parser.add_argument("-o", "--output", type=Path, default=Path("downloads"), help="download folder")
    parser.add_argument("--cookies-from-browser", help="browser name for authenticated/private media")
    return parser


def main(argv: list[str] | None = None) -> int:
    try:
        parser = build_parser()
        args = parser.parse_args(argv)
        bootstrap_media_tools()
        if not args.urls:
            args = interactive_args(parser)
        formats = parse_formats(args.formats or ["best"])
        return download(args.urls, formats, args.output.expanduser(), args.cookies_from_browser)
    except KeyboardInterrupt:
        print(style("\n  Stopped safely. Partial downloads can resume next time.", "33"))
        return 130
    except ValueError as error:
        parser.error(str(error))


if __name__ == "__main__":
    raise SystemExit(main())