#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-2.0-or-later
"""CDJ-900NXS firmware patcher: the official 1.31 update with the mods you
choose, as a new C900NXS.UPD for the player's own USB update.

    python cdj900nxs_patcher.py C900NXS.UPD                 # choose in a list, write out/...
    python cdj900nxs_patcher.py C900NXS.UPD --mods beatjump,cuecolours
    python cdj900nxs_patcher.py C900NXS.UPD --stock         # unmodified, to go back
    python cdj900nxs_patcher.py --list

  * Only the original update is accepted (its SHA-256 is checked).
  * Only the main application is replaced. The boot loader and recovery
    loader (flash 0x0-0x3FFFF), the S-record header and entry, and the DRIV
    and PANL sections are copied from the original byte for byte, and checked
    again on the finished file before it is written.
  * The new file says 1.32 in its MAIN header so that the player's updater
    does not skip it as the version it already has; the application itself
    still reports 1.31 (as arnoldasr's builds do). Only that label differs.
  * Before anything is built a disclaimer must be accepted: the mods and
    this patcher were developed with Claude Code (--accept where no one can
    type).
  * Nothing here talks to a player. You copy the file to a USB stick and run
    the update on the player yourself, at your own risk: nothing of this has
    been tested on hardware except the relabelled stock update (by arnoldasr),
    and recovery from an application that crashes at start-up is untested.

The mods come from the CDJ900NXSVirtual checkout next to this folder
(--repo): mods/patch_cdj900nxs.py, which also needs arnoldasr's checkout in
mods/cdj900nxs/upstream for his mods (extended_bpm, phasemeter, advanced_waveforms).
"""
import argparse
import os
import sys
import time

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import upd  # noqa: E402

DEFAULT_REPO = os.path.join(os.path.dirname(HERE), 'CDJ900NXSVirtual')
UPSTREAM_MODS = ('extended_bpm', 'phasemeter', 'advanced_waveforms')
NEEDS = {'hotcuekeys': 'cuecolours',    # the hot cues the keys play are read by cuecolours
         'beatjump': 'hotcuekeys'}      # hotcuekeys' engine makes the jumps
LABEL = '1.32'
# not ticked at the start of the list, and why
KNOWN_FAULTS = {}


def load_repo(repo):
    """patch_cdj900nxs and the LZSS packer from the emulator checkout."""
    mods_dir = os.path.join(repo, 'mods')
    fw_dir = os.path.join(repo, 'scripts', 'firmware')
    if not os.path.isfile(os.path.join(mods_dir, 'patch_cdj900nxs.py')):
        raise SystemExit('error: no mods/patch_cdj900nxs.py in %s (use --repo)' % repo)
    for d in (mods_dir, fw_dir):
        if d not in sys.path:
            sys.path.insert(0, d)
    import patch_cdj900nxs
    import lzss_encode
    return patch_cdj900nxs, lzss_encode


def available(patcher, repo):
    """(name, description, usable, why not) for every mod."""
    upstream = os.path.isdir(os.environ.get('CDJ900NXS_MODS_SRC')
                             or os.path.join(repo, 'mods', 'cdj900nxs', 'upstream'))
    out = []
    for name, what in patcher.MODS.items():
        if name in KNOWN_FAULTS:
            what = '%s  [%s]' % (what, KNOWN_FAULTS[name])
        ok = upstream or name not in UPSTREAM_MODS
        out.append((name, what, ok, '' if ok else "needs arnoldasr's checkout in mods/cdj900nxs/upstream"))
    return out


# --- choosing ------------------------------------------------------------------

def _keys():
    """A function returning 'up', 'down', 'space', 'enter', 'all', 'quit' per
    key press, or None when the terminal cannot be read key by key."""
    if not sys.stdin.isatty():
        return None
    if os.name == 'nt':
        import msvcrt

        def read():
            c = msvcrt.getwch()
            if c in ('\x00', '\xe0'):
                return {'H': 'up', 'P': 'down'}.get(msvcrt.getwch())
            return {' ': 'space', '\r': 'enter', 'a': 'all', 'q': 'quit', '\x1b': 'quit',
                    '\x03': 'quit'}.get(c.lower())
        return read
    import termios
    import tty

    def read():
        fd = sys.stdin.fileno()
        old = termios.tcgetattr(fd)
        try:
            tty.setraw(fd)
            c = sys.stdin.read(1)
            if c == '\x1b':
                if sys.stdin.read(1) == '[':
                    return {'A': 'up', 'B': 'down'}.get(sys.stdin.read(1))
                return 'quit'
        finally:
            termios.tcsetattr(fd, termios.TCSADRAIN, old)
        return {' ': 'space', '\r': 'enter', '\n': 'enter', 'a': 'all', 'q': 'quit',
                '\x03': 'quit'}.get(c.lower())
    return read


