dism_battle_ai.py

import collections
import os
import sys

PTRTYPE_SCRIPT = 0
PTRTYPE_BYTEARRAY = 1
PTRTYPE_HWORDARRAY = 2

if len(sys.argv) > 1:
    project_dir = sys.argv[1]
else:
    project_dir = os.path.join(os.environ.get('HOME', '.'), 'pokefirered')

names = [
    'AI_CheckBadMove',
    'AI_TryToFaint',
    'AI_CheckViability',
    'AI_SetupFirstTurn',
    'AI_Risky',
    'AI_PreferStrongestMove',
    'AI_PreferBatonPass',
    'AI_DoubleBattle',
    'AI_HPAware',
    'AI_Unknown',
    'AI_Ret',
    'AI_Ret',
    'AI_Ret',
    'AI_Ret',
    'AI_Ret',
    'AI_Ret',
    'AI_Ret',
    'AI_Ret',
    'AI_Ret',
    'AI_Ret',
    'AI_Ret',
    'AI_Ret',
    'AI_Ret',
    'AI_Ret',
    'AI_Ret',
    'AI_Ret',
    'AI_Ret',
    'AI_Ret',
    'AI_Ret',
    'AI_Roaming',
    'AI_Safari',
    'AI_FirstBattle'
]


macros = []
with open(os.path.join(project_dir, 'asm', 'macros', 'battle_ai_script.inc')) as fp:
    in_macro = False
    is_standalone_macro = False
    params = ()
    name = None
    for line in fp:
        if not in_macro:
            if line.startswith('\t.macro '):
                in_macro = True
                is_standalone_macro = False
                name = line.split()[1]
        else:
            if not is_standalone_macro and line.startswith(f'\t.byte 0x{len(macros):02x}'):
                is_standalone_macro = True
                continue
            if is_standalone_macro:
                if line.startswith('\t.byte \\'):
                    params += (1,)
                elif line.startswith('\t.2byte \\'):
                    params += (2,)
                elif line.startswith('\t.4byte \\'):
                    params += (4,)
            if line.startswith('\t.endm\n'):
                if is_standalone_macro:
                    macros.append((name, params))
                params = ()
                in_macro = False

# outfile = open(os.path.join(project_dir, 'tmp.s'), 'w')
with open(os.path.join(project_dir, 'data', 'battle_ai_scripts.s')) as fp:
    for line in fp:
        if line.startswith('\t.incbin "baserom.gba", 0x'):
            _, start, size = line.rstrip('\n').split(', ')
            start = int(start, 0)
            size = int(size, 0)
            break
    else:
        start = 0x1D9BF4
        size = 0x20FA

with open(os.path.join(project_dir, 'baserom.gba'), 'rb') as baserom:
    baserom.seek(start)
    data = baserom.read(size)

pointers = {}
n = 0
while True:
    if (4 * n + start) | 0x08000000 in pointers:
        break
    pointer = int.from_bytes(data[(n * 4):(n * 4 + 4)], 'little')
    print(f'\t.4byte {names[n]}')
    if pointer not in pointers:
        pointers[pointer] = (PTRTYPE_SCRIPT, names[n])
    n += 1

ptr = n * 4
last_ptrtype = PTRTYPE_SCRIPT
print_buffer = collections.defaultdict(list)
while ptr < size:
    addr = (ptr + start) | 0x08000000

    if addr not in pointers and data[ptr - 1] == macros.index(('end', ())):
        for i in range(4):
            if data[ptr + i] != 0 and addr + i in pointers:
                print_buffer[addr].append(f'\t.space {i}')
                ptr += i
                addr += i
                break

    ptrtype, label = pointers.get(addr, (last_ptrtype, None))
    if ptrtype == PTRTYPE_SCRIPT and data[ptr] < len(macros):
        macro, params = macros[data[ptr]]
        ptr += 1
        output = ()
        for p in params:
            value = int.from_bytes(data[ptr:ptr + p], 'little', signed=True)
            if p == 4 and (value & 0xFF000000) == 0x08000000:
                ptrtype2, label = pointers.get(value, (None, f'AI_sub_{value:X}'))
                if value not in pointers:
                    ptrtype2 = PTRTYPE_SCRIPT
                    if p == 4 and len(output) == 0:
                        if macro in ('if_in_bytes', 'if_not_in_bytes'):
                            label = f'AI_bytes_{value:X}'
                            ptrtype2 = PTRTYPE_BYTEARRAY
                        elif macro in ('if_in_words', 'if_not_in_words'):
                            label = f'AI_words_{value:X}'
                            ptrtype2 = PTRTYPE_HWORDARRAY
                    pointers[value] = (ptrtype2, label)
                value = label
            output += (value,)
            ptr += p
        if output:
            print_buffer[addr].append(f'\t{macro} ' + ', '.join(f'{xx}' for xx in output))
        else:
            print_buffer[addr].append(f'\t{macro}')
    elif ptrtype == PTRTYPE_BYTEARRAY:
        while data[ptr] != 255:
            print_buffer[addr].append(f'\t.byte {data[ptr]}')
            ptr += 1
        print_buffer[addr].append('\t.byte 0xFF')
        ptr += 1
    elif ptrtype == PTRTYPE_HWORDARRAY:
        print_buffer[addr].append('\t.align 1')
        while data[ptr] != 255 and data[ptr + 1] != 255:
            hword = int.from_bytes(data[ptr:ptr + 2], 'little')
            print_buffer[addr].append(f'\t.2byte 0x{hword:X}')
            ptr += 2
        print_buffer[addr].append('\t.2byte 0xFFFF')
        ptr += 2
    else:
        print_buffer[addr].append(f'\t.byte 0x{data[ptr]:02x}')
        ptr += 1

for addr in sorted(print_buffer):
    if addr in pointers:
        _, label = pointers[addr]
        print(f'\n{label}:: @ {addr:X}')
    for line in print_buffer[addr]:
        print(line)