cli.py

#!/usr/bin/env python3
"""Apple Music playlist CLI.

Usage:
    python cli.py match --songs '[{"artist": "Daft Punk", "title": "One More Time"}]'
    python cli.py search --query "Miles Davis Blue in Green"
    python cli.py create --name "Chill Jazz" --track-ids '["123", "456"]'
    python cli.py create --name "Chill Jazz" --track-ids '["123"]' --upsert
    python cli.py list
    python cli.py tracks --id p.ABC123
    python cli.py add --id p.ABC123 --track-ids '["123", "456"]'
    python cli.py remove --id p.ABC123 --track-ids '["i.XXX", "i.YYY"]'
    python cli.py rename --id p.ABC123 --name "New Name" --description "New desc"
    python cli.py reorder --id p.ABC123 --track-ids '["123", "456", "789"]'
    python cli.py yt-search --query "persona 5 cafe music"
    python cli.py yt-download --url "https://youtube.com/watch?v=..." --name "Cafe Leblanc"
"""

import argparse
import json
import hashlib
import re
import subprocess
import sys
import os
from pathlib import Path

import requests
from dotenv import load_dotenv

load_dotenv()

from apple_music import AppleMusicClient, TokenExpiredError, fetch_dev_token

CACHE_DIR = Path(__file__).parent / ".search_cache"
DOWNLOAD_DIR = Path.home() / "Music" / "YouTube Downloads"


def _cache_key(query: str, limit: int) -> str:
    return hashlib.sha256(f"{query}|{limit}".encode()).hexdigest()[:16]


def cached_search(am: AppleMusicClient, query: str, limit: int = 10) -> list[dict]:
    """Search with file-based cache. Avoids repeat API calls for the same query."""
    CACHE_DIR.mkdir(exist_ok=True)
    key = _cache_key(query, limit)
    cache_file = CACHE_DIR / f"{key}.json"
    if cache_file.exists():
        return json.loads(cache_file.read_text())
    results = am.search_song(query, limit=limit)
    cache_file.write_text(json.dumps(results))
    return results


def get_client() -> AppleMusicClient:
    user_token = os.environ.get("APPLE_USER_TOKEN", "").strip()
    if not user_token:
        print("APPLE_USER_TOKEN not set in .env")
        sys.exit(1)
    dev_token = os.environ.get("APPLE_DEV_TOKEN", "").strip()
    if not dev_token:
        dev_token = fetch_dev_token()
    return AppleMusicClient(
        dev_token=dev_token,
        user_token=user_token,
        storefront=os.environ.get("APPLE_STOREFRONT", "za"),
    )


# --- Scoring logic ---

REMIX_PATTERNS = re.compile(
    r"\b(remix|remixed|rmx|mix(?:ed)?|rework|bootleg|sped up|slowed|edit(?:ed)?)\b"
    r"|"
    r"\[.*(?:remix|mix|rework|version).*\]",
    re.IGNORECASE,
)

DJ_MIX_PATTERNS = re.compile(
    r"\b(dj mix|mixed|ministry of sound|fabric \d|ibiza|essential mix)\b",
    re.IGNORECASE,
)

# Words that are part of the original title, not remix indicators
SAFE_TITLE_WORDS = re.compile(
    r"\b(radio edit|original|remaster(?:ed)?|deluxe|bonus track)\b",
    re.IGNORECASE,
)


def score_result(result: dict, requested_artist: str, requested_title: str) -> int:
    """Score a search result. Higher = better match. Range roughly 0-100."""
    score = 50  # baseline

    name = result.get("name", "")
    artist = result.get("artist", "")
    album = result.get("album", "")

    # --- Artist match (biggest signal) ---
    req_lower = requested_artist.lower()
    art_lower = artist.lower()
    # Split both into word sets for comparison
    req_words = set(req_lower.split())
    art_words = set(art_lower.split())
    # Filter out short words (a, &, the, ft, feat, vs) for matching
    noise = {"a", "an", "the", "&", "ft", "ft.", "feat", "feat.", "vs", "vs.", "and"}
    req_significant = req_words - noise
    art_significant = art_words - noise
    if req_significant and req_significant <= art_significant:
        # All significant requested words are in the result artist
        score += 25
    elif req_significant and len(req_significant & art_significant) / len(req_significant) >= 0.5:
        # At least half the words match
        score += 10
    else:
        score -= 30  # wrong artist entirely

    # --- Remix/mix penalty ---
    # Check if the name has remix indicators that aren't in the requested title
    name_without_safe = SAFE_TITLE_WORDS.sub("", name)
    if REMIX_PATTERNS.search(name_without_safe):
        req_title_lower = requested_title.lower()
        # Only penalize if the user didn't ask for a remix
        if not REMIX_PATTERNS.search(req_title_lower):
            score -= 25

    # --- DJ mix / compilation album penalty ---
    if DJ_MIX_PATTERNS.search(album):
        score -= 15

    # --- Audio quality boost ---
    traits = result.get("audio_traits", [])
    if "atmos" in traits or "spatial" in traits:
        score += 5
    if "hi-res-lossless" in traits:
        score += 3
    elif "lossless" in traits:
        score += 1

    return score