def _toggle(chosen, name, mods):
    """Tick or untick one mod: ticking one ticks what it needs, unticking
    one unticks what needs it (so every mod can be unticked)."""
    if name in chosen:
        drop = {name}
        while True:
            more = {n for n, need in NEEDS.items() if need in drop} - drop
            if not more:
                break
            drop |= more
        chosen.difference_update(drop)
    else:
        chosen.add(name)
        need = NEEDS.get(name)
        while need and any(m[0] == need and m[2] for m in mods):
            chosen.add(need)
            need = NEEDS.get(need)


def choose(mods):
    """The interactive list: arrows move, space ticks, enter goes on."""
    read = _keys()
    if read is None:
        return choose_by_number(mods)
    if os.name == 'nt':
        os.system('')               # VT sequences on in the Windows console
    chosen = {m[0] for m in mods if m[2] and m[0] not in KNOWN_FAULTS}
    at = 0
    lines = len(mods) + 3

    def draw(first):
        if not first:
            sys.stdout.write('\x1b[%dA' % lines)
        sys.stdout.write('\x1b[2K  Mods (arrows move, space ticks, a = all/none, enter goes on, q quits)\n')
        for i, (name, what, ok, why) in enumerate(mods):
            mark = '[x]' if name in chosen else ('[ ]' if ok else ' - ')
            cursor = '>' if i == at else ' '
            text = '%s %s %-13s %s' % (cursor, mark, name, what if ok else why)
            sys.stdout.write('\x1b[2K' + ('\x1b[7m%s\x1b[0m' % text if i == at else text)[:200] + '\n')
        need = ', '.join('%s needs %s' % kv for kv in NEEDS.items())
        sys.stdout.write('\x1b[2K  (%s)\n\x1b[2K\n' % need)
        sys.stdout.flush()

    draw(True)
    while True:
        k = read()
        if k == 'up':
            at = (at - 1) % len(mods)
        elif k == 'down':
            at = (at + 1) % len(mods)
        elif k == 'space' and mods[at][2]:
            _toggle(chosen, mods[at][0], mods)
        elif k == 'all':
            usable = {m[0] for m in mods if m[2]}
            chosen = set() if chosen == usable else usable
        elif k == 'enter':
            break
        elif k == 'quit':
            raise SystemExit('nothing written')
        draw(False)
    return [m[0] for m in mods if m[0] in chosen]


def choose_by_number(mods):
    """For terminals that cannot be read key by key (mintty, pipes)."""
    chosen = {m[0] for m in mods if m[2] and m[0] not in KNOWN_FAULTS}
    while True:
        print('\n  Mods (all usable ones without known faults are ticked):')
        for i, (name, what, ok, why) in enumerate(mods, 1):
            mark = '[x]' if name in chosen else ('[ ]' if ok else ' - ')
            print('  %d %s %-13s %s' % (i, mark, name, what if ok else why))
        answer = input('  Numbers to tick/untick (e.g. 1 3), enter to go on, q to quit: ').strip()
        if not answer:
            break
        if answer.lower() == 'q':
            raise SystemExit('nothing written')
        for tok in answer.replace(',', ' ').split():
            if tok.isdigit() and 1 <= int(tok) <= len(mods) and mods[int(tok) - 1][2]:
                _toggle(chosen, mods[int(tok) - 1][0], mods)
    return [m[0] for m in mods if m[0] in chosen]


# --- the disclaimer ------------------------------------------------------------

DISCLAIMER = '''
  Please read before anything is built:

  The mods of this patcher (beatjump, cuecolours, hotcuekeys) and the patcher
  itself were developed with Claude Code, Anthropic's AI coding assistant,
  by reverse engineering the firmware in an emulator. They have not been
  reviewed by Pioneer DJ / AlphaTheta and have not been tested on a player
  except where the README says so. Flashing modified firmware can leave the
  player unusable. You use the result at your own risk; there is no warranty.
'''


def confirm(accepted, ask=input):
    """The disclaimer, which must be accepted before a build: by typing yes,
    or with --accept where no one can type."""
    print(DISCLAIMER)
    if accepted:
        print('  accepted with --accept')
        return
    if not sys.stdin.isatty():
        raise SystemExit('nothing written: accept the disclaimer by typing yes, or with --accept')
    try:
        answer = ask('  Type yes to accept and go on: ').strip().lower()
    except EOFError:
        answer = ''
    if answer not in ('yes', 'y', 'ja', 'j'):
        raise SystemExit('nothing written: the disclaimer was not accepted')


# --- building ------------------------------------------------------------------

