#!/usr/bin/env python3
"""YouTube Music search helper for shanty-dl.

Requires: pip install ytmusicapi

Usage:
    ytmusic_search.py search <query>           # general search
    ytmusic_search.py artist <artist>          # bulk search by artist (up to 100 songs)
    ytmusic_search.py track <artist> <title>   # search for a specific track

Output: JSON array of results to stdout, each with:
    - videoId: YouTube video ID
    - title: track title
    - artist: primary artist name
    - album: album name (if available)
    - duration_seconds: duration in seconds (if available)
"""

import json
import sys

from ytmusicapi import YTMusic


def search_general(yt: YTMusic, query: str) -> list[dict]:
    """Search YouTube Music for songs matching a general query."""
    try:
        results = yt.search(query, filter="songs", limit=5)
    except Exception as e:
        print(f"Search error: {e}", file=sys.stderr)
        return []
    return _format_results(results)


def search_artist_bulk(yt: YTMusic, artist: str) -> list[dict]:
    """Fetch up to 100 songs by an artist."""
    try:
        results = yt.search(artist, filter="songs", limit=100)
    except Exception as e:
        print(f"Search error: {e}", file=sys.stderr)
        return []
    return _format_results(results)


def search_track(yt: YTMusic, artist: str, title: str) -> list[dict]:
    """Search for a specific track by artist and title."""
    try:
        results = yt.search(f"{artist} {title}", filter="songs", limit=5)
    except Exception as e:
        print(f"Search error: {e}", file=sys.stderr)
        return []
    return _format_results(results)


def _format_results(results: list[dict]) -> list[dict]:
    """Convert ytmusicapi results to a simple JSON-friendly format."""
    formatted = []
    for r in results:
        if not r.get("videoId"):
            continue
        artists = r.get("artists", [])
        artist_name = artists[0]["name"] if artists else None
        album = r.get("album")
        album_name = album["name"] if isinstance(album, dict) else None
        formatted.append({
            "videoId": r["videoId"],
            "title": r.get("title", ""),
            "artist": artist_name,
            "album": album_name,
            "duration_seconds": r.get("duration_seconds"),
        })
    return formatted


def main():
    if len(sys.argv) < 3:
        print("Usage: ytmusic_search.py <search|artist|track> <args...>", file=sys.stderr)
        sys.exit(1)

    yt = YTMusic()
    command = sys.argv[1]

    if command == "search":
        query = " ".join(sys.argv[2:])
        results = search_general(yt, query)
    elif command == "artist":
        artist = " ".join(sys.argv[2:])
        results = search_artist_bulk(yt, artist)
    elif command == "track":
        if len(sys.argv) < 4:
            print("Usage: ytmusic_search.py track <artist> <title>", file=sys.stderr)
            sys.exit(1)
        # Artist is argv[2], title is the rest
        artist = sys.argv[2]
        title = " ".join(sys.argv[3:])
        results = search_track(yt, artist, title)
    else:
        print(f"Unknown command: {command}", file=sys.stderr)
        sys.exit(1)

    json.dump(results, sys.stdout)


if __name__ == "__main__":
    main()