def match_songs(am: AppleMusicClient, songs: list[dict]) -> list[dict]:
    """Batch search and score. Returns list of {request, pick, score, warnings}."""
    results = []
    for song in songs:
        query = f"{song['artist']} {song['title']}"
        candidates = cached_search(am, query, limit=10)
        if not candidates:
            results.append({
                "request": song,
                "pick": None,
                "score": 0,
                "warnings": ["no results found"],
            })
            continue

        scored = []
        for c in candidates:
            s = score_result(c, song["artist"], song["title"])
            scored.append((s, c))
        scored.sort(key=lambda x: x[0], reverse=True)

        best_score, best = scored[0]
        warnings = []

        # Flag potential issues
        noise = {"a", "an", "the", "&", "ft", "ft.", "feat", "feat.", "vs", "vs.", "and"}
        req_words = set(song["artist"].lower().split()) - noise
        pick_words = set(best["artist"].lower().split()) - noise
        if req_words and not (req_words <= pick_words):
            warnings.append(f"artist mismatch: requested '{song['artist']}', got '{best['artist']}'")

        name_without_safe = SAFE_TITLE_WORDS.sub("", best.get("name", ""))
        if REMIX_PATTERNS.search(name_without_safe) and not REMIX_PATTERNS.search(song["title"]):
            warnings.append("remix/mix version")

        if DJ_MIX_PATTERNS.search(best.get("album", "")):
            warnings.append("from DJ mix/compilation")

        results.append({
            "request": song,
            "pick": best,
            "score": best_score,
            "warnings": warnings,
        })
    return results


# --- Commands ---

def cmd_match(args):
    am = get_client()
    songs = json.loads(args.songs)
    results = match_songs(am, songs)

    warn_count = 0
    miss_count = 0
    for r in results:
        req = r["request"]
        pick = r["pick"]
        if pick is None:
            print(f"  ✗ {req['artist']} - {req['title']}  (no results)")
            miss_count += 1
        elif r["warnings"]:
            flags = ", ".join(r["warnings"])
            print(f"  ⚠ {pick['id']}  {pick['artist']} - {pick['name']}  ({flags})")
            warn_count += 1
        else:
            print(f"  ✓ {pick['id']}  {pick['artist']} - {pick['name']}")

    total = len(results)
    ok = total - warn_count - miss_count
    print(f"\n{ok} matched, {warn_count} warnings, {miss_count} missed out of {total}")


def cmd_search(args):
    am = get_client()
    results = am.search_song(args.query, limit=args.limit)
    if not results:
        print("No results.")
        return
    for r in results:
        mins = r["duration_ms"] // 60000
        secs = (r["duration_ms"] % 60000) // 1000
        traits = ", ".join(r["audio_traits"])
        print(f"  {r['id']}  {r['artist']} - {r['name']} ({r['album']}, {r['release_date']}, {mins}:{secs:02d}, {traits})")


def cmd_create(args):
    am = get_client()
    track_ids = json.loads(args.track_ids)

    # Upsert: find existing playlist by name instead of creating a duplicate
    if args.upsert:
        playlists = am.list_playlists()
        for p in playlists:
            if p["name"] == args.name:
                print(f"Found existing playlist: {args.name} ({p['id']})")
                if args.description:
                    am.update_playlist(p["id"], description=args.description)
                if track_ids:
                    am.add_tracks(p["id"], track_ids)
                    print(f"Added {len(track_ids)} tracks.")
                return

    playlist_id = am.create_playlist(args.name, args.description or "")
    print(f"Created playlist: {args.name} ({playlist_id})")

    if track_ids:
        am.add_tracks(playlist_id, track_ids)
        print(f"Added {len(track_ids)} tracks.")


def cmd_list(args):
    am = get_client()
    playlists = am.list_playlists()
    for p in playlists:
        print(f"  {p['id']}  {p['name']}")


def cmd_tracks(args):
    am = get_client()
    tracks = am.get_playlist_tracks(args.id)
    for t in tracks:
        print(f"  {t['id']}  {t['artist']} - {t['name']} ({t['album']})")


