import pygame


class Input:
    def __init__(self, debug=True):
        self.debug = debug
        self.controller = None

        if not self.debug:
            pygame.joystick.init()

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

    def get_input(self, event):
        if self.debug:
            return self.get_keyboard_input(event)
        else:
            return self.get_controller_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_RETURN:
                return "SELECT"
            if event.key == pygame.K_ESCAPE:
                return "QUIT"

        return None

    def get_controller_input(self, event):
        if not self.controller:
            return None

        if event.type == pygame.JOYAXISMOTION:
            if event.axis == 1:  # Y-axis
                if event.value < -0.5:  # Joystick up
                    return "UP"
                if event.value > 0.5:  # Joystick down
                    return "DOWN"
        elif event.type == pygame.JOYBUTTONDOWN:
            if event.button == 0:  # Button A or select
                return "SELECT"
            if event.button == 1:  # Button B
                return "QUIT"

        return None