mice/engine/unit_manager.py

import random
import uuid
from units import gas, rat, bomb, mine



from units.unit import UnitType

class UnitManager:
    def __init__(self, game):
        self.game = game

    def _spawnable_rat_positions(self):
        positions = []
        for y in range(1, self.game.map.height - 1):
            for x in range(1, self.game.map.width - 1):
                if not self.game.map.is_empty(x, y):
                    continue
                for dx, dy in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                    nx = x + dx
                    ny = y + dy
                    if self.game.map.in_bounds(nx, ny) and self.game.map.is_empty(nx, ny):
                        positions.append((x, y))
                        break
        return positions

    def has_weapon_at(self, position):
        """Check if there's a weapon (bomb, gas, mine) at the given position"""
        weapon_types = {UnitType.BOMB_TIMER, UnitType.BOMB_NUCLEAR, UnitType.GAS, UnitType.MINE}
        for unit in self.game.units.values():
            if unit.position == position and unit.type in weapon_types:
                return True
        return False

    def can_place_weapon_at(self, position):
        x, y = position
        if not self.game.map.in_bounds(x, y):
            return False
        if not self.game.map.is_empty(x, y):
            return False
        if self.has_weapon_at(position):
            return False
        return True
        
    def count_rats(self):
        count = 0
        rat_types = {UnitType.RAT_MALE, UnitType.RAT_FEMALE}
        for unit in self.game.units.values():
            if unit.type in rat_types:
                count += 1
        return count

    def refill_ammo(self):
        """Randomly refill ammo during gameplay."""
        import random
        for ammo_type, data in self.game.ammo.items():
            if ammo_type == "bomb":
                if random.random() < 0.02:
                    data["count"] = min(data["count"] + 1, data["max"])
            elif ammo_type == "mine":
                if random.random() < 0.05:
                    data["count"] = min(data["count"] + 1, data["max"])
            elif ammo_type == "gas":
                if random.random() < 0.01:
                    data["count"] = min(data["count"] + 1, data["max"])

    def spawn_gas(self, parent_id=None):
        if not self.can_place_weapon_at(self.game.pointer):
            return
        if self.game.ammo["gas"]["count"] <= 0:
            return
        self.game.ammo["gas"]["count"] -= 1
        self.game.render_engine.play_sound("GAS.WAV")
        self.spawn_unit(gas.Gas, self.game.pointer, parent_id=parent_id)

    def spawn_rat(self, position=None):
        if position is None:
            position = self.choose_start()
        if position is None:
            print("[flow] spawn_rat aborted: no valid spawn position", flush=True)
            return

        print(f"[flow] spawn_rat using position={position}", flush=True)
        
        # Don't spawn rats on top of weapons
        if self.has_weapon_at(position):
            # Try nearby positions
            for dx, dy in [(0,1), (1,0), (0,-1), (-1,0), (1,1), (-1,-1), (1,-1), (-1,1)]:
                alt_pos = (position[0] + dx, position[1] + dy)
                if not self.game.map.in_bounds(alt_pos[0], alt_pos[1]):
                    continue
                if self.game.map.is_empty(alt_pos[0], alt_pos[1]) and not self.has_weapon_at(alt_pos):
                    position = alt_pos
                    break
            else:
                # All nearby positions blocked, abort spawn
                print(f"[flow] spawn_rat aborted: weapon blocks spawn near {position}", flush=True)
                return
        
        rat_class = rat.Male if random.random() < 0.5 else rat.Female
        self.spawn_unit(rat_class, position)
        
    def spawn_bomb(self, position):
        if not self.can_place_weapon_at(position):
            return
        if self.game.ammo["bomb"]["count"] <= 0:
            return
        self.game.render_engine.play_sound("PUTDOWN.WAV")
        self.spawn_unit(bomb.Timer, position)
        self.game.ammo["bomb"]["count"] -= 1

    def spawn_nuclear_bomb(self, position):
        """Spawn a nuclear bomb at the specified position"""
        if self.game.ammo["nuclear"]["count"] <= 0:
            return
        if not self.can_place_weapon_at(position):
            return
        self.game.render_engine.play_sound("NUCLEAR.WAV")
        self.game.ammo["nuclear"]["count"] -= 1
        self.spawn_unit(bomb.NuclearBomb, position)
        
    def spawn_mine(self, position):
        if self.game.ammo["mine"]["count"] <= 0:
            return
        if not self.can_place_weapon_at(position):
            return
        self.game.render_engine.play_sound("PUTDOWN.WAV")
        self.game.ammo["mine"]["count"] -= 1
        self.spawn_unit(mine.Mine, position, on_bottom=True)

    def spawn_unit(self, unit, position, on_bottom=False, **kwargs):
        id = uuid.uuid4()
        if on_bottom:
            self.game.units = {id: unit(self.game, position, id, **kwargs), **self.game.units}
        else:
            self.game.units[id] = unit(self.game, position, id, **kwargs)

    def choose_start(self):
        if not hasattr(self.game, '_valid_positions') or self.game._valid_positions is None:
            self.game._valid_positions = self._spawnable_rat_positions()
            print(f"[flow] choose_start computed {len(self.game._valid_positions)} spawnable cells", flush=True)
        if not self.game._valid_positions:
            return None
        return random.choice(self.game._valid_positions)
    
    def get_unit_by_id(self, id):
        return self.game.units.get(id) or None