def cmd_add(args):
    am = get_client()
    track_ids = json.loads(args.track_ids)
    am.add_tracks(args.id, track_ids)
    print(f"Added {len(track_ids)} tracks.")


def cmd_remove(args):
    am = get_client()
    track_ids = json.loads(args.track_ids)
    am.remove_tracks(args.id, track_ids)
    print(f"Removed {len(track_ids)} tracks.")


def cmd_rename(args):
    am = get_client()
    name = args.name if args.name else None
    desc = args.description if args.description else None
    if name is None and desc is None:
        print("Nothing to do. Use --name and/or --description.")
        sys.exit(1)
    am.update_playlist(args.id, name=name, description=desc)
    print("Updated.")


def cmd_reorder(args):
    am = get_client()
    track_ids = json.loads(args.track_ids)
    new_id = am.replace_all_tracks(args.id, track_ids)
    print(f"Reordered playlist. New ID: {new_id}")


def _sanitize_filename(name: str) -> str:
    """Remove characters that cause problems in filenames and AppleScript."""
    # Replace fullwidth and other problematic unicode with ASCII equivalents
    replacements = {
        "\uff5c": "-", "\uff1a": "-", "\uff0f": "-",  # fullwidth |, :, /
        "|": "-", ":": "-", "/": "-", "\\": "-",
        '"': "'", "?": "", "*": "", "<": "", ">": "",
    }
    for old, new in replacements.items():
        name = name.replace(old, new)
    return name.strip()