def build(stock, names, patcher, lzss, label=LABEL, log=print):
    """(update bytes, application, list of checks)."""
    if names:
        try:
            app = patcher.patch(stock.application, names)
        except SystemExit as e:
            raise upd.UpdError(str(e))
        log('packing the application (LZSS, about a minute and a half)...')
        t = time.time()
        stream = lzss.pack(app)
        log('  %d -> %d bytes in %.0f s' % (len(app), len(stream), time.time() - t))
    else:
        import struct
        app = stock.application
        size = struct.unpack_from('<I', stock.flash, upd.APP_START)[0]
        stream = stock.flash[upd.APP_START + 4:upd.APP_START + 4 + size]
    raw = upd.build(stock, app, stream, label)
    checks = upd.verify(raw, stock, app, label)
    return raw, app, checks


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument('stock', nargs='?', help='the official C900NXS.UPD (firmware 1.31)')
    ap.add_argument('--out', help='new folder for the built C900NXS.UPD (default out/<date>-<mods>)')
    ap.add_argument('--mods', help='comma-separated mods instead of the list')
    ap.add_argument('--stock', dest='back', action='store_true',
                    help='no mods: the original application labelled %s, to go back to stock' % LABEL)
    ap.add_argument('--repo', default=DEFAULT_REPO, help='the CDJ900NXSVirtual checkout (default %(default)s)')
    ap.add_argument('--list', action='store_true', help='list the mods and exit')
    ap.add_argument('--accept', action='store_true',
                    help='accept the disclaimer (developed with Claude Code, at your own risk) without asking')
    args = ap.parse_args(argv)

    patcher, lzss = load_repo(args.repo)
    mods = available(patcher, args.repo)
    if args.list:
        for name, what, ok, why in mods:
            print('%-13s %s%s' % (name, what, '' if ok else '  (%s)' % why))
        return 0
    if not args.stock:
        ap.error('give the official C900NXS.UPD')

    print('CDJ-900NXS firmware patcher')
    raw = open(args.stock, 'rb').read()
    try:
        stock = upd.Stock(raw)
    except upd.UpdError as e:
        raise SystemExit('error: %s' % e)
    print('  %s: the original 1.31 update (SHA-256 %s...), checked' % (os.path.basename(args.stock),
                                                                       upd.STOCK_UPD_SHA256[:16]))
    print('  boot loader and recovery: %d records below 0x%X, kept as they are'
          % (len(stock.boot_records()), upd.APP_START))

    if args.back:
        names = []
    elif args.mods is not None:
        names = [n for n in args.mods.replace(' ', ',').split(',') if n]
        unknown = [n for n in names if n not in patcher.MODS]
        if unknown:
            raise SystemExit('error: unknown mod(s) %s; --list shows them' % ', '.join(unknown))
        blocked = [m[0] for m in mods if m[0] in names and not m[2]]
        if blocked:
            raise SystemExit('error: %s: %s' % (', '.join(blocked), mods[[m[0] for m in mods].index(blocked[0])][3]))
        for name, need in NEEDS.items():
            if name in names and need not in names:
                raise SystemExit('error: %s needs %s' % (name, need))
    else:
        names = choose(mods)
        if not names:
            print('  no mod chosen: use --stock for the unmodified update')
            return 1
    print('  mods: %s' % (', '.join(names) or 'none (stock application)'))
    confirm(args.accept)

    try:
        out_raw, app, checks = build(stock, names, patcher, lzss)
    except upd.UpdError as e:
        raise SystemExit('error: %s; nothing written' % e)
    folder = args.out or os.path.join(HERE, 'out', time.strftime('%Y%m%d-%H%M%S') + '-' +
                                      ('-'.join(names) if names else 'stock'))
    if os.path.exists(os.path.join(folder, 'C900NXS.UPD')):
        raise SystemExit('error: %s already holds a C900NXS.UPD; nothing written' % folder)
    os.makedirs(folder, exist_ok=True)
    path = os.path.join(folder, 'C900NXS.UPD')
    with open(path, 'wb') as f:
        f.write(out_raw)
    report = ['CDJ-900NXS firmware patcher', 'source: %s (SHA-256 %s)' % (args.stock, upd.STOCK_UPD_SHA256),
              'mods: %s' % (', '.join(names) or 'none'), 'C900NXS.UPD SHA-256 %s' % upd.sha256(out_raw),
              'application SHA-256 %s, %d bytes' % (upd.sha256(app), len(app)),
              'MAIN label %s (the application still reports 1.31)' % LABEL, 'checks:'] + \
             ['  - ' + c for c in checks]
    with open(os.path.join(folder, 'report.txt'), 'w', encoding='utf-8') as f:
        f.write('\n'.join(report) + '\n')
    print('  checked:')
    for c in checks:
        print('    - ' + c)
    print('  written: %s' % path)
    print('  SHA-256 %s' % upd.sha256(out_raw))
    print('  Copy it to the root of a FAT32 USB stick and run the update from the player.')
    print('  Untested on hardware: use at your own risk; keep this folder\'s report.txt.')
    return 0


if __name__ == '__main__':
    sys.exit(main())