#include <U8g2lib.h>
#include <Wire.h>

#define UP_BUTTON 2
#define DOWN_BUTTON 4

const unsigned long PADDLE_RATE = 33;
const unsigned long BALL_RATE = 16;
const uint8_t PADDLE_HEIGHT = 24;

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64

U8G2_SSD1306_128X64_NONAME_F_HW_I2C u8g2(U8G2_R0, /* reset=*/U8X8_PIN_NONE);

uint8_t ball_x = 64, ball_y = 32;
uint8_t ball_dir_x = 1, ball_dir_y = 1;
unsigned long ball_update;

unsigned long paddle_update;
const uint8_t CPU_X = 12;
uint8_t cpu_y = 16;

const uint8_t PLAYER_X = 115;
uint8_t player_y = 16;

void setup() {
    u8g2.begin();
    unsigned long start = millis();

    pinMode(UP_BUTTON, INPUT);
    pinMode(DOWN_BUTTON, INPUT);
    digitalWrite(UP_BUTTON, 1);
    digitalWrite(DOWN_BUTTON, 1);
    u8g2.clearBuffer();
    drawCourt();

    while (millis() - start < 2000);

    u8g2.sendBuffer();  // Send buffer to update the display
    ball_update = millis();
    paddle_update = ball_update;
}

void drawCourt() {
    // Draw the frame
    u8g2.drawFrame(0, 0, SCREEN_WIDTH, SCREEN_HEIGHT);
}

void loop() {
    drawCourt();
    bool update = false;
    unsigned long time = millis();

    static bool up_state = false;
    static bool down_state = false;

    up_state |= (digitalRead(UP_BUTTON) == LOW);
    down_state |= (digitalRead(DOWN_BUTTON) == LOW);

    if (time > ball_update) {
        uint8_t new_x = ball_x + ball_dir_x;
        uint8_t new_y = ball_y + ball_dir_y;

        if (new_x == 0 || new_x == 127) {
            ball_dir_x = -ball_dir_x;
            new_x += ball_dir_x + ball_dir_x;
        }

        if (new_y == 0 || new_y == 63) {
            ball_dir_y = -ball_dir_y;
            new_y += ball_dir_y + ball_dir_y;
        }

        if (new_x == CPU_X && new_y >= cpu_y && new_y <= cpu_y + PADDLE_HEIGHT) {
            ball_dir_x = -ball_dir_x;
            new_x += ball_dir_x + ball_dir_x;
        }

        if (new_x == PLAYER_X && new_y >= player_y && new_y <= player_y + PADDLE_HEIGHT) {
            ball_dir_x = -ball_dir_x;
            new_x += ball_dir_x + ball_dir_x;
        }

        // Draw the ball at the new position
        u8g2.drawPixel(ball_x, ball_y);

        // Update the ball position
        ball_x = new_x;
        ball_y = new_y;

        ball_update += BALL_RATE;

        update = true;
    }

    if (time > paddle_update) {
        // Update the paddles
        const uint8_t half_paddle = PADDLE_HEIGHT >> 1;
        if (cpu_y + half_paddle > ball_y) {
            cpu_y -= 1;
        }
        if (cpu_y + half_paddle < ball_y) {
            cpu_y += 1;
        }
        if (cpu_y < 1) cpu_y = 1;
        if (cpu_y + PADDLE_HEIGHT > 63) cpu_y = 63 - PADDLE_HEIGHT;

        if (up_state) {
            player_y = player_y - 1;
        }
        if (down_state) {
            player_y = player_y + 1;
        }
        up_state = down_state = false;
        if (player_y < 1) player_y = 1;
        if (player_y + PADDLE_HEIGHT > 63) player_y = 63 - PADDLE_HEIGHT;

        // Draw the paddles at the new positions
        u8g2.drawBox(CPU_X, cpu_y, 2, PADDLE_HEIGHT);
        u8g2.drawBox(PLAYER_X, player_y, 2, PADDLE_HEIGHT);

        update = true;
        paddle_update += PADDLE_RATE;
    }

    if (update) {
        // Send buffer to update the display
        u8g2.sendBuffer();
        delayMicroseconds(200);
        u8g2.clearBuffer();
    }
}