import os
import random
import time
from pynput import keyboard

base_path = "/sys/class/leds/rgb:kbd_backlight"
brightness_path = os.path.join(base_path, "brightness")

# Globals
snake = [(3, 3), (3, 2), (3, 1)]  # Initial snake position
direction = "RIGHT"  # Initial direction
food = None
game_running = True
grid_state = {}  # Tracks the current color of each cell

def main():
    if not init():
        return

    print("Starting Snake...")
    spawn_food()

    # Start the keyboard listener for input
    with keyboard.Listener(on_press=on_press, on_release=on_release) as listener:
        game_loop()
        listener.join()

def init():
    if not os.path.exists(brightness_path):
        print(f"Brightness path '{brightness_path}' not found. Are you sure the device supports it?")
        return False
    else:
        reset_grid()
        return True

def spawn_food():
    global food
    while True:
        food = (random.randint(0, 5), random.randint(0, 19))
        if food not in snake:  # Ensure food doesn't spawn on the snake
            break

def game_loop():
    global snake, food, direction, game_running

    while game_running:
        # Move the snake
        head = snake[0]
        if direction == "UP":
            new_head = (head[0] - 1, head[1])
        elif direction == "DOWN":
            new_head = (head[0] + 1, head[1])
        elif direction == "LEFT":
            new_head = (head[0], head[1] - 1)
        elif direction == "RIGHT":
            new_head = (head[0], head[1] + 1)

        # Check collisions
        if (
            new_head[0] < 0 or new_head[0] >= 6 or
            new_head[1] < 0 or new_head[1] >= 20 or
            new_head in snake
        ):
            for segment in snake:
                update_cell(segment, (255,255,255))

            game_running = False
            print("Game Over! Your score:", len(snake) - 3)
            break

        # Check food collision
        if new_head == food:
            snake.insert(0, new_head)  # Grow the snake
            spawn_food()  # Place new food
        else:
            snake.insert(0, new_head)
            tail = snake.pop()  # Move the snake
            if tail == (0,0):
                update_cell(tail, (255, 255, 0)) # Escape button
            else:
                update_cell(tail, (0, 0, 0)) # Clear the old tail

        # Update the grid
        update_cell(food, (0, 255, 0))  # Food color
        update_cell(new_head, (255, 0, 0))  # Snake head color

        time.sleep(0.3)  # Game speed

def reset_grid():
    """Reset the entire grid to black."""
    global grid_state
    grid_state = {}
    for key in range(120):
        grid_state[key] = (0, 0, 0)  # Set all cells to black
        set_colour((0, 0, 0), key)

    update_cell((0,0), (255, 255, 0)) # Escape button

def update_cell(coords, color):
    """Update a single cell if its color has changed."""
    global grid_state
    key = coors_to_key(coords[1], coords[0])
    if grid_state.get(key) != color:  # Only update if color has changed
        grid_state[key] = color
        set_colour(color, key)

def set_colour(color, key):
    write_value(key_to_path(key), f"{color[0]} {color[1]} {color[2]}")

def on_press(key):
    global direction
    try:
        if key == keyboard.Key.up and direction != "DOWN":
            direction = "UP"
        elif key == keyboard.Key.down and direction != "UP":
            direction = "DOWN"
        elif key == keyboard.Key.left and direction != "RIGHT":
            direction = "LEFT"
        elif key == keyboard.Key.right and direction != "LEFT":
            direction = "RIGHT"
    except AttributeError:
        pass  # Ignore other keys

def on_release(key):
    if key == keyboard.Key.esc:
        set_base_colors()
        global game_running
        game_running = False
        return False # Exit the game

def write_value(path, value):
    try:
        with open(path, 'w') as f:
            f.write(str(value))
    except Exception as e:
        print(f"Failed to set value: {e}")

def key_to_path(key):
    if key > 0:
        return f"{base_path}_{key}/multi_intensity"
    else:
        return f"{base_path}/multi_intensity"

def coors_to_key(x, y):
    return y * 20 + x

def lerp_rgb(color1, color2, steps = 120):
    if steps < 2:
        raise ValueError("Steps must be at least 2 to interpolate between colors.")

    # Calculate the step size for each channel
    step_colors = [
        [
            int(color1[i] + (color2[i] - color1[i]) * t / (steps - 1))
            for i in range(3)
        ]
        for t in range(steps)
    ]

    return step_colors

def set_base_colors():
    grad_start = (0,255,0)
    grad_end = (0,0,255)

    color_wheel = lerp_rgb(grad_start,grad_end)
    for i, color in enumerate(color_wheel, 0):
        set_colour(color, i)

if __name__ == "__main__":
    main()