proteus/game/engine/arcengine/base_game.py

"""
Module for the base game class in ARCEngine.
"""

import os
from abc import ABC
from typing import List, Optional, final

import numpy as np
from numpy import ndarray

from .camera import Camera
from .enums import ActionInput, FrameData, FrameDataRaw, GameAction, GameState
from .level import Level
from .sprites import Sprite

MAX_FRAME_PER_ACTION: int = 1000


class ARCBaseGame(ABC):
    """Base class for ARCEngine games that manages levels and camera.

    This is a base class that games should inherit from. and extend with game logic.
    It handles the game loop and rendering.

    Custom game logic should be implemented in the step() method.
    """

    _game_id: str
    _levels: list[Level]
    _clean_levels: list[Level]
    _current_level_index: int
    _camera: Camera
    _debug: bool
    _action: ActionInput
    _action_complete: bool
    _action_count: int
    _state: GameState
    _score: int
    _next_level: bool
    _full_reset: bool
    _win_score: int
    _available_actions: list[int]
    _placeable_sprite: Optional[Sprite]
    _seed: int

    def __init__(
        self,
        game_id: str,
        levels: List[Level],
        camera: Optional[Camera] = None,
        debug: bool = False,
        win_score: int = 1,
        available_actions: list[int] = [1, 2, 3, 4, 5, 6],
        seed: int = 0,
    ) -> None:
        """Initialize a new game.

        Args:
            levels: List of levels to initialize the game with. Each level will be cloned.
            camera: Optional camera to use. If not provided, a default 64x64 camera will be created.

        Raises:
            ValueError: If levels list is empty
        """
        if not levels:
            raise ValueError("Game must have at least one level")

        # Game ID should be set by subclasses
        self._game_id = game_id

        # Clone each level to prevent external modification
        self._levels = [level.clone() for level in levels]
        self._clean_levels = [level.clone() for level in levels]
        self._current_level_index = 0

        # Debug mode
        self._debug = debug

        # Camera
        # Use provided camera or create default
        self._camera = camera if camera is not None else Camera()

        # Game state
        self._state = GameState.NOT_PLAYED
        self._score = 0
        self._next_level = False
        self._action = ActionInput()
        self._action_complete = False
        self._action_count = 0
        self._full_reset = False
        self._win_score = win_score if win_score > 1 else len(levels)
        self.set_level(0)
        self._available_actions = available_actions
        self._placeable_sprite = None
        self._seed = seed

    def debug(self, message: str) -> None:
        """Debug mode.

        Args:
            message: The message to print
        """
        if self._debug:
            print(message)

    @property
    @final
    def current_level(self) -> Level:
        """Get the current level.

        Returns:
            Level: The current level
        """
        return self._levels[self._current_level_index]

    @property
    @final
    def camera(self) -> Camera:
        """Get the game's camera.

        Returns:
            Camera: The game's camera
        """
        return self._camera

    @property
    @final
    def game_id(self) -> str:
        """Get the game's ID.

        Returns:
            str: The game's ID
        """
        return self._game_id

    @property
    @final
    def win_score(self) -> int:
        """Get the game's max score.

        Returns:
            int: The game's max score
        """
        return self._win_score

    @final
    def set_level(self, index: int) -> None:
        """Set the current level by index.

        Args:
            index: The index of the level to set as current

        Raises:
            IndexError: If index is out of range
        """
        if not 0 <= index < len(self._levels):
            raise IndexError(f"Level index {index} out of range [0, {len(self._levels)})")
        self._current_level_index = index
        self._action_count = 0
        level = self.current_level
        if level.grid_size:
            self.camera.resize(level.grid_size[0], level.grid_size[1])
        self.on_set_level(level)

    def set_level_by_name(self, name: str) -> None:
        """Set the current level by name.

        Args:
            name: The name of the level to set as current
        """
        for index, level in enumerate(self._levels):
            if level.name == name:
                self.set_level(index)
                return
        raise ValueError(f"Level with name {name} not found")

    @property
    @final
    def level_index(self) -> int:
        """Get the current level index.

        Returns:
            int: The current level index
        """
        return self._current_level_index

    @final
    def perform_action(self, action_input: ActionInput, raw: bool = False) -> FrameData | FrameDataRaw:
        """Perform an action and return the resulting frame data.

        DO NOT OVERRIDE THIS METHOD, Your Game Logic should be in step()

        The base implementation:
        1. While the action is not complete, call step() and render frames
        2. Returns a FrameData object with the current state

        Args:
            action_input: The action to perform

        Returns:
            FrameData: The resulting frame data
        """
        self._full_reset = False
        if action_input.id == GameAction.RESET:
            self.handle_reset()
        elif self._state == GameState.GAME_OVER or self._state == GameState.WIN:
            return FrameData(
                game_id=self._game_id,
                frame=[],
                state=self._state,
                score=self._score,
                win_score=self._win_score,
                action_input=action_input,
                available_actions=self._available_actions,
            )

        self._set_action(action_input)

        frame_list: list[ndarray | list[list[int]]] = []

        count = 0

        while not self.is_action_complete():
            if count > MAX_FRAME_PER_ACTION:
                raise ValueError("Action took too many frames")
            count += 1
            if self._next_level:
                self._really_set_next_level()
            else:
                self.step()
            frame = self.camera.render(self.current_level.get_sprites())
            if raw:
                frame_list.append(frame)
            else:
                frame_list.append(frame.tolist())

        # Create and return FrameData
        if raw:
            frame_raw = FrameDataRaw()
            frame_raw.game_id = self._game_id
            frame_raw.frame = frame_list
            frame_raw.state = self._state
            frame_raw.levels_completed = self._score
            frame_raw.win_levels = self._win_score
            frame_raw.action_input = action_input
            frame_raw.full_reset = self._full_reset
            frame_raw.available_actions = self._available_actions
            return frame_raw

        return FrameData(
            game_id=self._game_id,
            frame=frame_list,
            state=self._state,
            levels_completed=self._score,
            win_levels=self._win_score,
            action_input=action_input,
            full_reset=self._full_reset,
            available_actions=self._available_actions,
        )

    @property
    @final
    def action(self) -> ActionInput:
        """Get the current action."""
        return self._action

    @final
    def _set_action(self, action_input: ActionInput) -> None:
        """Set the action to perform.

        Args:
            action_input: The action to perform
        """
        self._state = GameState.NOT_FINISHED
        self._action = action_input
        self._action_complete = False
        if action_input.id != GameAction.RESET:
            self._action_count += 1

    @final
    def complete_action(self) -> None:
        """Complete the action. Call this when the provided action is fully resolved"""
        self._action_complete = True

    @final
    def is_action_complete(self) -> bool:
        """Check if the action is complete.

        Returns:
            bool: True if the action is complete, False otherwise
        """
        return not self._next_level and self._action_complete

    @final
    def win(self) -> None:
        """Call this when the player has beaten the game."""
        self._state = GameState.WIN

    @final
    def lose(self) -> None:
        """Call this when the player has losses the game."""
        self._state = GameState.GAME_OVER

    def handle_reset(self) -> None:
        """Handle the reset of the game.

        If the action count is 0, perform a full reset.
        Otherwise, perform a level reset.
        """
        if os.getenv("ONLY_RESET_LEVELS") == "true" and self._state != GameState.WIN:
            self.level_reset()
        elif self._action_count == 0 or self._state == GameState.WIN:
            self.full_reset()
        else:
            self.level_reset()

    def full_reset(self) -> None:
        self._levels = [level.clone() for level in self._clean_levels]
        self._score = 0
        self._action_count = 0
        self._full_reset = True
        self.set_level(0)
        self._state = GameState.NOT_FINISHED

    def level_reset(self) -> None:
        self._levels[self._current_level_index] = self._clean_levels[self._current_level_index].clone()
        self.set_level(self._current_level_index)
        self._state = GameState.NOT_FINISHED

    def step(self) -> None:
        """Step the game.  This is where your game logic should be implemented.

        REQUIRED: Call complete_action() when the action is complete.
          It does not need to be called every step, but once the action is complete.
          The engine will keep calling step and rendering frames until the action is complete.
        """

        self.complete_action()

    def try_move(self, sprite_name: str, dx: int, dy: int) -> List[Sprite]:
        """Try to move a sprite and return a list of sprites it collides with.

        This method attempts to move the sprite by the given deltas and checks for collisions.
        If any collisions are detected, the sprite is not moved and the method returns a list
        of sprite names that were collided with.

        Args:
            sprite_name: The name of the sprite to move.
            dx: The change in x position (positive = right, negative = left).
            dy: The change in y position (positive = down, negative = up).

        Returns:
            A list of sprite names that the sprite collided with. If no collisions occurred,
            the sprite is moved and an empty list is returned.

        Raises:
            ValueError: If no sprite with the given name is found.
        """
        # Get the sprite to move
        sprites = self.current_level.get_sprites_by_name(sprite_name)
        if not sprites:
            raise ValueError(f"No sprite found with name: {sprite_name}")
        return self.try_move_sprite(sprites[0], dx, dy)  # Use the first sprite with this name

    def try_move_sprite(self, sprite: Sprite, dx: int, dy: int) -> List[Sprite]:
        """Try to move a sprite and return a list of sprites it collides with.

        This method attempts to move the sprite by the given deltas and checks for collisions.
        If any collisions are detected, the sprite is not moved and the method returns a list
        of sprite names that were collided with.

        Args:
            sprite_name: The name of the sprite to move.
            dx: The change in x position (positive = right, negative = left).
            dy: The change in y position (positive = down, negative = up).

        Returns:
            A list of sprite names that the sprite collided with. If no collisions occurred,
            the sprite is moved and an empty list is returned.

        Raises:
            ValueError: If no sprite with the given name is found.
        """  # Get the sprite to move
        # Store original position
        original_x = sprite.x
        original_y = sprite.y

        # Try the move
        sprite.move(dx, dy)

        # Check for collisions with all other sprites
        collisions = []
        for other in self.current_level.get_sprites():
            if sprite.collides_with(other):
                collisions.append(other)

        # If there were collisions, revert the move
        if collisions:
            sprite.set_position(original_x, original_y)

        return collisions

    def is_last_level(self) -> bool:
        """Check if the current level is the last level.

        Returns:
            bool: True if the current level is the last level, False otherwise
        """
        return self._current_level_index == len(self._levels) - 1

    def next_level(self) -> None:
        """Move to the next level."""
        self._score += 1
        if not self.is_last_level():
            self._next_level = True
        else:
            self.win()

    def _really_set_next_level(self) -> None:
        self.set_level(self._current_level_index + 1)
        self._next_level = False

    def on_set_level(self, level: Level) -> None:
        """Called when the level is set, use this to set level specific data."""
        pass

    def get_pixels_at_sprite(self, sprite: Sprite) -> ndarray:
        """Get the camera pixels at a sprite.

        Args:
            sprite: The sprite to get the pixels at

        Returns:
            list[list[int]]: The camera returned pixels at the sprite
        """
        return self.get_pixels(sprite.x - self.camera.x, sprite.y - self.camera.y, sprite.width, sprite.height)

    def get_pixels(self, x: int, y: int, width: int, height: int) -> ndarray:
        """Get the camera pixels at a given position.

        Args:
            x: The x position to get the pixels at
            y: The y position to get the pixels at
            width: The width of the area to get the pixels at
            height: The height of the area to get the pixels at

        Returns:
            list[list[int]]: The camera returned pixels at the given position and width/height
        """

        frame = self.camera._raw_render(self.current_level.get_sprites())
        return frame[y : y + height, x : x + width]

    def set_placeable_sprite(self, sprite: Sprite | None) -> None:
        """Set the placeable sprite.

        Args:
            sprite: The sprite to set as placeable
        """
        self._placeable_sprite = sprite

    def _get_graph_location(self) -> tuple[float, float, float] | None:
        """Get the location for the graph builder to use for this state.

        Returns:
            tuple[float, float, float] | None: The location this state should report as its location in the graph,
                or None if the graph builder should calculate the location.
        """
        return None

    def _get_hidden_state(self) -> ndarray:
        """Get the hidden state for the graph builder to use for this state.

        Returns:
            ndarray: Hidden state for the graph builder to use for this state.
        """
        return np.zeros((4, 4), dtype=np.int8)

    def _get_valid_actions(self) -> list[ActionInput]:
        """Get the valid actions for the current game state.

        Note: This method is for internal use only, the data here is never exposed
        via the API or to Users/Agents.

        Returns:
            list[int]: The valid actions for the current game state
        """
        valid_actions: list[ActionInput] = []

        for action in self._available_actions:
            match action:
                case 1 | 2 | 3 | 4 | 5:
                    valid_actions.append(ActionInput(id=GameAction.from_id(action)))
                case 6:
                    if self._placeable_sprite:
                        valid_actions.extend(self._get_valid_placeble_actions())
                    else:
                        valid_actions.extend(self._get_valid_clickable_actions())

        return valid_actions

    def _get_valid_placeble_actions(self) -> list[ActionInput]:
        """Get valid placeable actions from placeable areas.

        Returns:
            list[ActionInput]: List of valid placeable actions with screen coordinates
        """

        scale, x_offset, y_offset = self.camera._calculate_scale_and_offset()

        valid_actions: list[ActionInput] = []

        for area in self.current_level.placeable_areas:
            for y in range(area.y, area.y + area.height, area.y_scale):
                for x in range(area.x, area.x + area.width, area.x_scale):
                    action_input = ActionInput(id=GameAction.ACTION6.value, data={"x": x * scale + x_offset, "y": y * scale + y_offset})
                    valid_actions.append(action_input)

        return valid_actions

    def _get_valid_clickable_actions(self) -> list[ActionInput]:
        """Get valid clickable actions from sprites with the 'sys_click' tag.

        This method finds all sprites tagged with 'sys_click' and generates clickable actions
        based on their pixel data and the presence of the 'sys_every_pixel' tag.

        Returns:
            list[ActionInput]: List of valid clickable actions with screen coordinates
        """
        valid_actions: list[ActionInput] = []

        # Get all sprites with the 'sys_click' tag
        clickable_sprites = self.current_level.get_sprites_by_tag("sys_click")
        clickable_sprites.extend(self.current_level.get_sprites_by_tag("sys_place"))

        scale, x_offset, y_offset = self.camera._calculate_scale_and_offset()

        for sprite in clickable_sprites:
            if not self._is_sprite_clickable_now(sprite):
                continue

            # Check if sprite has the 'sys_every_pixel' tag
            has_every_pixel = "sys_every_pixel" in sprite._tags

            # Get the rendered sprite pixels (accounts for scale, rotation, etc.)
            rendered_pixels = sprite.render()

            if has_every_pixel:
                # Every non-negative pixel is a valid action
                for y in range(rendered_pixels.shape[0]):
                    for x in range(rendered_pixels.shape[1]):
                        if rendered_pixels[y, x] >= 0:
                            # Convert sprite-relative coordinates to screen coordinates
                            screen_x = (sprite._x + x) * scale + x_offset
                            screen_y = (sprite._y + y) * scale + y_offset

                            # Create ActionInput with ACTION6 (ComplexAction) and coordinates
                            action_input = ActionInput(id=GameAction.ACTION6.value, data={"x": screen_x, "y": screen_y})
                            valid_actions.append(action_input)
            else:
                # Find any single non-negative pixel to represent the entire sprite
                found_pixel = False
                for y in range(rendered_pixels.shape[0]):
                    for x in range(rendered_pixels.shape[1]):
                        if rendered_pixels[y, x] >= 0:
                            # Convert sprite-relative coordinates to screen coordinates
                            screen_x = (sprite._x + x) * scale + x_offset
                            screen_y = (sprite._y + y) * scale + y_offset

                            # Create ActionInput with ACTION6 (ComplexAction) and coordinates
                            action_input = ActionInput(id=GameAction.ACTION6.value, data={"x": screen_x, "y": screen_y})
                            valid_actions.append(action_input)
                            found_pixel = True
                            break
                    if found_pixel:
                        break

        return valid_actions

    def _is_sprite_clickable_now(self, sprite: Sprite) -> bool:
        """
        Check if a sprite is clickable now.  This method is designed to be overridden by games
        that have more complex clickable logic. (e.g. not all sprites flagged with 'sys_click'
        are clickable at all times)

        Args:
            sprite: The sprite to check

        Returns:
            bool: True if the sprite is clickable now, False otherwise
        """
        return True