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:
#  * keep a pickled file or a shelve db of the current state, and add a --resume option
#    for resuming a battle later
#  * an option for passing in multiple files that contain combatant definitions
#  * down combatants go into a separate list
#  * Ability for commands to accept partial data, and continue with selection from there


from dice import Dice
import sys


class CombatGroup():
    
    # What we're mostly getting here is a definition of the *members*
    # of the group... then we build them all and stick them in
    # the group
    @classmethod
    def from_input(cls):
        name = input_str("Name: ")
        hp = input_int('hp')
        init_mod = input_int('init mod', 0)
        ap = input_int('action points', 0)
        surges = input_int('healing surges', 0)

        recharges = []
        recharge = '-1'
        while True:
            recharge = input_str("recharge: ").split(',')
            if recharge == ['']:
                break
            else:
                recharges.append(recharge)

        count = input_int('count', 1)

        # Now make the combatants...
        members = []
        for i in range(count):
            members.append(Combatant(name, hp, pc=False, surges=surges, ap=ap, sw=0, recharges=recharges))

        if count > 1:
            name = name + 's'

        return CombatGroup(name, members, init_mod)


    def __init__(self, name, members, init_mod=0):
        self.name = name
        self.members = members
        self.init_mod = init_mod
        self.init = 0


    def roll_init(self):
        d = Dice.from_str('1d20+{}'.format(self.init_mod))
        self.set_init(d.roll()['total'])


    def set_init(self, init):
        self.init = init


    def add_member(self, c):
        self.members.append(c)


    def is_solo_group(self):
        return len(self.members) == 1


    def begin_turn(self):
        msg = None
        if self.is_solo_group():
            msg = '{} has initiative.'.format(self.name)
        else:
            msg = '{} have initiative.'.format(self.name)

        print msg


    def end_turn(self):
        for c in self.members:
            c.end_turn()



