import time
import board
import digitalio
import busio
import random
import adafruit_ssd1306
import neopixel
import usb_hid
from adafruit_hid.keyboard import Keyboard
from adafruit_hid.keycode import Keycode
from adafruit_hid.keyboard_layout_us import KeyboardLayoutUS

# Pin Definitions
ROW_PINS = [board.GP2, board.GP1, board.GP26, board.GP27]  # Row GPIO pins
COL_PINS = [board.GP4, board.GP28, board.GP29]            # Column GPIO pins
OLED_SDA = board.GP6                                      # OLED SDA pin
OLED_SCL = board.GP7                                      # OLED SCL pin
LED_PIN = board.GP3                                       # WS2812 LED pin (using GP0)

# OLED Display Settings
SCREEN_WIDTH = 128
SCREEN_HEIGHT = 32
OLED_ADDR = 0x3C

# LED Settings
NUM_LEDS = 12

# Key mapping definitions
KEYMAP = [
    ['7', '8', '9'],
    ['4', '5', '6'],
    ['1', '2', '3'],
    ['0', '.', 'Enter']
]

# HID Keycode mapping
KEY_TO_KEYCODE = {
    '7': Keycode.SEVEN,
    '8': Keycode.EIGHT,
    '9': Keycode.NINE,
    '4': Keycode.FOUR,
    '5': Keycode.FIVE,
    '6': Keycode.SIX,
    '1': Keycode.ONE,
    '2': Keycode.TWO,
    '3': Keycode.THREE,
    '0': Keycode.ZERO,
    '.': Keycode.PERIOD,
    'Enter': Keycode.ENTER
}

# LED mapping for each key
KEY_TO_LED = [
    [1, 8, 9],    # 7, 8, 9
    [2, 7, 10],   # 4, 5, 6
    [3, 6, 11],   # 1, 2, 3
    [4, 5, 12]    # 0, ., Enter
]

# Define LED adjacency map for ripple effect
# Each LED has up to 4 adjacent LEDs (could be less)
# -1 means no adjacency in that direction
LED_ADJACENCY = [
    [-1, -1, -1, -1],  # Dummy (index 0, not used)
    [2, 8, -1, -1],    # LED 1 (Key 7)
    [3, 7, 1, -1],     # LED 2 (Key 4)
    [4, 6, 2, -1],     # LED 3 (Key 1)
    [5, -1, 3, -1],    # LED 4 (Key 0)
    [-1, 6, 4, -1],    # LED 5 (Key .)
    [7, 11, 5, 3],     # LED 6 (Key 2)
    [8, 10, 6, 2],     # LED 7 (Key 5)
    [9, -1, 7, 1],     # LED 8 (Key 8)
    [-1, 10, 8, -1],   # LED 9 (Key 9)
    [11, -1, 9, 7],    # LED 10 (Key 6)
    [12, -1, 10, 6],   # LED 11 (Key 3)
    [-1, -1, 11, -1]   # LED 12 (Key Enter)
]

# Setup USB HID
kbd = Keyboard(usb_hid.devices)
layout = KeyboardLayoutUS(kbd)

# Initialize I2C for OLED
i2c = busio.I2C(OLED_SCL, OLED_SDA)
display = adafruit_ssd1306.SSD1306_I2C(SCREEN_WIDTH, SCREEN_HEIGHT, i2c, addr=OLED_ADDR)

# Initialize LED strip
pixels = neopixel.NeoPixel(LED_PIN, NUM_LEDS, brightness=0.5, auto_write=False)

# Initialize row pins as outputs
row_pins = []
for pin in ROW_PINS:
    row_pin = digitalio.DigitalInOut(pin)
    row_pin.direction = digitalio.Direction.OUTPUT
    row_pin.value = False
    row_pins.append(row_pin)

# Initialize column pins as inputs with pull-down resistors
col_pins = []
for pin in COL_PINS:
    col_pin = digitalio.DigitalInOut(pin)
    col_pin.direction = digitalio.Direction.INPUT
    col_pin.pull = digitalio.Pull.DOWN
    col_pins.append(col_pin)

