#!/usr/bin/env python3
"""
MUSCLE (U) bytecode-VM disassembly generator.

Emits a decoded listing for every BytecodeProgram / StandaloneBytecodeProgram and
patches prg.s so each program carries a readable token-by-token listing above its
unchanged .byte data. No byte is altered: the .byte lines are copied verbatim and
only comment lines are inserted.

Encoding (derived from RunBytecodeInterpreter $9AE8 and BytecodeHandlerTable $9237):
  * byte in $C8..$EF -> opcode, handler = table[byte - $C8]
  * any other byte   -> literal operand, pushed to $0033[Y] and accumulated in $003C[Y]
  * opcodes $CB $D1 $D7 $E1 additionally consume the NEXT stream byte inline
    (their handlers execute `lda ($1A),Y / iny`)
"""
import sys, re

ROM = open(sys.argv[1], 'rb').read()
PRG_S = sys.argv[2]
BASE = 0x8000

PROG_TABLE = 0x8CEF
HANDLER_TABLE = 0x9237
INTERP = 0x9AE8

HANDLERS = []
for i in range(40):
    o = HANDLER_TABLE - BASE + i * 2
    HANDLERS.append(ROM[o] | (ROM[o + 1] << 8))

TABLE = []
for i in range(38):
    o = PROG_TABLE - BASE + i * 2
    TABLE.append(ROM[o] | (ROM[o + 1] << 8))

# opcodes that consume one inline stream byte (verified against handler entry points)
INLINE = {0xCB, 0xD1, 0xD7, 0xE1}

# mnemonic + role, read off each handler's own straight-line path
# (up to its first `jmp $9AE8` / `jmp $9B25`).  Confidence is called out where a
# handler simply tail-calls a game routine.
OPS = {
    0xC8: ('VM_CALL_82E1', 'tail-calls the game routine at $82E1'),
    0xC9: ('LOOP_END', 'decrements the $7C loop counter; if still non-zero rewinds the '
                       'stream pointer from $7E, re-running the loop body'),
    0xCA: ('VM_CALL_8765', 'tail-calls the game routine at $8765'),
    0xCB: ('LOOP_BEGIN', 'reads the inline iteration count, stores it in $7C and the '
                         'body address in $7E'),
    0xCC: ('SET_STEP_12', 'sets the step byte $4B to $12'),
    0xCD: ('ACC_ADD_1_STEP', 'SaveStreamPos; adds 1 to the accumulator via $9653, then '
                             'folds the step byte into $0033[Y]'),
    0xCE: ('WAIT_FRAME_SLOT', 'advance-two-frames path: $4F-=1, $A4FF, default $F0=$80, '
                              'and kick a $D1=$FF sound sequence'),
    0xCF: ('SUB_TEN', '$4F-=10 and $F0=$D0'),
    0xD0: ('RESET_ACCUM', 'clears $F0 and zeroes the low operand bytes $0033[Y]/$0034[Y], '
                          'then starts the next actor'),
    0xD1: ('SET_RATE', 'stores the inline byte into $43 as this actor\'s rate'),
    0xD2: ('MOVE_DIAG', 'position update: $4F+=2, $9A3F, a $A660 helper call and two '
                       '$9A2E operand pushes'),
    0xD3: ('MOVE_OTHER', 'position update variant: two $9A24/$9A2E pushes and an $A660 '
                         'helper call'),
    0xD4: ('SET_STEP_2', 'SaveStreamPos, step byte $4B:=2 and add 1 to the accumulator'),
    0xD5: ('STEP_ACCUM', 'subtracts $4B>>1 from the accumulator with $FD/$06 saturation '
                         'clamping'),
    0xD6: ('TICK_STEP', 'step countdown: on reaching zero sets $0321=$FF'),
    0xD7: ('DRAW_SPRITE', 'uses the inline byte as a count for the $A660 sprite/blit helper'),
    0xD8: ('ANIMATE_CHECK', 'animation-frame test via $97FC/$A4DF, then a $99E2 call'),
    0xD9: ('ANIMATE_ADVANCE', 'animation-frame advance: $4F+=2, $5A:=$FF, $A660 call'),
    0xDA: ('SET_DIRECTION', 'sets the $5C direction flag from bit 0 of $24'),
    0xDB: ('PLAY_SEQ_BD20', 'runs sound sequence $BD20 through the $99E4 dispatcher'),
    0xDC: ('TIMER_FF', 'decrements the $60 timer; on zero jumps to $94CF, otherwise adds '
                       '$FF (-1) to the accumulator'),
    0xDD: ('TIMER_01', 'decrements the $60 timer; clears $62 and adds 1 to the accumulator'),
    0xDE: ('MOVE_ACTOR_B0BA', 'accumulator-indexed $A660 move plus a $99E4 call with '
                              'descriptor $B0BA'),
    0xDF: ('MOVE_ACTOR_DIR', 'direction-dependent position update over $F6/$F7'),
    0xE0: ('SET_FLAG_FF', 'stores $FF into $82'),
    0xE1: ('SET_WAIT', 'stores the inline byte into $7A as this actor\'s wait/step count'),
    0xE2: ('COND_FA', 'tests the $5A state register and the $4E flag before branching '
                      'through $96C8'),
    0xE3: ('COND_FB', 'the $FB half of the $96C8 condition test'),
    0xE4: ('ACCUM_PLUS_6', 'SaveStreamPos and adds 6 to the accumulator $3A[Y]'),
    0xE5: ('TICK_WAIT_STEP', 'decrements the $7A wait counter; runs the body while '
                             'non-zero'),
    0xE6: ('MOVE_ACTOR_ALT', 'alternate position update: $9585 test, $9A3F and two '
                              'operand pushes'),
    0xE7: ('SWAP_OPERANDS', 'exchanges the $39[Y]/$3A[Y] operand pair around the $A4FF '
                            'and $A4F2 helpers'),
    0xE8: ('STEP_DOWN', 'SaveStreamPos, $9732, step $4B:=2 and accumulate'),
    0xE9: ('STEP_DOWN_COND', 'as $E8 but gated on the $5A state sign bit'),
    0xEA: ('SET_ACTIVE_80', 'sets $F0 to $80 and falls into the $94CF accumulator clear'),
    0xEB: ('CLEAR_ACTIVE', 'clears $F0 when it currently holds $80'),
    0xEC: ('ADD_STEP_ACCUM', 'adds the step byte into $0033[Y], steps $3A[Y] down and '
                              'accumulates'),
    0xED: ('SUB_STEP_ACCUM', 'adds 1 to $3A[Y] and subtracts the step from $0033[Y]'),
    0xEE: ('ROT_STEP_ADD', 'rotates the step byte, adds it into $0033[Y], steps $3A[Y] '
                           'up by 2 and accumulates'),
    0xEF: ('ROT_STEP_SUB', 'rotates the step byte, complements it, subtracts from '
                           '$0033[Y], and calls $82ED'),
}

