import board
from digitalio import DigitalInOut, Direction, Pull
from enum import Enum, auto
import time

class Button(Enum):
    BTN_A = auto()
    BTN_B = auto()
    BTN_C = auto()
    DIR_U = auto()
    DIR_R = auto()
    DIR_D = auto()
    DIR_L = auto()

class ControlInput:
    _buttons = {
            Button.BTN_A: DigitalInOut(board.D5),
            Button.BTN_B: DigitalInOut(board.D6),
            Button.BTN_C: DigitalInOut(board.D4),
            Button.DIR_U: DigitalInOut(board.D17),
            Button.DIR_R: DigitalInOut(board.D23),
            Button.DIR_L: DigitalInOut(board.D27),
            Button.DIR_D: DigitalInOut(board.D22),
            }

    _pressed = set()

    def __init__(self):
        for button in self._buttons:
           self._buttons[button].direction = Direction.INPUT
           self._buttons[button].pull = Pull.UP

    def get_one_shot(self, timeout: float = 0):
        started = time.time()
        while True:
            time.sleep(0.05)
            for button in self._buttons:
                if not self._buttons[button].value and not button in self._pressed:
                    self._pressed.add(button)
                    return button
                elif self._buttons[button].value:
                    self._pressed.discard(button)
            if timeout > 0 and time.time() - started > timeout:
                return None