#!/usr/bin/env python3
"""setlist_samples — what samples does a set actually use?

Parses `s "..."` / `sound "..."` calls out of .tidal tracks, resolves each
sample name to its on-disk folder across the SuperDirt sample roots, and emits
a preload plan. This feeds the smart-preload boot step (task #26): warm exactly
the set's samples before the gig instead of lazy-loading them (crackle + xruns)
on first play mid-set — without loading the whole library.

Usage:
  tools/setlist_samples.py TRACK.tidal [TRACK2.tidal ...]
  tools/setlist_samples.py --last 10        # 10 most-recently-edited live/ tracks
  tools/setlist_samples.py --last 30 --emit-sc   # print SC preload snippet
"""
import re, sys, argparse
from pathlib import Path

HOME = Path.home()
# Sample roots, in the order start_and_midi.scd loads them.
SAMPLE_ROOTS = [
    HOME / ".local/share/SuperCollider/downloaded-quarks/Dirt-Samples",
    HOME / "Work/Sound/Samples/extra",
]
DRUM_MACHINES = HOME / "Work/Sound/Samples/tidal-drum-machines/machines"
LIVE_DIR = HOME / "Work/Sound/Tidal/live"

# ParVagues' dialect names samples via bare strings and `# "x"`, not `s "..."`,
# and mixes them with structure/notes/MIDI-ctrl strings in the same quotes. So
# we extract EVERY word-token from EVERY quoted string and let the on-disk
# sample library be the filter (a token is a sample iff a folder exists for it).
QUOTED = re.compile(r'"([^"]*)"')
WORD = re.compile(r'[A-Za-z][A-Za-z0-9_]*')
# Tidal boolean literals — collide with real Dirt-Samples/{t,f} folders but are
# never meant as samples in the mini-notation, so never preload them.
STOP = {'t', 'f'}
LINE_COMMENT = re.compile(r'--.*$', re.MULTILINE)
BLOCK_COMMENT = re.compile(r'\{-.*?-\}', re.DOTALL)


def extract_names(text):
    """Every candidate token appearing in a quoted string (resolution filters)."""
    text = BLOCK_COMMENT.sub(' ', text)
    text = LINE_COMMENT.sub('', text)
    names = set()
    for s in QUOTED.findall(text):
        names.update(WORD.findall(s))
    return names


def build_index():
    """name -> folder Path, across Dirt-Samples, extra/, and drum machines."""
    idx = {}
    for root in SAMPLE_ROOTS:
        if root.is_dir():
            for d in sorted(root.iterdir()):
                if d.is_dir():
                    idx.setdefault(d.name, d)
    # drum machines: <machine>/<folder>, name = folder basename minus '-'
    if DRUM_MACHINES.is_dir():
        for machine in sorted(DRUM_MACHINES.iterdir()):
            if machine.is_dir():
                for folder in sorted(machine.iterdir()):
                    if folder.is_dir():
                        idx.setdefault(folder.name.replace('-', ''), folder)
    return idx


def last_n_tracks(n):
    tracks = sorted(LIVE_DIR.rglob('*.tidal'),
                    key=lambda p: p.stat().st_mtime, reverse=True)
    return tracks[:n]


# Audio extensions SuperDirt will actually register, compared case-insensitively.
AUDIO_EXT = {'.wav', '.aiff', '.aif', '.flac', '.ogg'}


def bank_file_count(folder):
    """How many files SuperDirt will register for this folder.

    Mirrors `pathMatch(folder +/+ "*")` (DirtSoundLibrary.sc:180): non-recursive,
    and DOTFILES ARE EXCLUDED because a glob of `*` does not match them.

    That last detail is not pedantry. `Dirt-Samples/jbk_kick` is a symlink into the
    rhadamanthe pack, which carries an AppleDouble `._X.wav` twin beside every real
    `X.wav`. A naive `glob('*.wav')` counts 508 where SuperDirt registers 254 — and
    `ls` hides the twins, so the folder looks fine. Counting them would make the
    boot-time integrity check cry wolf on a perfectly good bank.
    """
    n = 0
    for p in folder.iterdir():
        if p.name.startswith('.'):
            continue
        if p.suffix.lower() in AUDIO_EXT and p.is_file():
            n += 1
    return n


def read_setlist(path):
    """One track name/path per line; `#` comments and blank lines ignored.

    A setlist FILE rather than a baked-in list, so reordering the show or swapping
    a track never means editing Python — and so the preload whitelist is derived
    from what will actually be PLAYED. See the warning in main() about --last.
    """
    tracks = []
    for raw in Path(path).read_text().splitlines():
        line = raw.split('#', 1)[0].strip()
        if line:
            tracks += resolve_track(line)
    return tracks


REPO = HOME / "Work/Sound/Tidal"


def resolve_track(arg):
    """Accept a full path OR a bare track name; find it under the repo."""
    p = Path(arg)
    if p.is_file():
        return [p]
    stem = p.name if p.name.endswith('.tidal') else p.name + '.tidal'
    matches = list(LIVE_DIR.rglob(stem)) or list(REPO.rglob(stem))
    if not matches:
        print(f'! no track named {stem!r} under live/ or the repo', file=sys.stderr)
    elif len(matches) > 1:
        print(f'! {stem!r} is ambiguous ({len(matches)} matches) — using {matches[0]}',
              file=sys.stderr)
    return matches[:1]