# State tracking
current_key_states = [[False for _ in range(3)] for _ in range(4)]
previous_key_states = [[False for _ in range(3)] for _ in range(4)]
led_brightness = 50  # Default LED brightness (0-100)
settings_mode = False
brightness_input = ""

# LED animation
class RippleState:
    def __init__(self):
        self.active = False
        self.origin_led = 0
        self.color = (0, 0, 0)
        self.distance = [255] * (NUM_LEDS + 1)
        self.wave_front = 0
        self.start_time = 0

MAX_RIPPLES = 5
ripples = [RippleState() for _ in range(MAX_RIPPLES)]

def wheel(pos):
    """Generate rainbow colors across 0-255 positions."""
    if pos < 85:
        return (255 - pos * 3, 0, pos * 3)
    elif pos < 170:
        pos -= 85
        return (0, pos * 3, 255 - pos * 3)
    else:
        pos -= 170
        return (pos * 3, 255 - pos * 3, 0)

def get_random_rgb_color():
    """Generate a random fully saturated color."""
    return wheel(random.randint(0, 255))

def calculate_ripple_distances(ripple):
    """Calculate distances from origin to all LEDs using breadth-first search."""
    # Initialize all distances to max value
    for i in range(NUM_LEDS + 1):
        ripple.distance[i] = 255  # Infinity
    
    # Distance to origin is 0
    ripple.distance[ripple.origin_led] = 0
    
    # Queue for BFS
    queue = []
    front, rear = 0, 0
    
    # Add origin to queue
    queue.append(ripple.origin_led)
    rear += 1
    
    # BFS to find distances
    while front < len(queue):
        current = queue[front]
        front += 1
        current_dist = ripple.distance[current]
        
        # Check all adjacent LEDs
        for i in range(4):
            next_led = LED_ADJACENCY[current][i]
            if next_led != -1 and ripple.distance[next_led] == 255:
                ripple.distance[next_led] = current_dist + 1
                queue.append(next_led)

def start_ripple(led_index):
    """Start a new ripple effect from the specified LED."""
    # Find an available ripple slot
    for ripple in ripples:
        if not ripple.active:
            ripple.active = True
            ripple.origin_led = led_index
            ripple.color = get_random_rgb_color()
            ripple.wave_front = 0
            ripple.start_time = time.monotonic()
            
            # Calculate distances from origin to all LEDs
            calculate_ripple_distances(ripple)
            return
    
    # If no slots available, replace the oldest ripple
    oldest_time = time.monotonic()
    oldest_index = 0
    
    for i, ripple in enumerate(ripples):
        if ripple.start_time < oldest_time:
            oldest_time = ripple.start_time
            oldest_index = i
    
    ripple = ripples[oldest_index]
    ripple.active = True
    ripple.origin_led = led_index
    ripple.color = get_random_rgb_color()
    ripple.wave_front = 0
    ripple.start_time = time.monotonic()
    
    # Calculate distances from origin to all LEDs
    calculate_ripple_distances(ripple)

def scan_keys():
    """Scan the key matrix and detect key presses/releases."""
    global current_key_states, previous_key_states, settings_mode, brightness_input
    
    # Copy current states to previous states
    for r in range(4):
        for c in range(3):
            previous_key_states[r][c] = current_key_states[r][c]
    
    # Reset current key states
    for r in range(4):
        for c in range(3):
            current_key_states[r][c] = False
    
    # Scan the matrix
    for row in range(4):
        # Set current row high
        row_pins[row].value = True
        
        # Small delay to stabilize
        time.sleep(0.001)
        
        # Check all columns in this row
        for col in range(3):
            # A high reading means the key is pressed
            current_key_states[row][col] = col_pins[col].value
        
        # Set row back to low
        row_pins[row].value = False
    
    # Check for key press events
    key_events = []
    
    # Process key presses
    for row in range(4):
        for col in range(3):
            # Key press detected
            if current_key_states[row][col] and not previous_key_states[row][col]:
                key = KEYMAP[row][col]
                key_events.append(key)
                
                # Start ripple effect for this key
                start_ripple(KEY_TO_LED[row][col])
                
                # Send key press via USB HID
                if not settings_mode:
                    # Press the key
                    kbd.press(KEY_TO_KEYCODE[key])
                else:
                    # Process settings mode input
                    process_settings_mode(key)
                
                # Check for settings mode combo (1, 9, and Enter pressed together)
                if (current_key_states[2][0] and  # 1
                    current_key_states[0][2] and  # 9
                    current_key_states[3][2]):    # Enter
                    settings_mode = True
                    brightness_input = ""
    
    # Process key releases
    for row in range(4):
        for col in range(3):
            # Key release detected
            if not current_key_states[row][col] and previous_key_states[row][col]:
                key = KEYMAP[row][col]
                
                # Release the key via USB HID if not in settings mode
                if not settings_mode:
                    kbd.release(KEY_TO_KEYCODE[key])
    
    return key_events