def cmd_yt_search(args):
    """Search YouTube and display results with clickable links."""
    limit = args.limit
    query = f"ytsearch{limit}:{args.query}"
    result = subprocess.run(
        ["yt-dlp", "--flat-playlist", "-J", query],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        print(f"yt-dlp error: {result.stderr}", file=sys.stderr)
        sys.exit(1)

    data = json.loads(result.stdout)
    entries = data.get("entries", [])
    if not entries:
        print("No results.")
        return

    for i, e in enumerate(entries, 1):
        dur = int(e.get("duration") or 0)
        hrs, remainder = divmod(dur, 3600)
        mins, secs = divmod(remainder, 60)
        dur_str = f"{hrs}:{mins:02d}:{secs:02d}" if hrs else f"{mins}:{secs:02d}"
        url = f"https://youtube.com/watch?v={e['id']}"
        print(f"  {i}. [{dur_str}] {e['title']}")
        print(f"     {url}")


def _tag_m4a(path: Path, title: str, artist: str | None, album: str | None, thumbnail_url: str | None):
    """Write metadata tags and artwork to an m4a file."""
    from mutagen.mp4 import MP4, MP4Cover

    audio = MP4(str(path))
    audio["\xa9nam"] = [title]
    if artist:
        audio["\xa9ART"] = [artist]
    if album:
        audio["\xa9alb"] = [album]

    # Embed thumbnail as cover art
    if thumbnail_url:
        try:
            resp = requests.get(thumbnail_url, timeout=15)
            resp.raise_for_status()
            fmt = MP4Cover.FORMAT_PNG if thumbnail_url.endswith(".png") else MP4Cover.FORMAT_JPEG
            audio["covr"] = [MP4Cover(resp.content, imageformat=fmt)]
        except Exception as e:
            print(f"  (artwork fetch failed: {e})")

    audio.save()


def cmd_yt_download(args):
    """Download audio from YouTube and import into Music.app."""
    DOWNLOAD_DIR.mkdir(parents=True, exist_ok=True)

    # Fetch video metadata (title, uploader, thumbnail)
    probe = subprocess.run(
        ["yt-dlp", "-J", "--no-download", args.url],
        capture_output=True, text=True,
    )
    if probe.returncode != 0:
        print(f"Failed to fetch video info: {probe.stderr}", file=sys.stderr)
        sys.exit(1)
    meta = json.loads(probe.stdout)

    title = args.name or meta.get("title", "Unknown")
    artist = args.artist or meta.get("uploader") or meta.get("channel")
    album = args.album or "YouTube"
    thumbnail_url = meta.get("thumbnail")

    filename = _sanitize_filename(title)
    output_path = DOWNLOAD_DIR / f"{filename}.m4a"

    # Download
    print(f"Downloading: {title}")
    dl = subprocess.run(
        [
            "yt-dlp", "-x", "--audio-format", "m4a", "--audio-quality", "0",
            "-o", str(output_path), args.url,
        ],
    )
    if dl.returncode != 0:
        print("Download failed.", file=sys.stderr)
        sys.exit(1)

    if not output_path.exists():
        print(f"Expected file not found: {output_path}", file=sys.stderr)
        sys.exit(1)

    # Tag metadata
    print(f"Tagging: {title} — {artist} ({album})")
    _tag_m4a(output_path, title, artist, album, thumbnail_url)

    print(f"Saved: {output_path}")

    # Import into Music.app
    if not args.no_import:
        print("Importing into Music.app ...")
        import_result = subprocess.run(
            [
                "osascript", "-e",
                f'tell application "Music" to add POSIX file "{output_path}"',
            ],
            capture_output=True, text=True,
        )
        if import_result.returncode != 0:
            print(f"Music.app import failed: {import_result.stderr}", file=sys.stderr)
            print("File is saved — you can drag it into Music.app manually.")
            sys.exit(1)
        print(f"Imported into Music.app: {title} — {artist}")
    else:
        print("Skipped Music.app import (--no-import).")


def main():
    parser = argparse.ArgumentParser(description="Apple Music playlist CLI")
    sub = parser.add_subparsers(dest="command", required=True)

    # match
    p_match = sub.add_parser("match", help="Find songs on Apple Music (batch search with scoring)")
    p_match.add_argument("--songs", required=True, help="JSON array of {artist, title}")

    # search
    p_search = sub.add_parser("search", help="Look up a specific song")
    p_search.add_argument("--query", required=True)
    p_search.add_argument("--limit", type=int, default=5)

    # create
    p_create = sub.add_parser("create", help="Make a playlist")
    p_create.add_argument("--name", required=True)
    p_create.add_argument("--description", default="")
    p_create.add_argument("--track-ids", required=True, help="JSON array of catalog song IDs")
    p_create.add_argument("--upsert", action="store_true", help="Add to existing playlist with same name")

    # list
    sub.add_parser("list", help="List all playlists")

    # tracks
    p_tracks = sub.add_parser("tracks", help="Show tracks in a playlist")
    p_tracks.add_argument("--id", required=True, help="Playlist ID")

    # add
    p_add = sub.add_parser("add", help="Add songs to a playlist")
    p_add.add_argument("--id", required=True, help="Playlist ID")
    p_add.add_argument("--track-ids", required=True, help="JSON array of catalog song IDs")

    # remove
    p_remove = sub.add_parser("remove", help="Remove songs from a playlist")
    p_remove.add_argument("--id", required=True, help="Playlist ID")
    p_remove.add_argument("--track-ids", required=True, help="JSON array of library-song IDs")

    # rename
    p_rename = sub.add_parser("rename", help="Rename or redescribe a playlist")
    p_rename.add_argument("--id", required=True, help="Playlist ID")
    p_rename.add_argument("--name", help="New playlist name")
    p_rename.add_argument("--description", help="New playlist description")

    # reorder
    p_reorder = sub.add_parser("reorder", help="Reorder a playlist")
    p_reorder.add_argument("--id", required=True, help="Playlist ID")
    p_reorder.add_argument("--track-ids", required=True, help="JSON array of catalog song IDs in desired order")

    # yt-search
    p_yts = sub.add_parser("yt-search", help="Search YouTube for audio")
    p_yts.add_argument("--query", required=True, help="YouTube search query")
    p_yts.add_argument("--limit", type=int, default=5, help="Number of results")

    # yt-download
    p_ytd = sub.add_parser("yt-download", help="Download from YouTube and import into Music.app")
    p_ytd.add_argument("--url", required=True, help="YouTube video URL")
    p_ytd.add_argument("--name", help="Custom title/filename (default: video title)")
    p_ytd.add_argument("--artist", help="Artist tag (default: YouTube uploader)")
    p_ytd.add_argument("--album", help="Album tag (default: 'YouTube')")
    p_ytd.add_argument("--no-import", action="store_true", help="Download only, skip Music.app import")

    args = parser.parse_args()

    try:
        {
            "match": cmd_match, "search": cmd_search, "create": cmd_create,
            "list": cmd_list, "tracks": cmd_tracks, "add": cmd_add,
            "remove": cmd_remove, "rename": cmd_rename, "reorder": cmd_reorder,
            "yt-search": cmd_yt_search, "yt-download": cmd_yt_download,
        }[args.command](args)
    except TokenExpiredError as e:
        print(f"\n{e}", file=sys.stderr)
        sys.exit(1)
    except requests.exceptions.HTTPError as e:
        if e.response is not None and e.response.status_code == 429:
            print("\n✗ Rate limited — retries exhausted. Wait 5-10 minutes and try again.", file=sys.stderr)
            sys.exit(2)
        raise


if __name__ == "__main__":
    main()