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

i2c = I2C(0, scl=Pin(22), sda=Pin(21))
display = ssd1306.SSD1306_I2C(128, 64, i2c)

button_up = Pin(32, Pin.IN, Pin.PULL_UP)
button_down = Pin(33, Pin.IN, Pin.PULL_UP)
button_left = Pin(34, Pin.IN, Pin.PULL_UP)
button_right = Pin(35, Pin.IN, Pin.PULL_UP)

class Snake:
    def __init__(self):
        self.reset()
    
    def reset(self):
        self.body = [[64, 32]]  
        self.direction = [1, 0] 
        self.length = 3
        self.food = self.generate_food()
        self.game_over = False
    
    def generate_food(self):
        while True:
            food = [random.randint(1, 126), random.randint(1, 62)]
            if food not in self.body:
                return food
    
    def move(self):
        if self.game_over:
            return
        
      
        head = self.body[0].copy()
        head[0] += self.direction[0] * 2
        head[1] += self.direction[1] * 2
        
        
        if head[0] <= 0 or head[0] >= 127 or head[1] <= 0 or head[1] >= 63:
            self.game_over = True
            return
        
        
        if head in self.body:
            self.game_over = True
            return
        
       
        self.body.insert(0, head)
        
        
        if abs(head[0] - self.food[0]) < 3 and abs(head[1] - self.food[1]) < 3:
            self.length += 1
            self.food = self.generate_food()
        
        
        while len(self.body) > self.length:
            self.body.pop()
    
    def change_direction(self, new_dir):
        
        if (new_dir[0] != -self.direction[0] or new_dir[1] != -self.direction[1]):
            self.direction = new_dir
    
    def draw(self):
        display.fill(0)  
        
       
        for segment in self.body:
            display.fill_rect(segment[0], segment[1], 2, 2, 1)
        
      
        display.fill_rect(self.food[0], self.food[1], 2, 2, 1)
        
      
        display.text(f'Score: {self.length-3}', 0, 0, 1)
        
        if self.game_over:
            display.text('GAME OVER', 30, 28, 1)
            display.text('Press to restart', 15, 40, 1)
        
        display.show()

snake = Snake()

while True:
    
    if not button_up.value():
        snake.change_direction([0, -1])
    elif not button_down.value():
        snake.change_direction([0, 1])
    elif not button_left.value():
        snake.change_direction([-1, 0])
    elif not button_right.value():
        snake.change_direction([1, 0])
    

    if snake.game_over and (not button_up.value() or not button_down.value() or 
                          not button_left.value() or not button_right.value()):
        snake.reset()
    
 
    snake.move()
    snake.draw()
    
    time.sleep(0.1)