import board
import busio
import digitalio
from adafruit_mcp230xx.mcp23017 import MCP23017

class KeyInterface:
    _MCP23017_OLATA = 0x14
    _MCP23017_OLATB = 0x15
    
    def set_olata(self, val: int):
        self.mcp._write_u8(self._MCP23017_OLATA, val)
    
    def set_olatb(self, val: int):
        self.mcp._write_u8(self._MCP23017_OLATB, val)
    
    def __init__(self, i2c):
        self._i2c = i2c
        
        # Turn on MCP
        self._mcp_reset_pin = digitalio.DigitalInOut(board.D4)
        self._mcp_reset_pin.direction = digitalio.Direction.OUTPUT
        self._mcp_reset_pin.value = True
        
        # Initialize MCP
        self.mcp = MCP23017(self._i2c)
        
        # Set up MCP pins
        self.mcp.iodirb = 255  # GPIO B (rows) all input
        self.mcp.ipolb = 255   # GPIO B inverted polarity
        self.mcp.gppub = 255   # GPIO B all pulled up
        self.set_olata(0) # GPIO A all output low
        
    def scan(self):
        result = []
        for i in range(0, 6):
            self.mcp.iodira = 0xFF ^ (0x01 << i)
            result.append(self.mcp.gpiob)
        self.mcp.iodira = 0xFF
        return result