4etools/battleman.py

#!/usr/bin/python
#
# battleman.py - RPG Battle Manager
#
# A table-top RPG battle flow manager
# Tuned pretty specifically to D&D 4e for now... need a templatized system
# to do anything fancier... may develop that at some point.
#
# future features:
#  * an option for passing in multiple files that contain combatant definitions

import sys
sys.path.append('lib/')

import cPickle as pickle
import argparse
import cmd
import os.path
import battle
from battle import CombatGroup
from battle import Combatant
import easyinput


def main():
    btl = battle.Battle()

    ### This is the pickling jar
    battle_pickle = None
    bp_io_failed = False
    BP_FILE = os.path.expanduser('~/.config/4etools/battleman/battle.pickle')
    ###

    # Make sure config directory exists
    if not os.path.exists(os.path.dirname(BP_FILE)):
        os.makedirs(os.path.dirname(BP_FILE))

    # Get command-line args
    settings = parse_args()


    print "Welcome to 4e Battle Manager.\n"
    
    # Resume battle if needed
    # fixme: critical: indexes get reused when restoring! Maybe switch to shelve?
    if settings.resume:
        try:
            with open(BP_FILE, 'r') as f:
                btl = pickle.load(f)
                battle_pickle = pickle.dumps(btl)
        except:
            print "Error: Couldn't resume. Quitting to preserve our pickle."
            sys.exit(1)

    else:
        # hard-coding test cases for now.
        # fixme: Eventually, use a state-saving text file that's easy to edit, or at least copy...
        adele = Combatant("Adele", hp=26, pc=True, surges=8, sw=1)
        adele_dict = {adele.index: adele}
        aristaire = Combatant("Aristaire", hp=20, pc=True, surges=6, sw=1)
        aristaire_dict = {aristaire.index: aristaire}

        foobolds = {}
        for i in range(5):
            c = Combatant("Foobold", hp=50)
            foobolds[c.index] = c

        barglins = {}
        for i in range(2):
            c = Combatant("Barglin", hp=50)
            barglins[c.index] = c

        orcs = {}
        for i in range(2):
            c = Combatant("Orc", hp=50)
            orcs[c.index] = c

        btl.add_group(CombatGroup("Adele", adele_dict, 2))
        btl.add_group(CombatGroup("Aristaire", aristaire_dict, 0))
        btl.add_group(CombatGroup("Foobolds", foobolds, 20))
        btl.add_group(CombatGroup("Barglins", barglins, 3))
        btl.add_group(CombatGroup("Orcs of Baz", orcs, 1))

    print btl

    cmd_parser = CommandParser(btl, battle_pickle, BP_FILE)
    cmd_parser.cmdloop()