LIT_NOTE = {
    0xB8: 'VRAM run fill selector (VramRunFillValues[0..11] span $B8-$C3)',
}


def rd(off):
    """Read one PRG byte by 16-bit PRG offset."""
    return ROM[off]


def cpu(off):
    return BASE + off


def decode(a, b):
    """Yield (addr, opcode_or_None, operand_or_None, raw_bytes)."""
    x = a
    while x < b:
        v = rd(x)
        if 0xC8 <= v <= 0xEF:
            if v in INLINE and x + 1 < b:
                yield (x, v, rd(x + 1), 2)
                x += 2
                continue
            yield (x, v, None, 1)
            x += 1
        else:
            yield (x, None, None, 1)
            x += 1


def listing(a, b, indent='    '):
    out = []
    cur = None
    nops = nlit = 0
    for addr, op, arg, ln in decode(a, b):
        if op is not None:
            nops += 1
            name, role = OPS[op]
            if arg is not None:
                txt = f'{name} ${arg:02X}'
            else:
                txt = name
            body = f'${cpu(addr):04X}  {txt:<28} -- {role}'
        else:
            nlit += 1
            v = rd(addr)
            note = LIT_NOTE.get(v, '')
            body = f'${cpu(addr):04X}  .byte ${v:02X}' + (f'                -- {note}' if note else '')
        out.append(indent + '; ' + body)
    return out, nops, nlit


def main():
    lines = open(PRG_S).read().split('\n')
    uniq = sorted(set(TABLE))
    stats = {}
    out = []
    i = 0
    N = len(lines)
    # The two blocks that only *looked* like bytecode programs are listed too,
    # so their byte streams keep their decoded listing; the listing says so
    # itself (no opcode byte in $C8-$EF).
    LAB = re.compile(r'^(BytecodeProgram_[0-9A-F]{2}|StandaloneBytecodeProgram(?:_[0-9])?'
                     r'|DeathSequenceVramStream|LevelOamTemplateTable):\s*$')
    DEC = re.compile(r'^\s*\.byte\s+\$([0-9A-F]{2})')
    while i < N:
        ln = lines[i]
        m = LAB.match(ln)
        if not m:
            out.append(ln)
            i += 1
            continue
        label = m.group(1)
        # collect the contiguous .byte run belonging to this program
        j = i + 1
        body = []
        while j < N and DEC.match(lines[j]):
            for t in lines[j].split(';')[0].split(','):
                t = t.strip()
                if t.lower().startswith('.byte'):
                    t = t[5:]
                t = t.strip().replace('$', '')
                if t:
                    body.append(int(t, 16))
            j += 1
        # program start = PRG offset from the first data line's annotation
        m2 = re.search(r'PRG:([0-9A-F]{4})', lines[i + 1])
        a = int(m2.group(1), 16)
        b = a + len(body)
        dec, nops, nlit = listing(a, b)
        idx = [k for k, t in enumerate(TABLE) if t == a]
        hdr = [f'; --- bytecode listing: {len(body)} bytes, {nops} opcodes, '
               f'{nlit} literal operands ---']
        if idx:
            hdr.append(f';     reached by BytecodeProgramPointerTable entry/entries {idx}')
        hdr.append(f';     handler for each opcode: BytecodeHandlerTable at $9237 '
                   f'(token - $C8); interpreter RunBytecodeInterpreter at $9AE8')
        if nops == 0:
            hdr.append(';     NOTE: this program contains no byte in $C8-$EF, so the')
            hdr.append(';           interpreter takes the literal path for every byte and the')
            hdr.append(';           whole block is operand data, not an instruction stream.')
        out.append(ln)
        out.extend(hdr)
        out.extend(dec)
        stats[label] = (nops, nlit)
        out.extend(lines[i + 1:j])
        i = j
    open(PRG_S, 'w').write('\n'.join(out))
    to = sum(v[0] for v in stats.values())
    tl = sum(v[1] for v in stats.values())
    print(f'decoded {len(stats)} programs: {to} opcodes, {tl} literal operands')
    for k in sorted(stats):
        print(f'   {k:<28} ops={stats[k][0]:>3} lits={stats[k][1]:>3}')


if __name__ == '__main__':
    main()