def emit_sc(resolved, unresolved, tracks):
    """Emit a SELF-CHECKING preload block for start_and_midi.scd to executeFile.

    Two things in here are load-bearing, both learned the hard way on 2026-07-28.

    1. `doNotReadYet` IS RESTORED TO TRUE AT THE END.
       The previous version of this emitter set it to false and left it false. That
       looks harmless and is catastrophic: `readFileIfNecessary` has exactly ONE
       caller (DirtSoundLibrary.sc:259-261) and it is gated on `doNotReadYet`. With
       the flag false, a bank that was registered header-only at boot is NEVER read
       — not "read late", never. So every bank this whitelist MISSED became
       permanently silent, emitting `Buffer UGen: no buffer data` on each event
       (19741 of them in one 6h session) with no error a performer would ever see.
       Restoring the flag makes a miss degrade to lazy-on-demand — a tiny first-hit
       delay — instead of silence. It turns this whitelist from load-bearing into an
       optimisation, which is the only safe thing for it to be on stage.

    2. A per-bank COUNT ASSERTION.
       Re-loading a bank REPLACES it (`addBuffer` defaults appendToExisting = false,
       DirtSoundLibrary.sc:42-46), so this is index-stable and `like_sugar:21` keeps
       meaning the same file. The assertion guards the one case that would break
       that: a whitelist gone stale against the folders on disk. A silently
       renumbered bank mid-set is worse than a silent one.
    """
    print('// preload — set-specific samples. GENERATED by tools/setlist_samples.py.')
    print('// DO NOT HAND-EDIT: regenerate via gig-up.sh, or')
    print('//   tools/setlist_samples.py --setlist setlist_opal2026.txt --emit-sc > preload.scd')
    print('//')
    print('// Warms the set\'s banks so they are read at BOOT instead of racing a')
    print('// first-hit async read mid-set. `doNotReadYet` is restored to true at the')
    print('// end so anything NOT listed here still lazy-loads on demand rather than')
    print('// going permanently silent — see emit_sc() in the generator for why that')
    print('// one line is the difference between a slow first hit and a dead orbit.')
    print('//')
    print(f'// {len(tracks)} track(s) scanned, {len(resolved)} banks, '
          f'{len(unresolved)} non-sample tokens ignored.')
    print()
    print('var ok = 0, bad = [], t0 = Main.elapsedTime;')
    print('~pvPreload = [')
    for name in sorted(resolved):
        folder = resolved[name]
        print(f'\t[ \\{name}, {bank_file_count(folder)}, "{folder}" ],')
    print('];')
    print("""
~dirt.doNotReadYet = false;              // read AUDIO now, not just WAV headers
"\\n=== PRELOAD: % banks, eager ===".format(~pvPreload.size).postln;
~pvPreload.do { |entry|
\tvar name = entry[0], expected = entry[1], folder = entry[2], got;
\t~dirt.loadSoundFiles(folder);
\tgot = (~dirt.soundLibrary.buffers[name] ? []).size;
\tif(got == expected) {
\t\tok = ok + 1;
\t} {
\t\tbad = bad.add("%: expected % files, got %".format(name, expected, got));
\t};
};
~dirt.doNotReadYet = true;               // CRITICAL: keep the lazy fallback alive
"=== PRELOAD: %/% banks OK in % s ===".format(
\tok, ~pvPreload.size, (Main.elapsedTime - t0).round(0.1)).postln;
if(bad.notEmpty) {
\t"=== PRELOAD COUNT MISMATCH — whitelist is STALE, regenerate it ===".postln;
\tbad.do { |m| ("    " ++ m).postln };
};""")


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument('tracks', nargs='*', type=Path)
    ap.add_argument('--last', type=int, metavar='N',
                    help='use the N most-recently-edited live/ tracks (NOT for gig '
                         'preload — see --setlist)')
    ap.add_argument('--setlist', metavar='FILE',
                    help='track list to warm, one per line (# comments ok). PREFER '
                         'THIS for a gig: --last guesses from mtime and silently '
                         'omits set tracks you have not edited recently.')
    ap.add_argument('--emit-sc', action='store_true',
                    help='print a SuperCollider preload snippet for the resolved folders')
    args = ap.parse_args()

    tracks = [p for arg in args.tracks for p in resolve_track(arg)]
    if args.setlist:
        tracks += read_setlist(args.setlist)
    if args.last:
        tracks += last_n_tracks(args.last)
    if not tracks:
        ap.error('give track names/paths, --setlist FILE, or --last N')
    # De-dup while preserving order: the same track can arrive from a setlist and
    # from --last, and loading a bank twice is wasted boot time.
    tracks = list(dict.fromkeys(tracks))

    idx = build_index()
    names, per_track = set(), {}
    for t in tracks:
        try:
            got = extract_names(t.read_text(errors='replace'))
        except OSError as e:
            print(f'! skip {t}: {e}', file=sys.stderr)
            continue
        per_track[t] = got
        names |= got

    resolved = {n: idx[n] for n in names if n in idx and n not in STOP}
    unresolved = sorted(names - resolved.keys())
    folders = sorted({str(p) for p in resolved.values()})

    if args.emit_sc:
        emit_sc(resolved, unresolved, tracks)
        return

    print(f'tracks scanned : {len(per_track)}')
    print(f'candidate toks : {len(names)}  (words seen inside all quoted strings)')
    print(f'SAMPLES found  : {len(resolved)}  -> {len(folders)} folders to warm')
    print(f'  (non-sample tokens = structure/notes/MIDI/synths, ignored: {len(unresolved)})')
    print('\nSAMPLES to preload  (name -> folder):')
    for n in sorted(resolved):
        print(f'  {n:20s} {resolved[n]}')


if __name__ == '__main__':
    main()