class CommandParser(cmd.Cmd):
    """Parse the commands from the command-line."""

    def __init__(self, btl, battle_pickle, BP_FILE):
        cmd.Cmd.__init__(self)

        self.btl = btl
        self.battle_pickle = battle_pickle
        self.BP_FILE = BP_FILE
        self.doc_header = 'Available commands (type help <command> for more help)'
        self.prompt = '\n> '


    def postloop(self):
        # Re-pickle and write if changed after every query. It's cheap
        # and we only have to run at user-speed anyway
        old_bp = self.battle_pickle
        self.battle_pickle = pickle.dumps(self.btl)

        if old_bp != self.battle_pickle:
            try:
                with open(self.BP_FILE, 'w') as f:
                    f.write(self.battle_pickle)
            except Exception:
                if not self.bp_io_failed:
                    print("Warning: can't write the battle pickle. Resuming later will fail.")
                    self.bp_io_failed = True


    # This allows us to do partial command completion without <tab>,
    # as long as 
    def default(self, line):
        cmd, data, line = self.parseline(line)
        cmds = self.completenames(cmd)
        num_cmds = len(cmds)
        if num_cmds == 1:
            getattr(self, 'do_'+cmds[0])(data)
        elif num_cmds > 1:
            sys.stdout.write('Error: Ambiguous command: {}'.format(cmd))
        else:
            print 'Error: Unrecognized command {}'.format(cmd)


    # We are overriding do_help to avoid printing info about
    # undocumented commands
    def do_help(self, arg):
        if arg:
            Cmd.cmd.do_help(arg)
        else:
            # Everything from here to the end is lifted straight
            # out of Cmd.cmd.do_help()
            names = self.get_names()
            cmds_doc = []
            cmds_undoc = []
            help = {}
            for name in names:
                if name[:5] == 'help_':
                    help[name[5:]]=1
            names.sort()
            # There can be duplicates if routines overridden
            prevname = ''
            for name in names:
                if name[:3] == 'do_':
                    if name == prevname:
                        continue
                    prevname = name
                    cmd=name[3:]
                    if cmd in help:
                        cmds_doc.append(cmd)
                        del help[cmd]
                    elif getattr(self, name).__doc__:
                        cmds_doc.append(cmd)
                    else:
                        cmds_undoc.append(cmd)
            self.stdout.write("%s\n"%str(self.doc_leader))
            self.print_topics(self.doc_header,   cmds_doc,   15,80)
            self.print_topics(self.misc_header,  help.keys(),15,80)
            # self.print_topics(self.undoc_header, cmds_undoc, 15,80)


    def do_add(self, line):
        """add [N]
        Add the specified number of groups"""

        data = parse_data(line)

        if len(data) >= 1:
            num_groups = int(data[0])
        else:
            num_groups = easyinput.input_int('number of groups')

        for i in range(1, num_groups+1):
            print "Adding group {}".format(i)
            self.btl.add_group(CombatGroup.from_input())


    def do_begin(self, line):
        """begin
        Begins the battle. Rolls initiative for NPCs and prompts for PCs"""

        self.btl.begin()


    def do_print(self, line):
        """print [index]
        Print detailed info for combatant with index, or combatant or group with initiative"""

        data = parse_data(line)

        if len(data) >= 1:
            c = self.btl.get_combatant(int(data[0]))
            if not c:
                print 'Error: Invalid combatant index.'
            else:
                print c.format_full_info()
        else:
            print self.btl.format_current_group()


    def do_list(self, line):
        """list
        Lists a summary of all of the combat groups and their members"""
    
        print self.btl.format_combatants()


    def do_damage(self, line):
        """damage [index] [amount]
        Deals damage to the specified combatant"""

        data = parse_data(line)

        c = battle.do_combatant_select(self.btl, data)
        if not c:
            return

        if len(data) >= 2:
            amount = int(data[1])
        else:
            amount = easyinput.input_int('damage')
            
        c.damage(amount)

    
    def do_heal(self, line):
        """heal [index] [amount]
        Heal hit points for the specified combatant"""
        
        data = parse_data(line)

        c = battle.do_combatant_select(self.btl, data)
        if not c:
            return

        if len(data) >= 2:
            amount = int(data[1])
        else:
            amount = easyinput.input_int('amount')
            
        c.heal(amount)


    def do_temp(self, line):
        """temp [index] [amount]
        Add temporary hit points to the specified combatant"""

        data = parse_data(line)

        c = battle.do_combatant_select(self.btl, data)
        if not c:
            return
        
        if len(data) >= 2:
            amount = int(data[1])
        else:
            amount = easyinput.input_int('amount')

        c.add_temp_hp(amount)


    def do_rmtemp(self, line):
        """rmtemp [index] [amount]
        Remove temporary hit points from the specified combatant"""

        do_stub()


    def do_surge(self, line):
        """surge [index] [heal]
        Combatant with index uses a healing surge. If heal is 0, don't heal the combatant"""

        data = parse_data(line)

        c = battle.do_combatant_select(self.btl, data)
        if not c:
            return

        heal = True
        if len(data) >= 2 and data[1] == '0':
            heal = False
        c.use_surge(heal)


    def do_wind(self, line):
        """wind [index]
        Use Second Wind for combatant"""

        data = parse_data(line)

        c = battle.do_combatant_select(self.btl, data)
        if not c:
            return
        c.use_second_wind()


    def do_cond(self, line):
        """cond [index] [name] [type] [duration] [start|end]
        Add a temporary condition to a combatant, optionally specifying the condition name, type (s or t), duration and what phase of the combatant's turn it expires on"""

        data = parse_data(line)

        duration = None
        end_type = 'e'

        c = battle.do_combatant_select(self.btl, data)
        if not c:
            return

        name = easyinput.do_data_input_str(data, 1, 'condition name')
        ctype = easyinput.do_data_input_str(data, 2, 'condition type', default='s', show_default=True)

        if ctype == 't':
            duration = easyinput.do_data_input_int(data, 3, 'duration')
            end_type = easyinput.do_data_input_str(data, 4, '(s)tart|(e)nd', default='e', show_default=True)

        c.add_condition(name, ctype, duration, end_type)


    def do_rmcond(self, line):
        """rmcond [index] [condition_index]
        Remove a condition from a combatant early."""

        data = parse_data(line)

        c = battle.do_combatant_select(self.btl, data)
        if not c:
            return

        if len(data) >= 2:
            index = int(data[1])
        else:
            cond = c.choose_condition()
            index = None
            if cond:
                index = cond['index']

        if index != None:
            c.remove_condition(index)


    def do_recharge(self, line):
        """recharge [index] [recharge_index]
        Use a rechargable power"""

        data = parse_data(line)

        c = battle.do_combatant_select(self.btl, data)
        if not c:
            return

        if len(data) >= 2:
            index = int(data[1])
        else:
            r = c.choose_recharge_power()
            index = None
            if r:
                index = r['index']

        if index != None:
            c.use_recharge_power(index)


    def do_wait(self, line):
        """wait
        Removes the specified combatant from the initiative roster and add them to the wait list."""

        data = parse_data(line)

        c = battle.do_combatant_select(self.btl, data)
        if not c:
            return

        self.btl.wait(c.index)


    def do_unwait(self, line):
        """unwait
        Removes the specified combatant from the wait list and adds them back into the initiative roster."""

        data = parse_data(line)

        c = battle.do_combatant_select(self.btl, data)
        if not c:
            return

        self.btl.unwait(c.index)


    def do_next(self, line):
        """next
        Steps to the next combatant in initiative order. This handles saving throws, effects that end at beginning and ends of turns, and round incrementing."""
        self.btl.next_combatant()


    def do_sync(self, line):
        """sync
        This function is still a stub"""
        
        # Since the postloop pickles & writes, we don't actually need to
        # do a damn thing here - just let it be an 'empty' command.
        pass


    def do_EOF(self, line):
        self.do_quit(line)


    def do_quit(self, line):
        """quit
        Exits the program. If a battle is in progress, it is temporarily saved and can be resumed by running the program with --resume next time."""

        return True



def parse_data(line):
    data = line.split(' ')
    if data == ['']:
        data = []
    return data


def do_stub():
    print "Sorry, this is a stub function"


def parse_args():
    parser = argparse.ArgumentParser(description='Command-line interface to manage battle data for D&D 4e', formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument('--resume', '-r', action='store_true', help='Resume the battle from the last run of the program')
    return parser.parse_args()


if __name__ == '__main__':
    main()