# Import necessary modules
from machine import Pin, I2C, ADC
from time import sleep_ms
from ELB import SSD1306_I2C, SnakeGame

# Initialize OLED display
i2c = I2C(0, scl=Pin(1), sda=Pin(0), freq=400000)
oled = SSD1306_I2C(128, 64, i2c)

# Joystick inputs
joystick_x = ADC(Pin(27))
joystick_y = ADC(Pin(26))

# Joystick center and movement range
ADC_CENTER = 32767
DEAD_ZONE = 4000

# Create snake game object
snake = SnakeGame(oled, button_pin=4)

# Read joystick and return a direction
def get_direction():
    # Read joystick and return snake direction
    x = joystick_x.read_u16()
    y = joystick_y.read_u16()

    # Move right or left
    if x < (ADC_CENTER - DEAD_ZONE):
        return snake.RIGHT
    if x > (ADC_CENTER + DEAD_ZONE):
        return snake.LEFT
    
    # Move up or down
    if y < (ADC_CENTER - DEAD_ZONE):
        return snake.DOWN
    if y > (ADC_CENTER + DEAD_ZONE):
        return snake.UP
    
    # No direction if joystick is in dead zone
    return None

# -------- Main Game Loop --------
while True:
    # Show start screen
    snake.start_game_message()
    
    # Wait for button press
    snake.wait_for_button()
    
    # Start new game
    snake.new_game()
    
    # Game loop - play until the snake crashes
    while snake.alive:
        # Get joystick input
        new_dir = get_direction()
        if new_dir:
            snake.update_direction(new_dir)
        
        # Move snake and check if game is still active
        if not snake.check_collision_and_move():
            break  # Oops! Snake crashed, exit game loop
            
        snake.draw()
        sleep_ms(10)
    
    # Show game over screen
    snake.game_over_message()
    
    # Wait for button press before restarting
    snake.wait_for_button()