def update_leds():
    """Update the LED strip with active ripple effects."""
    global ripples
    
    # Clear all LEDs
    for i in range(NUM_LEDS):
        pixels[i] = (0, 0, 0)
    
    # Update active ripples
    current_time = time.monotonic()
    
    for ripple in ripples:
        if ripple.active:
            # Calculate ripple progress
            ripple_age = current_time - ripple.start_time
            
            # Advance wave front based on time - one step every 100ms
            ripple.wave_front = int(ripple_age * 10)  # Convert to 100ms units
            
            # End ripple after it has traveled past all LEDs
            if ripple.wave_front > 10:
                ripple.active = False
                continue
            
            # Update each LED based on distance from origin
            for led in range(1, NUM_LEDS + 1):
                distance = ripple.distance[led]
                
                # Skip if LED is unreachable from origin
                if distance == 255:
                    continue
                
                # Check if this LED is at the wave front
                wave_diff = ripple.wave_front - distance
                
                if 0 <= wave_diff <= 3:
                    # LED is in the active part of the wave
                    intensity = 255 - (wave_diff * 85)  # Fade intensity as wave passes
                    
                    # Extract RGB components from original color and scale by intensity
                    r, g, b = ripple.color
                    r = (r * intensity) // 255
                    g = (g * intensity) // 255
                    b = (b * intensity) // 255
                    
                    # Apply color to LED, blending with existing colors
                    pixel_index = led - 1  # Convert to 0-based index
                    er, eg, eb = pixels[pixel_index]
                    
                    # Blend by taking maximum of each color component
                    nr = max(er, r)
                    ng = max(eg, g)
                    nb = max(eb, b)
                    
                    pixels[pixel_index] = (nr, ng, nb)
    
    # Show the updated LEDs
    pixels.show()

def update_oled():
    """Update the OLED display."""
    display.fill(0)
    display.text("Macropad", 0, 0, 1)
    
    if settings_mode:
        display.text("SETTINGS MODE", 0, 8, 1)
        display.text("LED Brightness:", 0, 16, 1)
        display.text(brightness_input if brightness_input else "0", 0, 24, 1)
        display.text("%", len(brightness_input) * 8, 24, 1)
    else:
        display.text("Ready", 64, 0, 1)
        display.text(f"Brightness: {led_brightness}%", 0, 16, 1)
    
    display.show()

def process_settings_mode(key):
    """Process input in settings mode."""
    global brightness_input, led_brightness, settings_mode
    
    if key == "Enter":
        # Enter key - save settings and exit settings mode
        if brightness_input:
            new_brightness = int(brightness_input)
            if 0 <= new_brightness <= 100:
                led_brightness = new_brightness
                pixels.brightness = led_brightness / 100
        settings_mode = False
    elif key in "0123456789":
        # Digit key - add to brightness input
        if len(brightness_input) < 3:  # Limit to 3 digits
            brightness_input += key
            # Ensure value is <= 100
            value = int(brightness_input)
            if value > 100:
                brightness_input = "100"
    elif key == ".":
        # Period key - clear brightness input
        brightness_input = ""

# Main loop
display.fill(0)
display.text("Macropad Ready", 0, 0, 1)
display.text(f"Brightness: {led_brightness}%", 0, 16, 1)
display.show()

while True:
    key_events = scan_keys()
    update_leds()
    update_oled()
    time.sleep(0.01)  # Small delay to prevent CPU hogging