splinter-keep/tools/music-fetch.py

#!/usr/bin/env python3
"""
tools/music-fetch.py — Search YouTube and download ambience music.
Wires downloaded tracks into session/ambience_options.md.

Requires:
    pip install yt-dlp
    ffmpeg (system install, for audio extraction)

Usage:
    python3 tools/music-fetch.py tavern
    python3 tools/music-fetch.py "deep fen" "swamp ambience D&D"
    python3 tools/music-fetch.py tavern --url "https://youtu.be/..."
    python3 tools/music-fetch.py tavern --replace
    python3 tools/music-fetch.py tavern --dry-run
"""

import argparse
import random
import re
import string
import sys
from pathlib import Path


# ── Paths ────────────────────────────────────────────────
BASE = Path(__file__).resolve().parent.parent
SESSION = BASE / 'session'
AUDIO_DIR = SESSION / 'audio'
OPTIONS_PATH = SESSION / 'ambience_options.md'
SOURCES_PATH = SESSION / 'ambience_sources.md'

# ── Default search queries per ambience ───────────────────
DEFAULT_QUERIES = {
    "tavern": "D&D fantasy tavern inn ambience background music",
    "forest": "D&D forest enchanted woodland ambience background music",
    "combat": "D&D combat battle epic fantasy music",
    "dungeon": "D&D dungeon cave ambience background music",
    "tension": "D&D tension suspense dark ambient music",
    "calm": "D&D calm peaceful relaxing ambience",
    "town": "D&D town village marketplace ambience",
    "wilds": "D&D wilderness nature ambient music",
}


# ── Helpers ──────────────────────────────────────────────
def _import_ytdlp():
    try:
        from yt_dlp import YoutubeDL
        return YoutubeDL
    except ImportError:
        print("Error: yt-dlp not found. Install it with: pip install yt-dlp")
        sys.exit(1)


def _check_ffmpeg():
    import shutil
    if not shutil.which("ffmpeg"):
        print("Error: ffmpeg not found. Install it (apt install ffmpeg, brew install ffmpeg, etc.)")
        sys.exit(1)


def sanitize_filename(name):
    safe = "-_.() " + string.ascii_letters + string.digits
    return "".join(c if c in safe else "_" for c in name).strip()


def make_stem(ambience, existing_files):
    """Generate next sequential stem, e.g. tavern_03"""
    nums = []
    for f in existing_files:
        m = re.search(rf"^{re.escape(ambience)}_(\d+)", f.stem)
        if m:
            nums.append(int(m.group(1)))
    idx = max(nums) + 1 if nums else 1
    return f"{ambience}_{idx:02d}"


# ── YouTube search / info ─────────────────────────────────
def search_videos(query, max_results=20):
    YoutubeDL = _import_ytdlp()
    ydl = YoutubeDL({"quiet": True, "no_warnings": True})
    info = ydl.extract_info(f"ytsearch{max_results}:{query}", download=False)
    return info.get("entries", [])


def get_video_info(url):
    YoutubeDL = _import_ytdlp()
    ydl = YoutubeDL({"quiet": True, "no_warnings": True})
    return ydl.extract_info(url, download=False)


# ── Download ──────────────────────────────────────────────
def download_audio(url, outtmpl):
    YoutubeDL = _import_ytdlp()
    ydl_opts = {
        "quiet": True,
        "no_warnings": True,
        "outtmpl": outtmpl,
        "format": "bestaudio/best",
        "postprocessors": [{
            "key": "FFmpegExtractAudio",
            "preferredcodec": "vorbis",
            "preferredquality": "192",
        }],
    }
    with YoutubeDL(ydl_opts) as ydl:
        ydl.download([url])


# ── ambience_options.md table helpers ─────────────────────
def _find_table_bounds(lines):
    """Return (start, end) of the first contiguous |...| table block."""
    start = end = None
    for i, line in enumerate(lines):
        s = line.strip()
        is_row = s.startswith("|") and s.endswith("|")
        if is_row and start is None:
            start = i
        if is_row:
            end = i + 1
        elif start is not None:
            break
    return start, end if end else start


