#!/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.
  * 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 (code, utility, overview).
"""
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 = ('code', 'utility', 'overview')
NEEDS = {'utility': 'code'}     # the PHASE METER switch chooses the code mod's meter
LABEL = '1.32'


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():
        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 _fix_needs(chosen, mods):
    for name, need in NEEDS.items():
        if name in chosen and need not in chosen and any(m[0] == need and m[2] for m in mods):
            chosen.add(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]}
    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 %-11s %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]:
            chosen ^= {mods[at][0]}
            _fix_needs(chosen, 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]}
    while True:
        print('\n  Mods (all usable ones 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 %-11s %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]:
                chosen ^= {mods[int(tok) - 1][0]}
        _fix_needs(chosen, mods)
    return [m[0] for m in mods if m[0] in chosen]


# --- 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')
    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('%-11s %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)'))

    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())