from machine import Pin, I2C
from ssd1306 import SSD1306_I2C
import utime

UP_BUTTON_PIN = 2
DOWN_BUTTON_PIN = 3
PADDLE_RATE = 64
BALL_RATE = 16
PADDLE_HEIGHT = 12
SCORE_LIMIT = 9

pix_res_x = 128
pix_res_y = 64

i2c_dev = I2C(1, scl=Pin(27), sda=Pin(26), freq=200000)
display = SSD1306_I2C(pix_res_x, pix_res_y, i2c_dev)

up_button = machine.Pin(UP_BUTTON_PIN, machine.Pin.IN, machine.Pin.PULL_UP)
down_button = machine.Pin(DOWN_BUTTON_PIN, machine.Pin.IN, machine.Pin.PULL_UP)

game_state = {
    'game_over': False,
    'win': False,
    'player_score': 0,
    'mcu_score': 0,
    'ball': {'x': 53, 'y': 26, 'dir_x': 1, 'dir_y': 1},
    'paddles': {'mcu': {'x': 12, 'y': 16}, 'player': {'x': 115, 'y': 16}},
    'timing': {'ball_update': utime.ticks_ms(), 'paddle_update': utime.ticks_ms()}
}

def draw_court():
    display.fill(0)
    display.rect(0, 0, 128, 54, 1)

def update_score():
    display.fill_rect(0, 56, 128, 8, 0)
    display.text(str(game_state['mcu_score']), 0, 56, 1)
    display.text(str(game_state['player_score']), 122, 56, 1)

def draw_paddle(x, y, color):
    display.fill_rect(x, y, 2, PADDLE_HEIGHT, color)

def draw_ball(x, y, color):
    display.pixel(x, y, color)

def update_ball():
    ball = game_state['ball']
    new_x = ball['x'] + ball['dir_x']
    new_y = ball['y'] + ball['dir_y']

    collision = check_collision(new_x, new_y)

    if collision == 'vertical':
        ball['dir_x'] = -ball['dir_x']
        new_x += 2 * ball['dir_x']
        if new_x < 64:
            game_state['player_score'] += 1
        else:
            game_state['mcu_score'] += 1

        if game_state['player_score'] == SCORE_LIMIT or game_state['mcu_score'] == SCORE_LIMIT:
            game_state['win'] = game_state['player_score'] > game_state['mcu_score']
            game_state['game_over'] = True

    if collision == 'horizontal':
        ball['dir_y'] = -ball['dir_y']
        new_y += 2 * ball['dir_y']

    draw_ball(ball['x'], ball['y'], 0)
    draw_ball(new_x, new_y, 1)
    ball['x'], ball['y'] = new_x, new_y

def check_collision(new_x, new_y):
    mcu_paddle = game_state['paddles']['mcu']
    player_paddle = game_state['paddles']['player']

    if (new_x == 0 or new_x == 127) or (
            (new_x == mcu_paddle['x'] and mcu_paddle['y'] <= new_y <= mcu_paddle['y'] + PADDLE_HEIGHT) or
            (new_x == player_paddle['x'] and player_paddle['y'] <= new_y <= player_paddle['y'] + PADDLE_HEIGHT)):
        return 'vertical'
    elif new_y == 0 or new_y == 53:
        return 'horizontal'
    else:
        return None

def update_paddle_positions():
    mcu_paddle = game_state['paddles']['mcu']
    player_paddle = game_state['paddles']['player']
    half_paddle = PADDLE_HEIGHT // 2

    # Update MCU paddle position
    draw_paddle(mcu_paddle['x'], mcu_paddle['y'], 0)

    if mcu_paddle['y'] + half_paddle > game_state['ball']['y']:
        mcu_paddle['y'] -= 1
    elif mcu_paddle['y'] + half_paddle < game_state['ball']['y']:
        mcu_paddle['y'] += 1

    # Limit MCU paddle position
    mcu_paddle['y'] = max(min(mcu_paddle['y'], 53 - PADDLE_HEIGHT), 1)
    draw_paddle(mcu_paddle['x'], mcu_paddle['y'], 1)

    # Update player paddle position
    draw_paddle(player_paddle['x'], player_paddle['y'], 0)

    if up_button.value() == 0:
        player_paddle['y'] -= 1
    if down_button.value() == 0:
        player_paddle['y'] += 1

    # Limit player paddle position
    player_paddle['y'] = max(min(player_paddle['y'], 53 - PADDLE_HEIGHT), 1)
    draw_paddle(player_paddle['x'], player_paddle['y'], 1)

def check_game_over():
    if game_state['player_score'] == SCORE_LIMIT or game_state['mcu_score'] == SCORE_LIMIT:
        game_state['win'] = game_state['player_score'] > game_state['mcu_score']
        game_state['game_over'] = True
        display_game_over()
        reset_game()

def update_ball():
    ball = game_state['ball']
    new_x = ball['x'] + ball['dir_x']
    new_y = ball['y'] + ball['dir_y']

    collision = check_collision(new_x, new_y)
    if collision == 'vertical':
        ball['dir_x'] = -ball['dir_x']
        new_x += ball['dir_x'] * 2
        update_score()
    elif collision == 'horizontal':
        ball['dir_y'] = -ball['dir_y']
        new_y += ball['dir_y'] * 2

    display.pixel(ball['x'], ball['y'], 0)
    display.pixel(new_x, new_y, 1)
    ball['x'] = new_x
    ball['y'] = new_y

def update_score():
    mcu_score = game_state['mcu_score']
    player_score = game_state['player_score']
    display.text(str(mcu_score), 0, 56, 0)
    display.text(str(player_score), 122, 56, 0)
    display.text(str(mcu_score), 0, 56, 1)
    display.text(str(player_score), 122, 56, 1)

def display_game_over():
    win = game_state['score']['player'] > game_state['score']['mcu']
    text = "YOU WIN!!" if win else "YOU LOSE!"
    display.fill(0)
    display.text(text, 40, 28, 1)
    display.show()
    utime.sleep(5)

def reset_game():
    game_state['ball']['x'] = 53
    game_state['ball']['y'] = 26
    game_state['ball']['dir_x'] = 1
    game_state['ball']['dir_y'] = 1
    game_state['paddles']['mcu']['y'] = 16
    game_state['paddles']['player']['y'] = 16
    game_state['score']['mcu'] = 0
    game_state['score']['player'] = 0
    game_state['game_over'] = False
    draw_court()

def main():

    while True:
        now = utime.ticks_ms()

        if now - game_state['timing']['ball_update'] > BALL_RATE:
            game_state['timing']['ball_update'] = now
            update_ball()
            check_game_over()

        if now - game_state['timing']['paddle_update'] > PADDLE_RATE:
            game_state['timing']['paddle_update'] = now
            update_paddle_positions()

        display.show()
        update_score()
        utime.sleep_ms(1)

if __name__ == '__main__':
    main()