class Combatant():
    next_index = 0

    def __init__(self, name, hp, pc=False, init_mod=0, surges=0, ap=0, sw=0, recharges=[]):
        self.name = name
        self.max_hp = hp
        self.hp = self.max_hp
        self.temp_hp = 0
        self.pc = pc
        self.surges = surges
        self.ap = ap
        self.sw = sw
        self.recharges = []
        self.conditions = {}
        self.index = Combatant.next_index
        self.next_condition_index = 0
        Combatant.next_index += 1


    def __str__(self):
        return "{}: {} ({})".format(self.index, self.name, self.format_health_summary())


    # cond_type can be 's' or 't', for 'save' or 'timed'. If it is 't', condition expires at the end of the players turn
    # 'duration' rounds from now
    def add_condition(self, name, cond_type, duration=None):
        condition = {}
        condition['name'] = name
        condition['cond_type'] = cond_type
        condition['duration'] = duration
        condition['index'] = self.next_condition_index
        if cond_type == 'timed' and duration == None:
            print('Error: specified a timed condition with no duration.')
            return
        self.conditions[self.next_condition_index] = condition
        self.next_condition_index += 1

        
    # Removes a condition, prints a message about it, and returns a copy of the condition
    def remove_condition(self, index):
        if index not in self.conditions:
            print "Error: invalid condition index."
            return None

        c = self.conditions.pop(index)
        print('{} is no longer affected by {}.'.format(self, c['name']))            

        return c


    def choose_condition(self):
        print self.format_condition_summary()
        index = input_int('choice')

        if index not in self.conditions:
            print 'Error: {} is not a valid index'.format(index)
            return self.choose_condition()

        return self.conditions[index]


    def tick_conditions(self):
        for c in self.conditions.values():
            if c['cond_type'] == 't':
                c['duration'] -= 1


    def end_turn(self):
        for (index, c) in self.conditions.items():
            if c['cond_type'] == 's':
                r = None
                if self.pc:
                    print('{}: save against {}.'.format(self, c['name']))
                    r = input_int('saving throw')
                else:
                    save_die = Dice.from_str('1d20')
                    r = save_die.roll()['total']

                if r >= 10:
                    self.remove_condition(index)
                else:
                    print('{} failed a save against {}.'.format(self, c['name']))

            elif c['cond_type'] == 't':
                c['duration'] -= 1
                if c['duration'] < 0:
                    self.remove_condition(index)
                else:
                    print('{} is still affected by {} ({} round{} left).'.format(self, c['name'], c['duration']+1, 's'[c['duration']==0:]))

        # fixme: still need to add recharges


    def damage(self, amount):
        was_bloodied = self.is_bloodied()

        if self.temp_hp > 0:
            self.temp_hp -= amount
            if self.temp_hp < 0:
                amount = abs(self.temp_hp)
                self.temp_hp = 0

        self.hp -= amount

        print '{} took {} points of damage.'.format(self, amount)

        if self.is_down():
            print('{} is down!'.format(self))
        elif self.is_bloodied() and not was_bloodied:
            print('{} is bloodied!'.format(self))


    def heal(self, amount):
        was_down = self.is_down()
        was_bloodied = self.is_bloodied()

        if self.hp < 0:
            self.hp = 0

        amount_healed = amount
        if self.hp + amount_healed > self.max_hp:
            amount_healed = self.max_hp - self.hp

        self.hp += amount_healed

        if was_down:
            print('{} regains consciousness.'.format(self))

        print('{} regained {} hit points.'.format(self, amount_healed))

        if was_bloodied and not self.is_bloodied():
            print('{} is no longer bloodied.'.format(self))
        elif was_bloodied and self.is_bloodied():
            print('{} is still bloodied.'.format(self))


    def add_temp_hp(self, amount):
        self.temp_hp += amount

        print '{} gained {} temporary hit points.'.format(self, amount)


    def use_surge(self, heal=True):
        if self.surges <= 0:
            print '{} has no healing surges.'.format(self)
            return
        
        self.surges -= 1
        noheal = ''
        if not heal:
            noheal = " (but didn't regain hit points)"
        print '{} spent a healing surge{}.'.format(self, noheal)

        if heal:
            self.heal(self.max_hp / 4)


    def use_second_wind(self):
        if self.sw <= 0:
            print "{} doesn't have a second wind.".format(self)
            return
        
        self.surges -= 1

        print '{} got their second wind!'.format(self)

        # Now the actual effects of the SW
        self.use_surge()
        self.add_condition('Second Wind (+2 all def)', 't', 1)


    def is_bloodied(self):
        return self.hp <= self.max_hp / 2


    def is_down(self):
        return self.hp <= 0


    def format_health_summary(self):
        bloodied = ''
        temp_info = ''
        if self.is_bloodied():
            bloodied = ', bloodied'
        if len(self.conditions):
            bloodied = bloodied + ', '

        if self.temp_hp > 0:
            temp_info = ', {} temp hp'.format(self.temp_hp)

        return '{} hp{}{}{}'.format(self.hp, temp_info, bloodied, ', '.join([x['name'] for x in self.conditions.values()]))


    def format_condition_summary(self):
        summary = ''
        for (index, c) in self.conditions.items():
            type_string = ''
            if c['cond_type'] == 's':
                type_string = 'Save Ends'
            elif c['cond_type'] == 't':
                type_string = '{} Round{}'.format(c['duration'], 's'[ c['duration']==1: ] )
            summary = summary + '{}: {} ({})\n'.format(index, c['name'], type_string)

        return summary.rstrip()



