#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-2.0-or-later
"""Writes builder/data.js, the data the web builder (builder.html) applies.

    usage: make_builder_data.py [--repo ../CDJ900NXSVirtual]

Our mods (cuecolours, hotcuekeys, ledsaver, credit) are made by the emulator
checkout's mods/patch_cdj900nxs.py for every combination the builder offers;
their bytes do not depend on which of arnoldasr's mods are chosen (checked
below). For each combination the bytes our mods change in the application
go into data.js, except what is the player's own firmware:

  copy  the UTILITY item table (hotcuekeys and ledsaver extend a copy of it)
        is copied from the user's own application in the browser;
  box   the HOT KEYS box is the QUANTIZE label's image recoloured with
        HOT KEYS written on it; the browser makes it from the user's image.

arnoldasr's mods (credited on the site), as patch_cdj900nxs.py splits them
up, are in data.js the same way: per combination of extended_bpm, phasemeter
and rgb_waveform the bytes they change, except the asset blob at 0x08FC2D00,
which starts with a copy of the firmware's resource table: the browser makes
it from the user's application (the port of patch_cdj900nxs.assets()).
The SHA-256 of the finished application for every combination is in data.js
too; the browser checks its result against it.

Before data.js is written, the browser's way of applying the data is redone
here in Python for every combination of arnoldasr's mods and ours and
compared byte for byte with patch_cdj900nxs.py's result.
"""
import argparse
import base64
import hashlib
import itertools
import json
import os
import struct
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
WEB = os.path.dirname(HERE)
OURS = ('cuecolours', 'hotcuekeys', 'ledsaver', 'credit')
NEEDS = {'hotcuekeys': 'cuecolours'}
UT_TABLE_STOCK = 0x08C6C8DC
UT_TABLE_LEN = 17 * 52
BASE = 0x08000000
ASSETS = None


def sha(b):
    return hashlib.sha256(b).hexdigest()


def valid(combo):
    return all(NEEDS[m] in combo for m in combo if m in NEEDS)


def runs(base, final):
    """[(address, bytes)] where final differs from base (short equal gaps joined)."""
    n = len(final)
    base = base + bytes(n - len(base))
    out, i = [], 0
    while i < n:
        if base[i] != final[i]:
            j = i
            while j < n and (base[j] != final[j] or (j + 16 < n and base[j:j + 16] != final[j:j + 16])):
                j += 1
            out.append((BASE + i, final[i:j]))
            i = j
        else:
            i += 1
    return out


def ops_for(stock, final, boxes):
    """The browser's operations for one combination: data, copy, box."""
    table = stock[UT_TABLE_STOCK - BASE:UT_TABLE_STOCK - BASE + UT_TABLE_LEN]
    out = [dict(box, op='box') for box in boxes]    # whole, outside the runs
    masked = bytearray(final)
    padded = stock + bytes(len(final) - len(stock))
    for box in boxes:
        off = box['addr'] - BASE
        masked[off:off + box['len']] = padded[off:off + box['len']]
    for addr, data in runs(stock, bytes(masked)):
        cuts = []                                   # (start, end, op) inside this run
        k = data.find(table)
        while k >= 0:
            cuts.append((k, k + UT_TABLE_LEN, {'op': 'copy', 'addr': addr + k, 'src': UT_TABLE_STOCK,
                                               'len': UT_TABLE_LEN}))
            k = data.find(table, k + UT_TABLE_LEN)
        cuts.sort(key=lambda c: c[0])
        pos = 0
        for lo, hi, op in cuts:
            if lo > pos:
                out.append({'op': 'data', 'addr': addr + pos, 'data': data[pos:lo]})
            out.append(op)
            pos = hi
        if pos < len(data):
            out.append({'op': 'data', 'addr': addr + pos, 'data': data[pos:]})
    return out


