wildencounters2json.py

# pokemap2json.py
# This script will process a pokeemerald project and generate the wild encounters JSON data for it.
# It creates:
#    src/data/wild_encounters.json
#
# In order to convert your project's existing data from src/data/wild_encounters.h:
# 1. Merge pret/master.
# 2. Run wildencounters2json.py on your project.
#    Example: "python wildencounters2json.py path/to/my/project"
# 3. Commit the src/data/wild_encounters.json file it creates.
# 4. Build your ROM (Make sure to run ./build_tools.sh, since there is a new JSON conversion tool in tools/jsonproc.)
#

import os
import sys
import re
import json
from collections import OrderedDict
import pprint

def get_header_label(line):
    line = line.strip()
    line = line[31:]
    return line[:line.find('[')]

def get_map_group(line):
    line = line.strip()
    if 'MAP_GROUP' in line:
        line = line[line.find('MAP_GROUP') + 10:]
        return line[:line.find(')')]
    else:
        line = line[line.find('=') + 1:]
        line = line.strip()
        return line[:line.find(',')]

def get_map_num(line):
    line = line.strip()
    if 'MAP_NUM' in line:
        line = line[line.find('MAP_NUM') + 8:]
        return line[:line.find(')')]
    else:
        line = line[line.find('=') + 1:]
        line = line.strip()
        return line[:line.find(',')]

def get_mons_info(line):
    line = line.strip()
    if 'NULL' in line:
        return 'NULL'

    match = re.match(r'.*&(.+),', line)
    if not match:
        print 'ERROR: could not find mons info label in line: ' + line
        sys.exit()

    return match.group(1)

def get_encounter_label(mons_info_label):
    match = re.match(r'(.*)_[^_]+', mons_info_label)
    if not match:
        print 'ERROR: could not determine base encounter label name from  mons info label: ' + mons_info_label
        sys.exit()

    return match.group(1)


def parse_header_tables(content):
    header_tables = []
    header_table = {}
    lines = content.splitlines();
    i = 0
    state = 'searching'
    while i < len(lines):
        line = lines[i]
        if state == 'searching':
            if line.startswith('const struct WildPokemonHeader'):
                header_table['label'] = get_header_label(line)
                header_table['entries'] = []
                state = 'find_next_of_struct'
        elif state == 'find_next_of_struct':
            if '.mapGroup' in line:
                entry = {}
                entry['map_group'] = get_map_group(line)
                entry['map_num'] = get_map_num(lines[i + 1])
                entry['land_mons_info'] = get_mons_info(lines[i + 2])
                entry['water_mons_info'] = get_mons_info(lines[i + 3])
                entry['rock_smash_mons_info'] = get_mons_info(lines[i + 4])
                entry['fishing_mons_info'] = get_mons_info(lines[i + 5])
                if entry['land_mons_info'] != 'NULL':
                    entry['label'] = get_encounter_label(entry['land_mons_info'])
                elif entry['water_mons_info'] != 'NULL':
                    entry['label'] = get_encounter_label(entry['water_mons_info'])
                elif entry['rock_smash_mons_info'] != 'NULL':
                    entry['label'] = get_encounter_label(entry['rock_smash_mons_info'])
                elif entry['fishing_mons_info'] != 'NULL':
                    entry['label'] = get_encounter_label(entry['fishing_mons_info'])
                header_table['entries'].append(entry)
                i += 6
            elif ';' in line:
                header_tables.append(header_table)
                header_table = {}
                state = 'searching'
        i += 1

    return header_tables