# data about the battle - includes combatant list, etc
class Battle():
    def __init__(self):
        self.combatant_hash = {}
        self.groups = []
        self.current = None
        self.round = None


    def __str__(self):
        ret = ''
        if self.is_started():
            ret = 'Battle underway, currently on round {}\n\n'.format(self.round)
        else:
            ret = 'Battle not yet started\n\n'

        ret = ret + 'Combatants\n==========\n'
        ret = ret + self.format_combatants()

        return ret
        

    def is_started(self):
        return self.current != None


    def add_group(self, group):
        self.groups.append(group)
        for c in group.members:
            self.combatant_hash[c.index] = c


    def get_current_group(self):
        if self.current != -1:
            return self.groups[self.current]
        else:
            return None


    def get_combatant(self, index):
        if index in self.combatant_hash:
            return self.combatant_hash[index]
        else:
            return None


    def choose_combatant(self):
        print self.format_combatants()
        index = input_int('choose combatant')
        if index not in self.combatant_hash:
            print 'Error: {} is not a valid index'.format(index)
            return self.choose_combatant()
        return self.combatant_hash[index]


    def begin(self):
        if self.is_started():
            print("Error: battle is already running")

        for g in self.groups:
            if g.is_solo_group() and g.members[0].pc:
                g.set_init(input_int('Initiative for {}'.format(g.name)))
            else:
                g.roll_init()

        self.groups.sort(reverse=True, key=lambda group: group.init)
        self.current = 0

        print '\nInitiative Roster:\n'
        for g in self.groups:
            print '{} ({})'.format(g.name, g.init)
        print ''

        self.next_round()
        self.get_current_group().begin_turn()


    # Returns a formatted string with all of the combatants
    def format_combatants(self):
        ret = ''

        for g in self.groups:
            if g.is_solo_group():
                ret = ret + '{}\n'.format(g.members[0])
            else:
                ret = ret + '{}:\n'.format(g.name)
                for c in g.members:
                    ret = ret + '  {}\n'.format(c)
        
        return ret.rstrip()


    # Returns a formatted string with just the current group
    def format_current_group(self):
        if self.validate_started():
            return self.validate_started()

        ret = ''

        g = self.groups[self.current]
        if g.is_solo_group():
            ret = ret + '{}\n'.format(g.members[0])
        else:
            ret = ret + '{}:\n'.format(g.name)
            for c in g.members:
                ret = ret + '  {}\n'.format(c)
        
        return ret.rstrip()


    def next_combatant(self):
        if self.validate_started():
            print self.validate_started()
            return

        g = self.get_current_group()
        g.end_turn()
        
        self.current += 1

        if self.current >= len(self.groups):
            self.current = 0
            self.next_round()

        g = self.get_current_group()
        g.begin_turn()


    def next_round(self):
        if self.round == None:
            self.round = 1
        else:
            self.round += 1

        # Decrement all timed conditions
        for c in self.combatant_hash.values():
            c.tick_conditions()

        print('Beginning round {}'.format(self.round))


    def validate_started(self):
        if not self.is_started():
            return 'Error: you can only run this command after starting the battle'
        return None


battle = Battle()

def main():
    # hard-coding test cases for now.
    # Eventually, use a state-saving text file that's easy to edit
    battle.add_group(CombatGroup("Adele", [Combatant("Adele", hp=26, pc=True, surges=8, sw=1)], 2))
    battle.add_group(CombatGroup("Aristaire", [Combatant("Aristaire", hp=20, pc=True, surges=6, sw=1)], 0))


    battle.add_group(CombatGroup("Foobolds", [Combatant("Foobold", hp=50), Combatant("Foobold", hp=50), Combatant("Foobold", hp=50), Combatant("Foobold", hp=50), Combatant("Foobold", hp=50)], 20))
    battle.add_group(CombatGroup("Barglins", [Combatant("Barglin", hp=1), Combatant("Barglin", hp=1)], 3))
    battle.add_group(CombatGroup("Orcs of Baz", [Combatant("Orc", hp=32), Combatant("Orc", hp=32)], 1))

    # ngroups = input_int('Number of enemy groups:')
    # for i in range(1, ngroups+1):
    #     print("Adding enemy group {}".format(i))
    #     battle.add_group(CombatGroup.from_input())

    print "Welcome to 4e Battle Manager.\n"
    print battle

    while True:
        do_prompt()


def do_prompt():
    print('')
    (comm, rdata) = input_str('', default='?', show_default=False, prompt_str='>').partition(' ')[::2]
    data = rdata.split(' ')

    # To simplify our checks later
    if data == ['']:
        data = None

    if comm == '?':
        do_help()
    elif comm == 'a':
        print('Sorry, this is still a stub function.')
    elif comm == 'p':
        print battle.format_current_group()
    elif comm == 'l':
        print battle.format_combatants()
    elif comm == 'b':
        battle.begin()
    elif comm == 'd':
        do_damage(data)
    elif comm == 'h':
        do_heal(data)
    elif comm == 't':
        do_add_temp_hp(data)
    elif comm == 's':
        do_surge(data)
    elif comm == 'so':
        do_surge(data, heal=False)
    elif comm == 'sw':
        do_second_wind(data)
    elif comm == 'c':
        do_add_condition(data)
    elif comm == 'r':
        do_remove_condition(data)
    elif comm == 'n':
        battle.next_combatant()
    elif comm == 'w':
        print('Sorry, this is still a stub function.')
    elif comm == 'q':
        sys.exit(0)