def apply_ops(app, stock, ops, hk, size):
    """What builder.js does with the operations (redone here to check them):
    the application grown with zeros to `size` first."""
    out = bytearray(app) + bytes(max(0, size - len(app)))
    for o in ops:
        off = o['addr'] - BASE
        if o['op'] == 'data':
            out[off:off + len(o['data'])] = o['data']
        elif o['op'] == 'copy':
            out[off:off + o['len']] = stock[o['src'] - BASE:o['src'] - BASE + o['len']]
        elif o['op'] == 'assets':
            out[off:off + o['len']] = ASSETS(stock)[0][:o['len']]
        else:
            _, _, px = hk.box_pixels(stock, o['res'], o['bg'], o['ink'])
            out[off:off + o['len']] = px
    return bytes(out)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--repo', default=os.path.join(os.path.dirname(WEB), 'CDJ900NXSVirtual'))
    args = ap.parse_args()
    sys.path.insert(0, os.path.join(args.repo, 'mods'))
    sys.path.insert(0, os.path.join(args.repo, 'mods', 'cdj900nxs'))
    import patch_cdj900nxs as p
    import hotcuekeys as hk
    import cuecolours
    with open(os.path.join(args.repo, 'extract', 'cdj900nxs', 'main_unpacked.bin'), 'rb') as f:
        stock = f.read()[:p.APP_SIZE]
    if sha(stock) != p.APP_SHA256:
        raise SystemExit('not the 1.31 application')

    global ASSETS
    ASSETS = p.assets
    _, labels = hk.routine(hk.DECODER, cuecolours.HOT_TABLE)
    w, h, on = hk.box_pixels(stock, hk.RES_QUANTIZE, 0xF800, 0x0000)
    boxes = [{'addr': labels['PIX_ON'], 'len': len(on), 'res': hk.RES_QUANTIZE, 'bg': 0xF800, 'ink': 0x0000},
             {'addr': labels['PIX_OFF'], 'len': len(on), 'res': hk.RES_QUANTIZE_OFF, 'bg': 0x0000, 'ink': 0x4208}]

    combos = [c for n in range(1, len(OURS) + 1) for c in itertools.combinations(OURS, n) if valid(c)]
    variants, chunks = {}, {}

    def encode(ops):
        enc = []
        for o in ops:
            if o['op'] == 'data':
                key = sha(o['data'])[:16]
                chunks[key] = base64.b64encode(o['data']).decode()
                enc.append(['data', o['addr'], key])
            elif o['op'] == 'copy':
                enc.append(['copy', o['addr'], o['src'], o['len']])
            elif o['op'] == 'assets':
                enc.append(['assets', o['addr'], o['len']])
            else:
                enc.append(['box', o['addr'], o['res'], o['bg'], o['ink']])
        return enc

    # arnoldasr's mods, split up: per combination their bytes, but the asset
    # blob (it starts with a copy of the firmware's resource table) made by
    # the browser from the user's own application
    upstream, up_apps = {}, {(): stock}
    for k in range(1, len(p.UPSTREAM) + 1):
        for up in itertools.combinations(p.UPSTREAM, k):
            app = p.patch(stock, list(up))
            blob_len = len(app) - (p.DATA_BASE - BASE)
            masked = bytearray(app)
            masked[p.DATA_BASE - BASE:] = bytes(blob_len)
            ops = ops_for(stock, bytes(masked), [])
            ops.append({'op': 'assets', 'addr': p.DATA_BASE, 'len': blob_len})
            if apply_ops(stock, stock, ops, hk, len(app)) != app:
                raise SystemExit('%s: the operations do not give the patcher result' % ','.join(up))
            upstream[','.join(up)] = {'size': len(app), 'ops': encode(ops)}
            up_apps[up] = app
            print('%-40s %3d operations' % (','.join(up), len(ops)))
    for combo in combos:
        final = p.patch(stock, list(combo))
        ops = ops_for(stock, final, boxes if 'hotcuekeys' in combo else [])
        if apply_ops(stock, stock, ops, hk, len(final)) != final:
            raise SystemExit('%s: the operations do not give the patcher\'s result' % ','.join(combo))
        variants[','.join(combo)] = {'size': len(final), 'ops': encode(ops)}
        print('%-40s %3d operations' % (','.join(combo), len(ops)))

    # the browser applies ours on top of whatever of arnoldasr's was applied:
    # check every pair, and note the result's SHA-256 for the browser to check
    expect = {}
    for k in range(0, len(p.UPSTREAM) + 1):
        for up in itertools.combinations(p.UPSTREAM, k):
            base = up_apps[up]
            if up:
                expect[','.join(up)] = sha(base)
            for combo in combos:
                final = p.patch(stock, list(up) + list(combo))
                alone = p.patch(stock, list(combo))
                ops = ops_for(stock, alone, boxes if 'hotcuekeys' in combo else [])
                if apply_ops(base, stock, ops, hk, len(alone)) != final:
                    raise SystemExit('%s + %s: differs from the patcher' % (','.join(up), ','.join(combo)))
                expect[','.join(list(up) + list(combo))] = sha(final)
    print('checked: every combination of arnoldasr and ours gives the patcher bytes')

    glyphs = {c: rows for c, rows in hk.GLYPHS.items()}
    data = {
        'stockUpdSha256': '',            # filled from the patcher's upd.py below
        'stockAppSha256': p.APP_SHA256,
        'appSize': p.APP_SIZE,
        'dataBase': p.DATA_BASE,
        'upstream': upstream,
        'ours': variants,
        'expect': expect,
        'box': {'resTable': hk.RES_TABLE, 'width': w, 'height': h, 'glyphs': glyphs,
                'text': 'HOT KEYS', 'gap': 3, 'space': 5},
        'chunks': chunks,
    }
    sys.path.insert(0, os.path.join(os.path.dirname(WEB), 'CDJ900NXSPatcher'))
    import upd
    data['stockUpdSha256'] = upd.STOCK_UPD_SHA256
    for key, b64 in chunks.items():      # no firmware in what we ship
        blob = base64.b64decode(b64)
        for i in range(0, max(0, len(blob) - 47), 8):
            win = blob[i:i + 48]
            if len(set(win)) > 2 and stock.find(win) >= 0:
                raise SystemExit('chunk %s holds firmware bytes (at %d)' % (key, i))
    out = os.path.join(WEB, 'builder', 'data.js')
    os.makedirs(os.path.dirname(out), exist_ok=True)
    with open(out, 'w', encoding='utf-8', newline='\n') as f:
        f.write('// Generated by tools/make_builder_data.py -- do not edit.\n')
        f.write('window.BUILDER_DATA = ')
        json.dump(data, f, separators=(',', ':'))
        f.write(';\n')
    print('wrote %s (%d bytes, %d + %d variants, %d chunks)' % (out, os.path.getsize(out), len(upstream), len(variants), len(chunks)))


if __name__ == '__main__':
    main()