jwp-mission-dot.py

#!/usr/bin/env python3
# vim: set expandtab tabstop=4 shiftwidth=4:

# Borderlands 3 Data Processing Scripts
# Copyright (C) 2022 CJ Kucera
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of the development team nor the
#       names of its contributors may be used to endorse or promote products
#       derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL CJ KUCERA BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import os
import sys
import json
import html
import argparse
import textwrap
import subprocess

# Default external commands we call
cmd_serialize = '/home/pez/bin/ueserialize'
cmd_dot = '/usr/bin/dot'
cmd_view = '/usr/bin/feh'

class Element:

    def __init__(self, data):
        self.data = data

        # dot defaults
        self.dot_shape = 'ellipse'
        self.dot_color = 'white'
        self.dot_styles = []

        # From data
        self.obj_type = data['export_type']
        self.export = data['_jwp_export_idx']
        self.name = data['_jwp_object_name']
        self.dot_name = f'export_{self.export}'

    def dot_label(self):
        lines = [self.name]
        lines.extend(self._dot_label())
        lines.append(f'<i>(export {self.export})</i>')
        return '<br/>'.join(lines)

    def _dot_label(self):
        return []

    def dot_node(self):
        attrs = [
                'label=<{}>'.format(self.dot_label()),
                'shape={}'.format(self.dot_shape),
                'style="{}"'.format(','.join(['filled', *self.dot_styles])),
                'fillcolor={}'.format(self.dot_color),
                ]
        return '{} [{}];'.format(
                self.dot_name,
                ' '.join(attrs),
                )

class Objective(Element):


    def __init__(self, data):
        super().__init__(data)
        self.dot_color = 'aquamarine'
        self.count = data['ObjectiveCount']
        if 'FormattedProgressMessage' in data \
                and 'FormatText' in data['FormattedProgressMessage']:
            self.progress_message = data['FormattedProgressMessage']['FormatText']['string']
        else:
            self.progress_message = None
        if 'bInvisible' in data:
            self.invisible = data['bInvisible']
            if self.invisible:
                self.dot_styles.append('dashed')
        else:
            self.invisible = False
        if 'bAutoClearWhenNoLongerDormant' in data:
            self.auto_clear = data['bAutoClearWhenNoLongerDormant']
        else:
            self.auto_clear = False
        self.guid = data['ObjectiveGuid']
        self.obj_index = None

    def _dot_label(self):
        lines = []
        if self.invisible:
            lines.append('<i>(invisible)</i>')
        if self.obj_index is not None:
            lines.append(f'Objective idx {self.obj_index}')
        if self.progress_message:
            lines.append(html.escape(self.progress_message, quote=True))
        lines.append(f'<font face="Courier">{self.guid}</font>')
        return lines

class ObjectiveSet(Element):

    def __init__(self, data):
        super().__init__(data)
        self.dot_shape = 'box'
        self.dot_color = 'cadetblue1'
        self.objectives = []
        for objective in data['Objectives']:
            self.objectives.append(objective['export'])
        if 'NextSet' in data and data['NextSet']['export'] != 0:
            self.next_set = data['NextSet']['export']
        else:
            self.next_set = None
        if 'bCanCompleteMission' in data:
            self.can_complete = data['bCanCompleteMission']
            if self.can_complete:
                self.dot_shape = 'octagon'
                self.dot_color = 'coral'
        else:
            self.can_complete = False
        self.guid = data['ObjectiveSetGuid']
        # Not reading in ObjOrderPos
        self.obj_set_index = None

    def _finish(self, objectives, objective_sets):
        for idx, objective_idx in list(enumerate(self.objectives)):
            self.objectives[idx] = objectives[objective_idx]
            self.objectives[idx].obj_index = idx
        if self.next_set is not None:
            self.next_set = objective_sets[self.next_set]

    def dot_links(self):
        to_ret = []
        if self.next_set is not None:
            to_ret.append('{} -> {} [weight=3];'.format(
                self.dot_name,
                self.next_set.dot_name,
                ))
        for objective in self.objectives:
            to_ret.append('{} -> {};'.format(
                self.dot_name,
                objective.dot_name,
                ))
        return to_ret

    def _dot_label(self):
        lines = []
        if self.obj_set_index is not None:
            lines.append(f'Set idx {self.obj_set_index}')
        lines.append(f'<font face="Courier">{self.guid}</font>')
        return lines

class Phase(Element):

    def __init__(self, data):
        super().__init__(data)
        self.dot_shape = 'house'
        self.dot_color = 'darkgoldenrod1'
        self.objective_sets = []
        for objective_set in data['ObjectiveSets']:
            self.objective_sets.append(objective_set['export'])
        self.phase_index = data['PhaseIndex']

    def _finish(self, objective_sets):
        for idx, objective_set_idx in list(enumerate(self.objective_sets)):
            self.objective_sets[idx] = objective_sets[objective_set_idx]
            self.objective_sets[idx].obj_set_index = idx

    def dot_links(self):
        to_ret = []
        for objective_set in self.objective_sets:
            to_ret.append('{} -> {};'.format(
                self.dot_name,
                objective_set.dot_name,
                ))
        return to_ret

