import pygame


class Input:
    def __init__(self, debug=False):
        pygame.joystick.init()

        self.controller = None

        if pygame.joystick.get_count() > 0:
            self.controller = pygame.joystick.Joystick(0)

            self.controller.init()

        self.button_states = {}  # Current states
        self.prev_button_states = {}  # Previous states
        self.prev_axis_states = {"x": 0, "y": 0}
        self.prev_hat_states = (0, 0)  # D-pad

    def get_input(self, event):
        if self.controller:
            controller_input = self._get_controller_input(event)

            if controller_input:
                return controller_input

        return self._get_keyboard_input(event)

    def _get_keyboard_input(self, event):
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_UP:
                return "UP"

            if event.key == pygame.K_DOWN:
                return "DOWN"

            if event.key == pygame.K_LEFT:
                return "LEFT"

            if event.key == pygame.K_RIGHT:
                return "RIGHT"

            if event.key == pygame.K_RETURN:
                return "SELECT"

            if event.key == pygame.K_ESCAPE:
                return "QUIT"

        return None

    def _get_analog_state(self):
        if self.controller:
            x = self.controller.get_axis(0)  # X-axis
            y = self.controller.get_axis(1)  # Y-axis

            if x <= -0.1:
                return "RIGHT"

            if x >= 0.1:
                return "LEFT"

            if y <= -0.1:
                return "UP"

            if y >= 0.1:
                return "DOWN"

        return None

    def _get_controller_input(self, event):
        # Analog stick
        if event.type == pygame.JOYAXISMOTION:
            return self._get_analog_state()
        # D-pad
        elif event.type == pygame.JOYHATMOTION:
            prev_hat = self.prev_hat_states
            curr_hat = event.value
            self.prev_hat_states = curr_hat

            if prev_hat[0] == 0 and curr_hat[0] == -1:
                return "LEFT"

            if prev_hat[0] == 0 and curr_hat[0] == 1:
                return "RIGHT"

            if prev_hat[1] == 0 and curr_hat[1] == 1:
                return "UP"

            if prev_hat[1] == 0 and curr_hat[1] == -1:
                return "DOWN"
        # Buttons
        elif event.type == pygame.JOYBUTTONDOWN:
            button = event.button
            self.button_states[button] = True

            # Debounced
            if (
                button not in self.prev_button_states
                or not self.prev_button_states[button]
            ):
                if button == 0:
                    return "DOWN"

                if button == 1:
                    return "LEFT"

                if button == 2:
                    return "UP"

                if button == 3:
                    return "RIGHT"

                if button == 4:
                    return "START"

                if button == 5:
                    return "SELECT"
        elif event.type == pygame.JOYBUTTONUP:
            button = event.button
            self.button_states[button] = False

        self.prev_button_states = self.button_states.copy()

        return None