Grab useful supplies off the ground

mosfet/jobs/grab_supplies.py

@@ -16,6 +16,15 @@ from mosfet.info import items
 from mosfet.info import mobs
 
 class GrabSuppliesStates:
+    DEFAULT_SUPPLIES = {
+        tuple(items.BED_IDS): (1, 1),
+        (items.CHEST_ID,): (1, 1),
+        tuple(items.FOOD_IDS): (1, 1),
+    }
+
+    def default_supplies(self):
+        return dict(self.DEFAULT_SUPPLIES)
+
     def idle(self):
         return None
 
@@ -30,31 +39,166 @@ class GrabSuppliesStates:
             self.state = self.cleanup
             return
 
+        self.checked_ground_items = set()
+        self.checked_ground_groups = set()
+        self.ground_stack_counts = {}
+        self.ground_entity_id = None
+        self.ground_position = None
         if self.supplies:
             self.state = self.check_supplies
         else:
             print('Aborting getting supplies, none specified')
             self.state = self.cleanup
 
+    def _has_nearby_ground_items(self, item_ids, maximum_stacks):
+        if (maximum_stacks and
+                self.ground_stack_counts.get(item_ids, 0) >= maximum_stacks):
+            return False
+        p = utils.pint(self.g.pos)
+        return any(
+            obj.entity_id not in self.checked_ground_items and
+            utils.phyp(p, (obj.x, obj.y, obj.z)) <= 16
+            for obj in self.g.world.find_objects(item_ids))
+
+    def has_nearby_ground_supply(self):
+        return any(
+            self._has_nearby_ground_items(item_ids, limits[1])
+            for item_ids, limits in self.supplies.items())
+
+    def _record_ground_pickup(self):
+        self.checked_ground_items.add(self.ground_entity_id)
+        self.ground_stack_counts[self.target_items] = (
+            self.ground_stack_counts.get(self.target_items, 0) + 1)
+        if (self.target_maximum_stacks and
+                self.ground_stack_counts[self.target_items] >=
+                self.target_maximum_stacks):
+            self.checked_ground_groups.add(self.target_items)
+        self.g.path = []
+        self.ground_entity_id = None
+        self.ground_position = None
+
     def check_supplies(self):
-        # check if we need to grab anything
-        for items, limits in self.supplies.items():
+        # Search every configured item group on the ground before going to a
+        # barrel; a missing bed must not hide a nearby chest or food drop.
+        missing = []
+        for item_ids, limits in self.supplies.items():
             minimum, maximum = limits
-            print('Checking items:', items)
-            num_items = self.g.game.count_items(items)
+            print('Checking items:', item_ids)
+            num_items = self.g.game.count_items(item_ids)
             print('Have:', num_items)
 
+            if (item_ids not in self.checked_ground_groups and
+                    self._has_nearby_ground_items(item_ids, maximum)):
+                print('Found matching supplies on the ground')
+                self.target_items = item_ids
+                self.target_minimum = minimum
+                self.target_maximum_stacks = maximum
+                self.state = self.find_ground_items
+                return
+
             if num_items >= minimum:
                 print('Have enough, skipping')
                 continue
 
-            print('Need at least one item')
+            missing.append((item_ids, minimum, maximum))
+            if item_ids not in self.checked_ground_groups:
+                self.target_items = item_ids
+                self.target_minimum = minimum
+                self.target_maximum_stacks = maximum
+                self.state = self.find_ground_items
+                return
+
+        if missing:
+            self.target_items, self.target_minimum, self.target_maximum_stacks = missing[0]
+            print('No more ground drops; checking barrels for',
+                  self.target_items)
             self.state = self.find_barrels
             return
 
         print('Aborting, dont need any supplies')
         self.state = self.cleanup
 
+    def find_ground_items(self):
+        stack_count = self.ground_stack_counts.get(self.target_items, 0)
+        if self.target_maximum_stacks and stack_count >= self.target_maximum_stacks:
+            self.checked_ground_groups.add(self.target_items)
+            self.state = self.check_supplies
+            return
+
+        print('Looking for supplies on the ground...')
+        p = utils.pint(self.g.pos)
+        ground_items = self.g.world.find_objects(self.target_items)
+        ground_items.sort(key=lambda obj: utils.phyp(
+            p, (obj.x, obj.y, obj.z)))
+
+        for obj in ground_items:
+            if obj.entity_id in self.checked_ground_items:
+                continue
+            current_count = self.g.game.count_items(self.target_items)
+            if utils.phyp(p, (obj.x, obj.y, obj.z)) > 16:
+                continue
+            target = utils.pint((obj.x, obj.y, obj.z))
+            self.ground_entity_id = obj.entity_id
+            self.ground_position = target
+            self.ground_required_count = max(
+                self.target_minimum, current_count + 1)
+            if utils.phyp(p, (obj.x, obj.y, obj.z)) <= 1.25:
+                self.g.path = []
+                self.wait_time = 1.0
+                self.state = self.wait_for_ground_item
+                return
+
+            navpath = self.g.world.path_to_place(p, target)
+            if not navpath:
+                self.checked_ground_items.add(obj.entity_id)
+                continue
+
+            print('Going to ground item', obj.item_id, target)
+            self.g.path = navpath
+            self.state = self.going_to_ground_item
+            return
+
+        self.checked_ground_groups.add(self.target_items)
+        self.state = self.check_supplies
+
+    def going_to_ground_item(self):
+        obj = self.g.objects.get(self.ground_entity_id)
+        have_enough = (self.g.game.count_items(self.target_items) >=
+                       self.ground_required_count)
+        if have_enough:
+            print('Picked up required ground supplies')
+            self._record_ground_pickup()
+            self.state = self.check_supplies
+        elif (obj is None or
+              utils.pint(self.g.pos) == self.ground_position or
+              utils.phyp(self.g.pos, (obj.x, obj.y, obj.z)) <= 1.25):
+            self.wait_time = 1.0
+            self.state = self.wait_for_ground_item
+
+    def wait_for_ground_item(self):
+        if (self.g.game.count_items(self.target_items) >=
+                self.ground_required_count):
+            print('Picked up required ground supplies')
+            self._record_ground_pickup()
+            self.state = self.check_supplies
+            return
+
+        if self.wait_time > 0:
+            self.wait_time -= utils.TICK
+            return
+
+        if self.ground_entity_id not in self.g.objects:
+            # Entity removal and inventory updates can arrive in either order.
+            self.state = self.check_supplies
+            return
+
+        print('Could not pick up ground item; trying another')
+        self.checked_ground_items.add(self.ground_entity_id)
+        self.g.path = []
+        self.ground_entity_id = None
+        self.ground_position = None
+        self.state = self.find_ground_items
+
     def find_barrels(self):
         print('Finding barrels...')
         w = self.g.world
@@ -215,7 +359,7 @@ class GrabSuppliesStates:
         self.g = global_state
         self.state = self.idle
 
-        self.supplies = {}
+        self.supplies = self.default_supplies()
         self.barrels = []
         self.checked_barrels = []
         self.bad_barrels = []