class Mission:

    def __init__(self, json_path):
        self.json_path = json_path
        with open(self.json_path) as df:
            self.data = json.load(df)
        self.objectives = {}
        self.objective_sets = {}
        self.phases = {}
        self._parse()

    def _parse(self):
        for export in self.data:
            match export['export_type']:
                case 'MissionObjective':
                    objective = Objective(export)
                    self.objectives[objective.export] = objective
                case 'MissionObjectiveSet':
                    objective_set = ObjectiveSet(export)
                    self.objective_sets[objective_set.export] = objective_set
                case 'MissionPhase':
                    phase = Phase(export)
                    self.phases[phase.export] = phase

        for objective_set in self.objective_sets.values():
            objective_set._finish(self.objectives, self.objective_sets)
        for phase in self.phases.values():
            phase._finish(self.objective_sets)

    def to_dot(self):
        to_ret = []

        to_ret.append('// Phases')
        for phase in self.phases.values():
            to_ret.append(phase.dot_node())
        to_ret.append('')

        to_ret.append('// Objective Sets')
        for objective_set in self.objective_sets.values():
            to_ret.append(objective_set.dot_node())
        to_ret.append('')

        to_ret.append('// Objectives')
        for objective in self.objectives.values():
            to_ret.append(objective.dot_node())
        to_ret.append('')

        to_ret.append('// Phase Links')
        for phase in self.phases.values():
            to_ret.extend(phase.dot_links())
        to_ret.append('')

        to_ret.append('// Objective Set Links')
        for objective_set in self.objective_sets.values():
            to_ret.extend(objective_set.dot_links())
        to_ret.append('')

        return "\n".join(to_ret)

def main():

    global dot_output
    global cmd_serialize
    global cmd_dot
    global cmd_view

    # Arguments
    parser = argparse.ArgumentParser(
            description='Generate graphviz dot graphs of mission trees',
            )

    parser.add_argument('-r', '--render',
            default='svg',
            choices={'svg', 'png', 'jpg', 'gif', 'none'},
            help='Graph render type',
            )

    parser.add_argument('-s', '--serialize',
            type=str,
            default=cmd_serialize,
            help='Command to use for object serialization',
            )

    parser.add_argument('-d', '--dot',
            type=str,
            default=cmd_dot,
            help='Path to graphviz dot executable',
            )

    parser.add_argument('-v', '--view',
            type=str,
            default=cmd_view,
            help='Path to app to view rendered version (if --render is not `none`)',
            )

    parser.add_argument('filename',
            nargs=1,
            type=str,
            help='Mission filename to serialize and graph',
            )

    args = parser.parse_args()
    dot_output = args.render
    cmd_serialize = args.serialize
    cmd_dot = args.dot
    cmd_view = args.view
    filename = args.filename[0]

    # Grab the filename to process
    if filename.endswith('.'):
        filename = filename[:-1]
    if '.' in filename:
        filename_base, ext = filename.rsplit('.', 1)
        if ext not in {'json', 'uasset', 'uexp'}:
            raise RuntimeError('Unknown filename: {}'.format(filename))
        filename = filename_base

    # Serialize it (might be already serialized, but don't bother checking)
    subprocess.run([cmd_serialize, 'serialize', filename])

    # Make sure it worked
    json_path = '{}.json'.format(filename)
    if not os.path.exists(json_path):
        raise RuntimeError('Could not find {}'.format(json_path))

    # Parse the JSON
    mission = Mission(json_path)

    # Now generate a DOT graph
    dot_path = '{}.dot'.format(filename)
    with open(dot_path, 'wt') as odf:
        if '/' in filename:
            obj_name = filename.split('/')[-1]
        else:
            obj_name = filename
        print('digraph {} {{'.format(obj_name), file=odf)
        print('', file=odf)

        print('// Main Graph Label', file=odf)
        print('labelloc = "t";', file=odf)
        print('fontsize = 16;', file=odf)
        print('label = <{}>'.format(obj_name), file=odf)
        print('', file=odf)

        print(mission.to_dot(), file=odf)

        print('}', file=odf)

    # Now generate graphviz
    final_path = '{}.{}'.format(filename, dot_output)
    subprocess.run([cmd_dot, '-T{}'.format(dot_output), dot_path, '-o', final_path])

    # ... and display it, if it worked
    if os.path.exists(final_path):
        print('Wrote to {}!'.format(final_path))
        if dot_output in {'png', 'jpg', 'gif', 'svg'}:
            subprocess.run([cmd_view, final_path])
    else:
        print('ERROR: {} was not written'.format(final_path))

if __name__ == '__main__':
    main()