from snake_up_movement_state import SnakeUpMovementState
from snake_down_movement_state import SnakeDownMovementState
from snake_left_movement_state import SnakeLeftMovementState
from snake_right_movement_state import SnakeRightMovementState
from input_manager import InputManager


class SnakeStateMachine:
    def __init__(self):
        self.__state = SnakeLeftMovementState()

        self.__next_state_dict = {
            InputManager.UP: SnakeUpMovementState(),
            InputManager.DOWN: SnakeDownMovementState(),
            InputManager.LEFT: SnakeLeftMovementState(),
            InputManager.RIGHT: SnakeRightMovementState()
        }

    
    def update(self, user_input):
        self.__state = self.__next_state_dict.get(user_input, self.__state)


    def get_state(self):
        return self.__state