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

# Joystick Initialization
def init_joystick(x_pin, y_pin, btn_pin):
    # Define the ADC for X and Y axes
    adc_x = ADC(Pin(x_pin))
    adc_y = ADC(Pin(y_pin))

    # On the ESP32, ADC has a range of 0-4095 for 0-3.3V
    adc_x.atten(ADC.ATTN_11DB)
    adc_y.atten(ADC.ATTN_11DB)
    
    # Set up the button pin
    btn = Pin(btn_pin, Pin.IN, Pin.PULL_UP)
    return adc_x, adc_y, btn

# OLED Initialization
def init_oled(sda, scl, width, height):
    i2c = I2C(sda=Pin(sda), scl=Pin(scl))
    oled = ssd1306.SSD1306_I2C(width, height, i2c)
    return oled

# Display Data on OLED
def display_data(oled, x, y, btn_state):
    oled.fill(0)  # Clear the display
    oled.text('X: {}'.format(x), 0, 0)
    oled.text('Y: {}'.format(y), 0, 10)
    oled.text('Button: {}'.format('Pressed' if btn_state == 0 else 'Released'), 0, 20)
    oled.show()

# Read Joystick Values
def read_joystick(adc_x, adc_y, btn):
    # Read the values
    x_value = adc_x.read()
    y_value = adc_y.read()
    btn_state = btn.value()
    return x_value, y_value, btn_state

# Initialize joystick and OLED
adc_x, adc_y, btn = init_joystick(34, 35, 32)
oled = init_oled(21, 22, 128, 32)  # Assuming a 128x32 OLED display

# --- Game Logic ---

class PongGame:
    def __init__(self, width, height):
        self.width = width
        self.height = height
        
        self.paddle_width = 4
        self.paddle_height = 16
        self.ball_size = 4
        
        self.player_pos = self.height // 2
        self.ai_pos = self.height // 2
        self.ball_pos = [self.width // 2, self.height // 2]
        self.ball_vel = [2, 2]

    def update(self):
        # Update ball position
        self.ball_pos[0] += self.ball_vel[0]
        self.ball_pos[1] += self.ball_vel[1]

        # Ball collision with top and bottom
        if self.ball_pos[1] <= 0 or self.ball_pos[1] >= self.height:
            self.ball_vel[1] = -self.ball_vel[1]
            
        # Ball collision with paddles
        if (self.ball_pos[0] <= self.paddle_width and self.player_pos <= self.ball_pos[1] <= self.player_pos + self.paddle_height) or \
           (self.ball_pos[0] >= self.width - self.paddle_width and self.ai_pos <= self.ball_pos[1] <= self.ai_pos + self.paddle_height):
            self.ball_vel[0] = -self.ball_vel[0]

        # Update AI paddle position
        if self.ball_pos[1] > self.ai_pos + self.paddle_height / 2:
            self.ai_pos += 2
        elif self.ball_pos[1] < self.ai_pos + self.paddle_height / 2:
            self.ai_pos -= 2

        # Ensure paddles stay within screen
        self.player_pos = max(min(self.player_pos, self.height - self.paddle_height), 0)
        self.ai_pos = max(min(self.ai_pos, self.height - self.paddle_height), 0)

    def draw(self, oled):
        oled.fill(0)
        # Draw paddles
        oled.fill_rect(0, self.player_pos, self.paddle_width, self.paddle_height, 1)
        oled.fill_rect(self.width - self.paddle_width, self.ai_pos, self.paddle_width, self.paddle_height, 1)
        # Draw ball
        oled.fill_rect(self.ball_pos[0], self.ball_pos[1], self.ball_size, self.ball_size, 1)
        oled.show()

# --- Main Code ---

# Initialize
adc_x, adc_y, btn = init_joystick(34, 35, 32)
oled = init_oled(21, 22, 128, 64)

game = PongGame(128, 64)

while True:
    x, y, btn_state = read_joystick(adc_x, adc_y, btn)
    
    # Update player paddle position based on joystick
    game.player_pos = int((y / 4096) * (64 - game.paddle_height))
    
    game.update()
    game.draw(oled)
    time.sleep(0.05)