#!/usr/bin/env python3
"""Extract the raw firmware image from a KMI QuNexus firmware .syx file.

Container format (mirrors SysExEncDecode in qt-qunexus/source/midiio/sysexencdecode.cpp):

  F0 00 01 5F 7A 19 [pad 00...] 01 <7-in-8 encoded stream> F7

  * 00 01 5F 7A = manufacturer id bytes, 19 = product, 00 = format
  * 01          = SX_PACKET_START (sysexencdecode.cpp:33)
  * the encoded stream packs 7 data bytes as their low 7 bits followed by one
    byte holding their high bits (bit j = high bit of data byte j) --
    midi_sx_encode_char(), sysexencdecode.cpp:1830, SX_ENCODE_LEN = 7

  The decoded stream is a packet preamble (00 02 <cat> <type> <crc16>) followed
  by framed records:  03 <len+1> <binary Intel HEX record>
  where the record is  3A LL AAAA TT <data...> CC  (checksum = two's complement
  of the sum of LL..data), i.e. Intel HEX in binary rather than ASCII form.
"""

import sys
import argparse


def sysex_messages(data):
    msgs, i = [], 0
    while True:
        s = data.find(b"\xf0", i)
        if s < 0:
            break
        e = data.find(b"\xf7", s)
        if e < 0:
            break
        msgs.append(data[s:e + 1])
        i = e + 1
    return msgs


def decode_7in8(buf):
    """Undo midi_sx_encode_char(): 7 low-7-bit bytes, then a high-bits byte."""
    out = bytearray()
    for i in range(0, len(buf) - 7, 8):
        hi = buf[i + 7]
        for j in range(7):
            out.append(buf[i + j] | (0x80 if (hi >> j) & 1 else 0))
    return bytes(out)


def decode_message(msg):
    """Strip the sysex header up to SX_PACKET_START and 7-in-8 decode the body."""
    body = msg[1:-1]                      # drop F0 / F7
    i = 6                                 # manufacturer id (4) + product + format
    while i < len(body) and body[i] == 0x00:
        i += 1                            # padding before the packet start
    if i >= len(body) or body[i] != 0x01:  # SX_PACKET_START
        return b""
    return decode_7in8(body[i + 1:])


def parse_hex_records(stream):
    """Scan the decoded stream for checksum-valid binary Intel HEX records."""
    records, i, skipped = [], 0, 0
    while i < len(stream):
        if stream[i] != 0x3A:
            i += 1
            skipped += 1
            continue
        if i + 5 > len(stream):
            break
        ln = stream[i + 1]
        end = i + 5 + ln            # 3A LL AA AA TT data...
        if end >= len(stream):
            break
        rec = stream[i + 1:end]     # LL AAAA TT data (checksummed span)
        if (sum(rec) + stream[end]) & 0xFF != 0:
            i += 1                  # not a real record, keep scanning
            skipped += 1
            continue
        records.append((stream[i + 4],                      # type
                        (stream[i + 2] << 8) | stream[i + 3],  # address
                        bytes(rec[4:])))                    # data
        i = end + 1
    return records, skipped


def build_image(records):
    """Apply Intel HEX records (types 00/01/02/04) to a sparse address space."""
    mem, base, eof = {}, 0, False
    for rtype, addr, data in records:
        if rtype == 0x00:
            for k, b in enumerate(data):
                mem[base + addr + k] = b
        elif rtype == 0x01:
            eof = True
        elif rtype == 0x02 and len(data) == 2:
            base = ((data[0] << 8) | data[1]) << 4
        elif rtype == 0x04 and len(data) == 2:
            base = ((data[0] << 8) | data[1]) << 16
    return mem, eof


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("syx")
    ap.add_argument("-o", "--out", help="write flat binary image here")
    ap.add_argument("--hex", help="also write a standard ASCII .hex file here")
    ap.add_argument("--fill", default="0xFF", help="gap fill byte (default 0xFF)")
    args = ap.parse_args()

    data = open(args.syx, "rb").read()
    msgs = sysex_messages(data)
    stream = b"".join(decode_message(m) for m in msgs)
    records, skipped = parse_hex_records(stream)
    mem, eof = build_image(records)

    print(f"file            : {args.syx}")
    print(f"sysex messages  : {len(msgs)}")
    print(f"decoded stream  : {len(stream)} bytes ({skipped} non-record bytes skipped)")
    print(f"hex records     : {len(records)}  (EOF record seen: {eof})")
    if not mem:
        print("no data records found", file=sys.stderr)
        return 1

    lo, hi = min(mem), max(mem)
    types = sorted({t for t, _, _ in records})
    print(f"record types    : {[hex(t) for t in types]}")
    print(f"address range   : 0x{lo:08X} - 0x{hi:08X}  ({hi - lo + 1} bytes span)")
    print(f"bytes covered   : {len(mem)}  (gaps: {hi - lo + 1 - len(mem)})")

    fill = int(args.fill, 0)
    img = bytes(mem.get(a, fill) for a in range(lo, hi + 1))
    if args.out:
        open(args.out, "wb").write(img)
        print(f"wrote           : {args.out} ({len(img)} bytes, base 0x{lo:08X})")
    if args.hex:
        lines, base = [], None
        for a in range(lo, hi + 1, 16):
            chunk = bytes(mem.get(a + k, fill) for k in range(min(16, hi + 1 - a)))
            upper = a >> 16
            if upper != base:
                base = upper
                rec = bytes([2, 0, 0, 4, upper >> 8, upper & 0xFF])
                lines.append(":" + (rec + bytes([(-sum(rec)) & 0xFF])).hex().upper())
            rec = bytes([len(chunk), (a >> 8) & 0xFF, a & 0xFF, 0]) + chunk
            lines.append(":" + (rec + bytes([(-sum(rec)) & 0xFF])).hex().upper())
        lines.append(":00000001FF")
        open(args.hex, "w").write("\n".join(lines) + "\n")
        print(f"wrote           : {args.hex}")
    return 0


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