def parse_table():
    """Parse ambience_options.md table into {name: [filename, ...]}"""
    if not OPTIONS_PATH.exists():
        return {}
    lines = OPTIONS_PATH.read_text().splitlines()
    start, end = _find_table_bounds(lines)
    if start is None:
        return {}
    entries = {}
    for line in lines[start:end]:
        s = line.strip()
        parts = [p.strip() for p in s.split("|")]
        parts = [p for p in parts if p]
        if len(parts) < 2:
            continue
        if all(c in "-:| " for c in s):
            continue
        if parts[0].lower() == "ambience":
            continue
        name = parts[0].lower()
        raw = parts[1]
        if raw == "(stops all music)" or raw == "(no files)":
            entries[name] = []
        else:
            entries[name] = [AUDIO_DIR / f.strip() for f in raw.split(",") if f.strip()]
    return entries


def rebuild_table(entries):
    """Render entries as a markdown table."""
    table = [
        "| Ambience | Files |",
        "|----------|-------|",
    ]
    if "silence" in entries or not entries:
        table.append("| silence | (stops all music) |")
    for name in sorted(entries.keys()):
        if name == "silence":
            continue
        files = entries[name]
        if files:
            table.append(f"| {name} | {', '.join(f.name for f in files)} |")
        else:
            table.append(f"| {name} | (no files) |")
    return "\n".join(table)


def update_options_file(entries):
    """Replace the table in ambience_options.md while preserving surrounding text."""
    if not OPTIONS_PATH.exists():
        OPTIONS_PATH.write_text("# Ambience Options\n\n" + rebuild_table(entries) + "\n")
        return
    lines = OPTIONS_PATH.read_text().splitlines()
    start, end = _find_table_bounds(lines)
    new_table = rebuild_table(entries)
    if start is None:
        # No existing table — append after a blank line
        body = "\n".join(lines) + "\n\n" + new_table + "\n"
    else:
        body = "\n".join(lines[:start]) + "\n" + new_table + "\n" + "\n".join(lines[end:])
        body = body.strip() + "\n"
    OPTIONS_PATH.write_text(body)


# ── ambience_sources.md helpers ──────────────────────────
def _find_sources_table_bounds(lines):
    """Return (start, end) of the sources table block."""
    start = end = None
    for i, line in enumerate(lines):
        s = line.strip()
        is_row = s.startswith("|") and s.endswith("|")
        if is_row and start is None:
            start = i
        if is_row:
            end = i + 1
        elif start is not None:
            break
    return start, end if end else start


def parse_sources():
    """Parse ambience_sources.md into {filename: (ambience, url)}"""
    if not SOURCES_PATH.exists():
        return {}
    lines = SOURCES_PATH.read_text().splitlines()
    start, end = _find_sources_table_bounds(lines)
    if start is None:
        return {}
    sources = {}
    for line in lines[start:end]:
        s = line.strip()
        parts = [p.strip() for p in s.split("|")]
        parts = [p for p in parts if p]
        if len(parts) < 3:
            continue
        if all(c in "-:| " for c in s):
            continue
        if parts[0].lower() == "file":
            continue
        filename, ambience, url = parts[0], parts[1], parts[2]
        sources[filename] = (ambience, url)
    return sources


def rebuild_sources_table(sources):
    """Render sources dict as a markdown table."""
    table = [
        "| File | Ambience | Source |",
        "|------|----------|--------|",
    ]
    for filename in sorted(sources.keys()):
        ambience, url = sources[filename]
        table.append(f"| {filename} | {ambience} | {url} |")
    return "\n".join(table)