def parse_mon_info(entry, info_label, content):
    if info_label == 'NULL':
        return

    for line in content.splitlines():
        line = line.strip()
        if info_label in line:
            match = re.match(r'.*{(.+), (.+)}', line)
            if not match:
                print 'ERROR: could not parse WildPokemonInfo info struct in line: ' + line
                sys.exit()

            encounter_rate = int(match.group(1))
            label = match.group(2)
            entry['mons_data'][info_label] = {}
            entry['mons_data'][info_label]['encounter_rate'] = encounter_rate
            entry['mons_data'][info_label]['label'] = label
            break

    if info_label not in entry['mons_data']:
        print 'ERROR: could not find mon info label' + info_label
        sys.exit()

    wild_mons_label = 'const struct WildPokemon ' + entry['mons_data'][info_label]['label']
    lines = content.splitlines();
    state = 'searching'
    i = 0
    mons = []
    while True:
        line = lines[i]
        if state == 'searching':
            if wild_mons_label in line:
                state = 'processing'
        elif state == 'processing':
            if ';' in line:
                break
            if not line.startswith('{'):
                match = re.match(r'.*{(.+), (.+), (.+)},', line)
                if not match:
                    print 'ERROR: could not parse WildPokemon data struct in line: ' + line
                    sys.exit()

                min_level = int(match.group(1).strip())
                max_level = int(match.group(2).strip())
                species = match.group(3).strip()
                mons.append({'min_level': min_level, 'max_level': max_level, 'species': species})
        i += 1

    entry['mons_data'][info_label]['mons'] = mons


def parse_mon_infos(header_tables, content):
    for header_table in header_tables:
        for entry in header_table['entries']:
            entry['mons_data'] = {}
            parse_mon_info(entry, entry['land_mons_info'], content)
            parse_mon_info(entry, entry['water_mons_info'], content)
            parse_mon_info(entry, entry['rock_smash_mons_info'], content)
            parse_mon_info(entry, entry['fishing_mons_info'], content)

def build_final_mons_info(info):
    data = OrderedDict()
    data['encounter_rate'] = info['encounter_rate']
    data['mons'] = []
    for mon in info['mons']:
        data['mons'].append(OrderedDict([
            ('min_level', mon['min_level']),
            ('max_level', mon['max_level']),
            ('species', mon['species'])]))

    return data

def prepare_json(header_tables):
    data = {}
    data['wild_encounter_groups'] = []
    for header_table in header_tables:
        encounter_group = OrderedDict()
        encounter_group['label'] = header_table['label']
        encounter_group['for_maps'] = header_table['label'] == 'gWildMonHeaders'
        encounter_group['encounters'] = []
        for encounter in header_table['entries'][:-1]:
            entry = OrderedDict()
            if encounter_group['for_maps']:
                entry['map'] = 'MAP_' + encounter['map_group']
                entry['base_label'] = encounter['label']
            else:
                entry['base_label'] = header_table['label'].replace("WildMonHeaders", "") + '_' + encounter['map_num']

            if encounter['land_mons_info'] != 'NULL':
                entry['land_mons'] = build_final_mons_info(encounter['mons_data'][encounter['land_mons_info']])
            if encounter['water_mons_info'] != 'NULL':
                entry['water_mons'] = build_final_mons_info(encounter['mons_data'][encounter['water_mons_info']])
            if encounter['rock_smash_mons_info'] != 'NULL':
                entry['rock_smash_mons'] = build_final_mons_info(encounter['mons_data'][encounter['rock_smash_mons_info']])
            if encounter['fishing_mons_info'] != 'NULL':
                entry['fishing_mons'] = build_final_mons_info(encounter['mons_data'][encounter['fishing_mons_info']])

            encounter_group['encounters'].append(entry)

        data['wild_encounter_groups'].append(encounter_group)

    return data        


def main():
    if len(sys.argv) != 2:
        print 'USAGE: python wildencounters2json.py <project_path>'
        sys.exit()

    project_root = sys.argv[1]
    if not os.path.exists(project_root):
        print 'ERROR: Project directory "%s" does not exist' % (project_root)
        sys.exit()

    print 'Converting wild_encounters data to json for project "%s"' % (project_root)
    filepath = os.path.join(project_root, 'src/data/wild_encounters.h')
    if not os.path.exists(filepath):
        print 'ERROR: wild encounters file "%s" does not exist' % (filepath)
        sys.exit()

    with open(filepath) as f:
        content = f.read()

    header_tables = parse_header_tables(content)
    parse_mon_infos(header_tables, content)
    data = prepare_json(header_tables)
    
    filepath = os.path.join(project_root, 'src/data/wild_encounters.json')
    with open(filepath, 'w') as f:
        json.dump(data, f, indent=2, separators=(',', ': '))
    print 'Generated "%s"' % (filepath)


if __name__ == '__main__':
    main()