drift/scripts/search.py

#!/usr/bin/env python3
"""Fuzzy search artists in playlists.db and show their similar artists."""

import curses
import sqlite3
import sys
from pathlib import Path


def find_db():
    """Look for playlists.db in cwd, then script's parent dir."""
    for base in [Path.cwd(), Path(__file__).resolve().parent.parent]:
        p = base / "playlists.db"
        if p.exists():
            return str(p)
    print("Could not find playlists.db", file=sys.stderr)
    sys.exit(1)


def load_artists(db_path):
    conn = sqlite3.connect(db_path)
    rows = conn.execute(
        "SELECT mbid, COALESCE(name, mbid) FROM artists ORDER BY name"
    ).fetchall()
    conn.close()
    return rows  # [(mbid, display_name), ...]


def get_similar(db_path, mbid):
    conn = sqlite3.connect(db_path)
    rows = conn.execute(
        "SELECT similar_name, match_score FROM similar_artists "
        "WHERE artist_mbid = ?1 ORDER BY match_score DESC",
        (mbid,),
    ).fetchall()
    conn.close()
    return rows


def fuzzy_match(query, name):
    """Simple fuzzy: all query chars appear in order in name."""
    name_lower = name.lower()
    qi = 0
    for ch in name_lower:
        if qi < len(query) and ch == query[qi]:
            qi += 1
    return qi == len(query)


def run_tui(stdscr, db_path):
    curses.curs_set(0)
    curses.use_default_colors()
    curses.init_pair(1, curses.COLOR_BLACK, curses.COLOR_CYAN)
    curses.init_pair(2, curses.COLOR_CYAN, -1)
    curses.init_pair(3, curses.COLOR_WHITE, -1)

    artists = load_artists(db_path)
    query = ""
    selected = 0
    scroll = 0

    while True:
        stdscr.erase()
        h, w = stdscr.getmaxyx()

        # filter
        q = query.lower()
        if q:
            filtered = [(m, n) for m, n in artists if fuzzy_match(q, n)]
        else:
            filtered = artists

        selected = max(0, min(selected, len(filtered) - 1))

        # prompt
        prompt = f" > {query}"
        stdscr.addnstr(0, 0, prompt, w, curses.color_pair(2) | curses.A_BOLD)
        count_str = f"  {len(filtered)}/{len(artists)}"
        if len(prompt) + len(count_str) < w:
            stdscr.addstr(0, len(prompt), count_str, curses.color_pair(3))

        # artist list
        list_h = h - 1
        if list_h < 1:
            stdscr.refresh()
            continue

        if selected < scroll:
            scroll = selected
        if selected >= scroll + list_h:
            scroll = selected - list_h + 1

        for i in range(list_h):
            idx = scroll + i
            if idx >= len(filtered):
                break
            _, name = filtered[idx]
            attr = curses.color_pair(1) if idx == selected else curses.A_NORMAL
            stdscr.addnstr(i + 1, 0, f" {name}", w, attr)

        stdscr.refresh()

        key = stdscr.get_wch()

        if key == "\x1b":  # Esc
            return
        elif key == curses.KEY_UP or key == "\x10":  # Up / Ctrl-P
            selected = max(0, selected - 1)
        elif key == curses.KEY_DOWN or key == "\x0e":  # Down / Ctrl-N
            selected = min(len(filtered) - 1, selected + 1)
        elif key == "\n" or key == curses.KEY_ENTER:
            if filtered:
                show_similar(stdscr, db_path, filtered[selected])
        elif key in (curses.KEY_BACKSPACE, "\x7f", "\x08"):
            query = query[:-1]
            selected = 0
            scroll = 0
        elif isinstance(key, str) and key.isprintable():
            query += key
            selected = 0
            scroll = 0


def show_similar(stdscr, db_path, artist):
    mbid, name = artist
    similar = get_similar(db_path, mbid)

    curses.curs_set(0)
    stdscr.erase()
    h, w = stdscr.getmaxyx()

    title = f" Similar to {name}"
    stdscr.addnstr(0, 0, title, w, curses.color_pair(2) | curses.A_BOLD)
    stdscr.addnstr(h - 1, 0, " [q] back", w, curses.color_pair(3))

    scroll = 0
    list_h = h - 2

    while True:
        for i in range(list_h):
            stdscr.move(i + 1, 0)
            stdscr.clrtoeol()
            idx = scroll + i
            if idx >= len(similar):
                continue
            sname, score = similar[idx]
            line = f"  {score:5.2f}  {sname}"
            stdscr.addnstr(i + 1, 0, line, w)

        stdscr.refresh()
        key = stdscr.get_wch()

        if key in ("q", "Q", "\x1b"):
            return
        elif key == curses.KEY_UP or key == "\x10":
            scroll = max(0, scroll - 1)
        elif key == curses.KEY_DOWN or key == "\x0e":
            if scroll + list_h < len(similar):
                scroll += 1


def main():
    db_path = find_db()
    curses.wrapper(run_tui, db_path)


if __name__ == "__main__":
    main()