#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.14"
# dependencies = ["kaitaistruct>=0.11", "pyrekordbox>=0.4.4"]
# ///
"""Decode a rekordbox USB export so we can spot how the analyze-path
bucket/id is derived from track metadata.

Reads `export.pdb` from a rekordbox-built USB stick (or `data/q1-rekordbox`)
plus optionally cross-references the rekordbox machine's `master.db` to
correlate Mac-internal UUIDs with the USB-export short hex ids.

Usage:
    uv run tools/decode_rb_export.py <usb-root> [<master.db>]

`<usb-root>` is the directory holding `PIONEER/rekordbox/export.pdb`.
With `<master.db>` we also list the UUID + ContentLink for each track so
the per-track UUID→USB-id derivation can be eye-balled across rows.
"""
import re
import sys
from pathlib import Path

# Make the in-project kaitai schemas importable without an editable install.
HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parent / "src"))

from kaitai.rekordbox_pdb import RekordboxPdb  # noqa: E402
from kaitaistruct import KaitaiStream           # noqa: E402

ANLZ_PATH_RE = re.compile(r"/PIONEER/USBANLZ/([^/]+)/([^/]+)/ANLZ0000\.DAT")


def iter_rows(pdb: RekordboxPdb, page_type):
    """Walk a table's page list, yielding every present row body.
    Tolerates the page chain pointing past EOF (real exports do this on
    the last page — `next_page.index >= num_pages` is a sentinel)."""
    for tref in pdb.tables:
        if tref.type != page_type:
            continue
        pref = tref.first_page
        seen = set()
        while pref is not None and pref.index not in seen:
            seen.add(pref.index)
            try:
                page = pref.body
            except Exception:
                break  # past-EOF next_page (free-pool sentinel)
            if page.is_data_page:
                for rg in page.row_groups:
                    for rref in rg.rows:
                        if rref.present:
                            yield rref.body
            try:
                nxt = page.next_page
            except Exception:
                break
            if nxt is None or nxt.index == 0 or nxt.index == pref.index:
                break
            pref = nxt


def read_pdb_tracks(pdb_path: Path) -> list[dict]:
    """Return one dict per track row: id, title, file_path, analyze_path,
    plus every opaque/seed field plus parsed bucket+anlz_id."""
    with pdb_path.open("rb") as f:
        pdb = RekordboxPdb(False, KaitaiStream(f))
        tracks = []
        for r in iter_rows(pdb, RekordboxPdb.PageType.tracks):
            ap = r.analyze_path.body.text
            m = ANLZ_PATH_RE.match(ap)
            bucket = anlz_id = ""
            if m:
                bucket, anlz_id = m.group(1), m.group(2)
            tracks.append({
                "track_id":            r.id,
                "title":               r.title.body.text,
                "file_path":           r.file_path.body.text,
                "analyze_path":        ap,
                "bucket":              bucket,
                "anlz_id":             anlz_id,
                "opaque_track_id_lo":  r.opaque_track_id_lo,
                "opaque_seed_lo":      r.opaque_seed_lo,
                "opaque_seed_hi":      r.opaque_seed_hi,
                "tempo":               r.tempo,
                "duration":            r.duration,
                "file_size":           r.file_size,
                "bitmask":             r.bitmask,
            })
        return tracks


def maybe_attach_master_db(tracks: list[dict], master_db_path: Path) -> None:
    """Look up each track in master.db by file path and attach UUID/etc."""
    try:
        from pyrekordbox import Rekordbox6Database
        from pyrekordbox.db6 import tables as t
    except ImportError as e:
        print(f"(pyrekordbox unavailable: {e}; skipping master.db lookup)", file=sys.stderr)
        return
    db = Rekordbox6Database(path=str(master_db_path))
    for tr in tracks:
        # Look up by file name as a fallback — full FolderPath may differ if
        # the rekordbox machine mounts the source from a different prefix.
        title = tr["title"]
        rows = db.session.query(t.DjmdContent).filter(t.DjmdContent.Title == title).limit(3).all()
        if not rows:
            tr["rb_uuid"] = ""
            tr["rb_content_link"] = 0
            tr["rb_file_id"] = ""
            tr["rb_master_song_id"] = ""
            continue
        # Prefer one whose FolderPath ends with the same file name
        fname = Path(tr["file_path"]).name
        match = next((r for r in rows if (r.FolderPath or "").endswith(fname)), rows[0])
        tr["rb_uuid"] = match.UUID or ""
        tr["rb_content_link"] = match.ContentLink or 0
        tr["rb_file_id"] = match.rb_file_id or ""
        tr["rb_master_song_id"] = match.MasterSongID or ""


def fmt_int(v, w=10):
    if not isinstance(v, int):
        return f"{str(v):<{w}}"
    return f"{v:>{w-2}d}=0x{v:08x}"[-w:].rjust(w)


def main(argv: list[str]) -> int:
    if len(argv) < 1:
        sys.exit(__doc__)
    usb_root = Path(argv[0]).expanduser()
    pdb_path = usb_root / "PIONEER" / "rekordbox" / "export.pdb"
    if not pdb_path.exists():
        sys.exit(f"export.pdb not found at {pdb_path}")

    tracks = read_pdb_tracks(pdb_path)
    print(f"# Read {len(tracks)} tracks from {pdb_path}")

    if len(argv) >= 2:
        master_db = Path(argv[1]).expanduser()
        if master_db.exists():
            maybe_attach_master_db(tracks, master_db)
            print(f"# Cross-referenced against {master_db}")
        else:
            print(f"# (master.db at {master_db} not found, skipping)")

    # Output: one TSV row per track plus a summary at the end.
    cols = ["track_id", "bucket", "anlz_id", "opaque_track_id_lo",
            "opaque_seed_lo", "opaque_seed_hi",
            "rb_uuid", "rb_content_link", "rb_file_id", "rb_master_song_id",
            "title"]
    print()
    print("\t".join(cols))
    for tr in tracks:
        row = []
        for c in cols:
            v = tr.get(c, "")
            if isinstance(v, int):
                row.append(f"{v}=0x{v:x}" if v > 255 else str(v))
            else:
                row.append(str(v)[:64])
        print("\t".join(row))

    # Quick correlation tests: print numeric distances + ratios so a
    # multi-track export shows obvious patterns at a glance.
    if len(tracks) >= 2:
        print("\n# Successive anlz_id deltas (probe for sequential counter):")
        ids = []
        for tr in tracks:
            try:
                ids.append(int(tr["anlz_id"], 16))
            except (ValueError, TypeError):
                pass
        ids.sort()
        for a, b in zip(ids, ids[1:]):
            print(f"  {a:08x}  →  {b:08x}   Δ = {b-a}")
        print(f"\n# anlz_id range: {min(ids):08x} .. {max(ids):08x}")
        print(f"# bucket set: {sorted({tr['bucket'] for tr in tracks if tr['bucket']})}")

    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))