def update_sources_file(sources):
    """Replace the table in ambience_sources.md while preserving surrounding text."""
    if not SOURCES_PATH.exists():
        SOURCES_PATH.write_text("# Ambience Sources\n\n" + rebuild_sources_table(sources) + "\n")
        return
    lines = SOURCES_PATH.read_text().splitlines()
    start, end = _find_sources_table_bounds(lines)
    new_table = rebuild_sources_table(sources)
    if start is None:
        body = "\n".join(lines) + "\n\n" + new_table + "\n"
    else:
        body = "\n".join(lines[:start]) + "\n" + new_table + "\n" + "\n".join(lines[end:])
        body = body.strip() + "\n"
    SOURCES_PATH.write_text(body)


# ── Main ──────────────────────────────────────────────────
def main():
    parser = argparse.ArgumentParser(
        description="Fetch ambience music from YouTube and wire into the Chaos TUI.",
    )
    parser.add_argument("ambience", help="Ambience name (e.g. tavern, forest, combat)")
    parser.add_argument("query", nargs="?", default=None,
                        help="Search query (auto-generated from ambience name if omitted)")
    parser.add_argument("--url", default=None,
                        help="Specific YouTube URL (overrides search)")
    parser.add_argument("--replace", action="store_true",
                        help="Replace existing tracks for this ambience instead of appending")
    parser.add_argument("--dry-run", action="store_true",
                        help="Show what would be fetched without downloading")
    parser.add_argument("--max-results", type=int, default=20,
                        help="Number of search results to consider (default: 20)")
    args = parser.parse_args()

    name = args.ambience.strip().lower()
    if not name:
        print("Error: ambience name is required.")
        sys.exit(1)

    _check_ffmpeg()

    # Resolve query
    query = args.query
    if not query:
        if name in DEFAULT_QUERIES:
            query = DEFAULT_QUERIES[name]
        else:
            query = f"D&D {name} ambience background music"
    print(f"  query: {query}")

    # Parse existing table
    entries = parse_table()
    existing = entries.get(name, [])

    if args.replace:
        existing = []

    # Get candidate video(s)
    if args.url:
        YoutubeDL = _import_ytdlp()
        ydl = YoutubeDL({"quiet": True, "no_warnings": True})
        video = ydl.extract_info(args.url, download=False)
        candidates = [video]
        source_desc = args.url
    else:
        candidates = search_videos(query, args.max_results)
        if not candidates:
            print("No results found.")
            sys.exit(1)
        source_desc = f"top {len(candidates)} results for \"{query}\""
        # Pick randomly from the pool
        video = random.choice(candidates)
        candidates = [video]

    video = candidates[0]
    title = video.get("title", "unknown")
    url = video.get("webpage_url", video.get("url", args.url or "?"))
    duration = video.get("duration", 0)
    mins, secs = divmod(duration, 60) if duration else (0, 0)

    print(f"  picked: \"{title}\"")
    print(f"  url:    {url}")
    print(f"  length: {mins}:{secs:02d}")

    if args.dry_run:
        stem = make_stem(name, existing)
        print(f"  target: {AUDIO_DIR / (stem + '.ogg')}")
        print("  (dry-run, nothing saved)")
        return

    # Download
    stem = make_stem(name, existing)
    outtmpl = str(AUDIO_DIR / stem)  # no extension — yt-dlp handles it

    print(f"  downloading...", end=" ", flush=True)
    download_audio(url, outtmpl)
    target = AUDIO_DIR / f"{stem}.ogg"
    print("done.")

    if not target.exists():
        print(f"Warning: expected file not found at {target}")
        print(f"Check session/audio/ for the actual output.")
        sys.exit(1)

    # Wire into options config
    existing.append(target)
    entries[name] = existing
    update_options_file(entries)

    # Record source
    sources = parse_sources()
    if args.replace:
        # Remove stale entries for this ambience
        sources = {k: v for k, v in sources.items() if v[0] != name}
    sources[target.name] = (name, url)
    update_sources_file(sources)

    total = len(existing)
    print(f"  saved:  {target}")
    print(f"  wired:  {name} now has {total} track{'s' if total != 1 else ''}")
    print(f"  source: recorded in ambience_sources.md")


if __name__ == "__main__":
    main()