# snake_game.py
#
# MicroPython Snake game for qpad-xiao RP2040
# - 6 STEPTIME capacitive pads for direction + pause/reset
# - SSD1306 OLED for gameplay
# - WS2812 NeoPixel shows score colors + game over
#
# Pinning matches your which_pad_game.py example exactly!

from steptime import STEPTIME
from ws2812 import WS2812
from ssd1306 import SSD1306_I2C

import machine
from machine import Pin, I2C, freq
import utime
import urandom

# ----- Steptime needs high clock -----
freq(250000000)

# ----- OLED -----
i2c = I2C(1, scl=Pin(7), sda=Pin(6), freq=200000)
oled = SSD1306_I2C(128, 64, i2c)

# ----- NeoPixel -----
power = machine.Pin(11, machine.Pin.OUT)
power.value(1)

BLACK  = (0, 0, 0)
RED    = (255, 0, 0)
GREEN  = (0, 255, 0)
BLUE   = (0, 0, 255)
PURPLE = (180, 0, 255)
CYAN   = (0, 255, 255)
YELLOW = (255, 150, 0)

SCORE_COLORS = [RED, YELLOW, GREEN, CYAN, BLUE, PURPLE]

led = WS2812(12, 1, 0.5, 6)

def set_led_color(color):
    led.pixels_fill(color)
    led.pixels_show()

# ----- Touch pads -----
PAD_PINS = [26, 27, 1, 2, 3, 4]

for p in PAD_PINS:
    Pin(p, Pin.IN, Pin.PULL_UP)

channels = []
for i, pin in enumerate(PAD_PINS):
    channels.append([STEPTIME(i, pin), [1e6], [0]])

loop = 200
settle = 20000
thresh = 10000

def read_pads():
    """Return deltas for all 6 pads."""
    deltas = []
    for ch in channels:
        sm, min_val, printable = ch
        sm.put(loop)
        sm.put(settle)
        result = 4294967296 - sm.get()
        if result < min_val[0]:
            min_val[0] = result
        d = result - min_val[0]
        printable[0] = d
        deltas.append(d)
    return deltas

def strongest_pressed(deltas):
    """Return index of strongest pad above thresh, else -1."""
    best_i = -1
    best_v = thresh
    for i, v in enumerate(deltas):
        if v > best_v:
            best_v = v
            best_i = i
    return best_i

# ----- Snake constants -----
CELL_SIZE = 8
GRID_W = 128 // CELL_SIZE  # 16
GRID_H = 64 // CELL_SIZE   # 8

UP    = (0, -1)
DOWN  = (0, 1)
LEFT  = (-1, 0)
RIGHT = (1, 0)

PAD_UP, PAD_DOWN, PAD_LEFT, PAD_RIGHT, PAD_PAUSE, PAD_RESET = 0, 3, 1, 2, 4, 5

# ----- Game class -----
class SnakeGame:
    def __init__(self):
        self.reset()

    def reset(self):
        self.snake = [(GRID_W // 2, GRID_H // 2)]
        self.dir = RIGHT
        self.food = None
        self.place_food()
        self.alive = True
        self.score = 0

    def place_food(self):
        while True:
            x = urandom.getrandbits(8) % GRID_W
            y = urandom.getrandbits(8) % GRID_H
            if (x, y) not in self.snake:
                self.food = (x, y)
                return

    def set_dir_from_pad(self, pad_idx):
        if pad_idx == PAD_UP and self.dir != DOWN:
            self.dir = UP
        elif pad_idx == PAD_DOWN and self.dir != UP:
            self.dir = DOWN
        elif pad_idx == PAD_LEFT and self.dir != RIGHT:
            self.dir = LEFT
        elif pad_idx == PAD_RIGHT and self.dir != LEFT:
            self.dir = RIGHT

    def step(self):
        if not self.alive:
            return
        head_x, head_y = self.snake[0]
        dx, dy = self.dir
        nx = head_x + dx
        ny = head_y + dy
        new_head = (nx, ny)
        # Wall collision
        if nx < 0 or nx >= GRID_W or ny < 0 or ny >= GRID_H:
            self.alive = False
            return
        # Self collision
        if new_head in self.snake:
            self.alive = False
            return
        self.snake.insert(0, new_head)
        if new_head == self.food:
            self.score += 1
            self.place_food()
        else:
            self.snake.pop()

    def draw(self, oled):
        oled.fill(0)
        # Snake body
        for x, y in self.snake:
            for dx in range(CELL_SIZE):
                for dy in range(CELL_SIZE):
                    oled.pixel(x * CELL_SIZE + dx, y * CELL_SIZE + dy, 1)
        # Food
        if self.food:
            fx, fy = self.food
            for dx in range(CELL_SIZE):
                for dy in range(CELL_SIZE):
                    oled.pixel(fx * CELL_SIZE + dx, fy * CELL_SIZE + dy, 1)
        # Score
        oled.text("Score: {}".format(self.score), 0, 0, 1)
        if not self.alive:
            oled.text("GAME OVER", 30, 24, 1)
            oled.text("Pad6=RESET", 20, 36, 1)
        oled.show()

# ----- Game state -----
game = SnakeGame()
paused = False
running = True

last_pressed = -1
last_press_ms = 0
DEBOUNCE_MS = 150

last_step_ms = 0
STEP_INTERVAL = 250  # ms, adjust speed here

# ----- Init -----
set_led_color(BLUE)
oled.fill(0)
oled.text("SNAKE GAME", 20, 0)
oled.text("Pads: UDLR PS", 0, 12)
oled.text("Touch any to start", 0, 24)
oled.show()

# ----- Main loop -----
while running:
    deltas = read_pads()
    pressed = strongest_pressed(deltas)

    # Edge detection + debounce
    now_ms = utime.ticks_ms()
    pressed_event = False
    if pressed != -1 and last_pressed == -1:
        if utime.ticks_diff(now_ms, last_press_ms) > DEBOUNCE_MS:
            pressed_event = True
            last_press_ms = now_ms

    last_pressed = pressed

    if pressed_event:
        if pressed == PAD_PAUSE:
            paused = not paused
        elif pressed == PAD_RESET:
            game.reset()
            paused = False
            set_led_color(BLUE)
        else:
            game.set_dir_from_pad(pressed)

    # Auto step
    now_ms = utime.ticks_ms()
    if not paused and utime.ticks_diff(now_ms, last_step_ms) >= STEP_INTERVAL:
        last_step_ms = now_ms
        game.step()

    # Update LED
    if game.alive:
        idx = min(game.score, len(SCORE_COLORS) - 1)
        set_led_color(SCORE_COLORS[idx])
    else:
        set_led_color(RED)

    # Render
    game.draw(oled)

    utime.sleep_ms(10)