from machine import Pin, I2C
from ssd1306 import SSD1306_I2C
import time
import random

# OLED display setup
WIDTH = 128
HEIGHT = 64
i2c = I2C(0, scl=Pin(22), sda=Pin(21))
oled = SSD1306_I2C(WIDTH, HEIGHT, i2c)

# Button and Buzzer setup
button1 = Pin(16, Pin.IN, Pin.PULL_UP)  # Clockwise direction
button2 = Pin(17, Pin.IN, Pin.PULL_UP)  # Counterclockwise direction
buzzer = Pin(18, Pin.OUT)

# Game variables
snake = [(64, 32), (64, 34), (64, 36)]  # Initial snake position
direction = 0  # 0=UP, 1=RIGHT, 2=DOWN, 3=LEFT
food = (random.randint(1, WIDTH // 2 - 1) * 2, random.randint(1, HEIGHT // 2 - 1) * 2)
score = 0
game_over = False

def generate_food():
    """Generate random food position."""
    while True:
        new_food = (random.randint(1, WIDTH // 2 - 1) * 2, random.randint(1, HEIGHT // 2 - 1) * 2)
        if new_food not in snake:  # Ensure food is not on the snake
            return new_food

def check_collision():
    """Check for wall or self-collision."""
    global game_over
    head = snake[0]
    # Wall collision
    if head[0] < 0 or head[0] >= WIDTH or head[1] < 0 or head[1] >= HEIGHT:
        game_over = True
    # Self-collision
    if head in snake[1:]:
        game_over = True

def update_snake():
    """Update the snake's position."""
    global food, score
    head = snake[0]
    if direction == 0:  # Up
        new_head = (head[0], head[1] - 2)
    elif direction == 1:  # Right
        new_head = (head[0] + 2, head[1])
    elif direction == 2:  # Down
        new_head = (head[0], head[1] + 2)
    elif direction == 3:  # Left
        new_head = (head[0] - 2, head[1])

    snake.insert(0, new_head)  # Add new head

    # Check if food is eaten
    if new_head == food:
        score += 1
        buzzer.on()
        time.sleep(0.1)
        buzzer.off()
        food = generate_food()
    else:
        snake.pop()  # Remove tail if no food eaten

def draw_rect(x, y, w, h, color):
    """Draw a rectangle using horizontal and vertical lines."""
    for i in range(w):
        oled.pixel(x + i, y, color)  # Top line
        oled.pixel(x + i, y + h - 1, color)  # Bottom line
    for i in range(h):
        oled.pixel(x, y + i, color)  # Left line
        oled.pixel(x + w - 1, y + i, color)  # Right line

def draw_game():
    """Draw the snake, food, and score."""
    oled.fill(0)  # Clear display

    # Draw snake
    for segment in snake:
        draw_rect(segment[0], segment[1], 2, 2, 1)

    # Draw food
    draw_rect(food[0], food[1], 2, 2, 1)

    # Draw score
    oled.text(f"Score: {score}", 0, 0)

    oled.show()

def game_over_screen():
    """Display Game Over screen."""
    oled.fill(0)
    oled.text("GAME OVER!", 20, 20)
    oled.text(f"Score: {score}", 20, 40)
    oled.show()
    buzzer.on()
    time.sleep(0.5)
    buzzer.off()

# Initialize game
def init_game():
    global snake, direction, food, score, game_over
    snake = [(64, 32), (64, 34), (64, 36)]
    direction = 0
    food = generate_food()
    score = 0
    game_over = False

# Main game loop
init_game()
while True:
    if game_over:
        game_over_screen()
        time.sleep(2)
        init_game()
        continue

    # Button handling
    if not button1.value():  # Button 1 pressed
        direction = (direction + 1) % 4
        time.sleep(0.2)  # Debounce
    if not button2.value():  # Button 2 pressed
        direction = (direction - 1) % 4
        time.sleep(0.2)  # Debounce

    update_snake()
    check_collision()
    draw_game()
    time.sleep(0.1)  # Adjust game speed