state.py

from constants import AppState


class StateMachine:
    def __init__(self):
        self.state: AppState = AppState.SETUP
        self.click_count: int = 0

    def transition(self, new_state: AppState) -> None:
        self.state = new_state

    def can_place(self) -> bool:
        return self.state == AppState.SETUP

    def can_start(self, grid) -> bool:
        return (self.state == AppState.SETUP
                and grid.start is not None
                and grid.end is not None)

    def handle_grid_click(self, row: int, col: int, grid) -> None:
        if self.state != AppState.SETUP:
            return
        if not grid.in_bounds(row, col):
            return

        if self.click_count == 0:
            # Place start
            grid.start = (row, col)
            grid.obstacles.discard((row, col))
            self.click_count = 1
        elif self.click_count == 1:
            # Place end (can't be same as start)
            if (row, col) != grid.start:
                grid.end = (row, col)
                grid.obstacles.discard((row, col))
                self.click_count = 2
        else:
            # Toggle obstacles (can't place on start/end)
            if (row, col) != grid.start and (row, col) != grid.end:
                grid.toggle_obstacle(row, col)

    def reset(self) -> None:
        self.state = AppState.SETUP
        self.click_count = 0