def do_help():
    print("""Possible commands:
?  - print this help menu (yay, you already figured that one out)
a  - add more combatants (works during battle) [stub]
b  - begin the battle
l  - list combatants
p  - print info for combatant/group with initiative
d  - deal damage to someone
h  - heal someone
t  - add temporary hit points
s  - use a healing surge
so - use a healing surge, but don't regain hit points
sw - use a second wind
c  - apply a condition
r  - remove a condition (this can also happen automatically)
n  - next (end the current combat group's turn)
w  - wait (remove a combatant from the initiative order and into a separate pool) [stub]
q  - quit""")


def do_damage(data):
    if (data):
        if len(data) != 2:
            print ('Error: wrong number of arguments.')

        c = battle.get_combatant(int(data[0]))
        if not c:
            print ('Error: Invalid combatant index.')
            return

        amount = int(data[1])
    else:
        c = battle.choose_combatant()
        amount = input_int('damage')

    c.damage(amount)


def do_heal(data):
    if (data):
        if len(data) != 2:
            print ('Error: wrong number of arguments')

        c = battle.get_combatant(int(data[0]))
        if not c:
            print ('Error: Invalid combatant index.')
            return

        amount = int(data[1])
    else:
        c = battle.choose_combatant()
        amount = input_int('amount')

    c.heal(amount)


def do_add_temp_hp(data):
    if (data):
        if len(data) != 2:
            print ('Error: wrong number of arguments')

        c = battle.get_combatant(int(data[0]))
        if not c:
            print ('Error: Invalid combatant index.')
            return

        amount = int(data[1])
    else:
        c = battle.choose_combatant()
        amount = input_int('amount')

    c.add_temp_hp(amount)


def do_surge(data, heal=True):
    if (data):
        if len(data) != 1:
            print ('Error: wrong number of arguments.')

        c = battle.get_combatant(int(data[0]))
        if not c:
            print ('Error: Invalid combatant index.')
            return

    else:
        c = battle.choose_combatant()

    c.use_surge(heal)


def do_second_wind(data):
    if (data):
        if len(data) != 1:
            print ('Error: wrong number of arguments.')

        c = battle.get_combatant(int(data[0]))
        if not c:
            print ('Error: Invalid combatant index.')
            return

    else:
        c = battle.choose_combatant()

    c.use_second_wind()


def do_add_condition(data):
    duration = None

    if data:
        if len(data) < 2 or (len(data) == 3 and data[2] == 't'):
            print ('Error: wrong number of arguments.')

        c = battle.get_combatant(int(data[0]))
        name = data[1]
        ctype = 's'
        if len(data) > 2:
            ctype = data[2]

        if ctype == 't':
            duration = int(data[3])

    else:
        c = battle.choose_combatant()
        name = input_str('condition name')
        ctype = input_str('condition type', default='s', show_default=True)
        duration=None
        if ctype == 't':
            duration = input_int('duration')

    c.add_condition(name, ctype, duration)


def do_remove_condition(data):
    if data:
        if len(data) != 2:
            print ('Error: wrong number of arguments.')        
            return

        c = battle.get_combatant(int(data[0]))
        if not c:
            print ('Error: Invalid combatant index.')
            return

        index = int(data[1])

    else:
        c = battle.choose_combatant()
        index = c.choose_condition()['index']

    c.remove_condition(index)


def input_str(prompt, default=None, show_default=False, prompt_str=':'):
    full_prompt = prompt
    if default != None and show_default:
        full_prompt = full_prompt + ' [{}]'.format(default)
    full_prompt = full_prompt + '{} '.format(prompt_str)
    
    data = raw_input(full_prompt)
    if not data:
        if default == None:
            print 'Error: you must provide a value!'
            return input_str(prompt, default, show_default, prompt_str)
        else:
            return default
    else:
        return data


def input_int(prompt, default=None, show_default=True, prompt_str=':'):
    return int(input_str(prompt, default, show_default))


if __name__ == '__main__':
    main()