#include <avr/io.h>
#include <avr/pgmspace.h>
#include <util/delay.h>

// Pin definitions
#define I2C_SDA         PB0                   // Serial data pin
#define I2C_SCL         PB2                   // Serial clock pin
#define BUTTON_PIN      PB3                   // Button input pin

// I2C macros
#define I2C_SDA_HIGH()  DDRB &= ~(1<<I2C_SDA) // release SDA   -> pulled HIGH by resistor
#define I2C_SDA_LOW()   DDRB |=  (1<<I2C_SDA) // SDA as output -> pulled LOW  by MCU
#define I2C_SCL_HIGH()  DDRB &= ~(1<<I2C_SCL) // release SCL   -> pulled HIGH by resistor
#define I2C_SCL_LOW()   DDRB |=  (1<<I2C_SCL) // SCL as output -> pulled LOW  by MCU

// OLED definitions
#define OLED_ADDR       0x78                  // OLED write address
#define OLED_CMD_MODE   0x00                  // Set command mode
#define OLED_DAT_MODE   0x40                  // Set data mode
#define OLED_INIT_LEN   17                    // OLED initialization command length

// Font and initialization data
const uint8_t OLED_INIT_CMD[] PROGMEM = {
    0xA8, 0x1F,       // Set multiplex (HEIGHT-1): 0x1F for 128x32
    0x22, 0x00, 0x03, // Set min and max page
    0x20, 0x01,       // Set vertical memory addressing mode
    0xDA, 0x02,       // Set COM pins hardware configuration to sequential
    0x8D, 0x14,       // Enable charge pump
    0xAF,             // Switch on OLED
    0x00, 0x10, 0xB0, // Set cursor at home position
    0xA1, 0xC8        // Flip the screen
};

const uint8_t OLED_FONT[] PROGMEM = {
    0x7F, 0x41, 0x7F, // 0
    0x00, 0x00, 0x7F, // 1
    0x79, 0x49, 0x4F, // 2
    0x41, 0x49, 0x7F, // 3
    0x0F, 0x08, 0x7E, // 4
    0x4F, 0x49, 0x79, // 5
    0x7F, 0x49, 0x79, // 6
    0x03, 0x01, 0x7F, // 7
    0x7F, 0x49, 0x7F, // 8
    0x4F, 0x49, 0x7F  // 9
};

// I2C functions
void I2C_init(void) {
    DDRB &= ~((1 << I2C_SDA) | (1 << I2C_SCL)); // Set pins as input
    PORTB &= ~((1 << I2C_SDA) | (1 << I2C_SCL)); // Pull low when output
}

void I2C_write(uint8_t data) {
    for (uint8_t i = 8; i; i--) {
        I2C_SDA_LOW();
        if (data & 0x80) I2C_SDA_HIGH();
        I2C_SCL_HIGH();
        data <<= 1;
        I2C_SCL_LOW();
    }
    I2C_SDA_HIGH(); // Release SDA for ACK
    I2C_SCL_HIGH();
    asm("nop");
    I2C_SCL_LOW();
}

void I2C_start(uint8_t addr) {
    I2C_SDA_LOW();
    I2C_SCL_LOW();
    I2C_write(addr);
}

void I2C_stop(void) {
    I2C_SDA_LOW();
    I2C_SCL_HIGH();
    I2C_SDA_HIGH();
}

// OLED functions
void OLED_init(void) {
    I2C_init();
    I2C_start(OLED_ADDR);
    I2C_write(OLED_CMD_MODE);
    for (uint8_t i = 0; i < OLED_INIT_LEN; i++) {
        I2C_write(pgm_read_byte(&OLED_INIT_CMD[i]));
    }
    I2C_stop();
}

uint8_t OLED_stretch(uint8_t b) {
    b = ((b & 2) << 3) | (b & 1);
    b |= b << 1;
    b |= b << 2;
    return b;
}

void OLED_printDigit(uint8_t digit) {
    uint8_t i, j, k, b;
    uint8_t sb[4];
    digit += digit << 1;
    for (i = 8; i; i--) I2C_write(0x00);
    for (i = 3; i; i--) {
        b = pgm_read_byte(&OLED_FONT[digit++]);
        for (j = 0; j < 4; j++, b >>= 2) sb[j] = OLED_stretch(b);
        j = 4;
        if (i == 2) j = 6;
        while (j--) {
            for (k = 0; k < 4; k++) I2C_write(sb[k]);
        }
    }
}
// Set OLED brightness/contrast
void OLED_setBrightness(uint8_t brightness) {
    I2C_start(OLED_ADDR);           // Start transmission
    I2C_write(OLED_CMD_MODE);       // Set command mode
    I2C_write(0x81);                // Contrast control command
    I2C_write(brightness);          // Brightness level (0x00 to 0xFF)
    I2C_stop();                     // Stop transmission
}
// Clear the entire OLED display
void OLED_clear(void) {
    I2C_start(OLED_ADDR);
    I2C_write(OLED_DAT_MODE);
    for (uint16_t i = 0; i < 512; i++) { // Clear all pixels (128x32 OLED has 512 bytes of memory)
        I2C_write(0x00);
    }
    I2C_stop();
}

void OLED_printNumber(uint16_t number) {
    OLED_clear(); // Clear the screen before printing the new number
    I2C_start(OLED_ADDR);
    I2C_write(OLED_DAT_MODE);
    OLED_printDigit((number / 10000000) % 10);
    OLED_printDigit((number / 1000000) % 10);
    OLED_printDigit((number / 100000) % 10);
    OLED_printDigit((number / 10000) % 10);
    OLED_printDigit((number / 1000) % 10);
    OLED_printDigit((number / 100) % 10);
    OLED_printDigit((number / 10) % 10);
    OLED_printDigit(number % 10);
    I2C_stop();
}

// Debounced button press
uint8_t isButtonPressed(void) {
    if (!(PINB & (1 << BUTTON_PIN))) { // Button pressed (active low)
        _delay_ms(50); // Debounce
        if (!(PINB & (1 << BUTTON_PIN))) {
            while (!(PINB & (1 << BUTTON_PIN))); // Wait for release
            return 1;
        }
    }
    return 0;
}

// Main function
int main(void) {
    // Variables
    uint16_t counter = 0;

    // Setup
    OLED_init();
    DDRB &= ~(1 << BUTTON_PIN); // Set button pin as input
    PORTB |= (1 << BUTTON_PIN); // Enable pull-up resistor for button
OLED_setBrightness(0x10);
            OLED_printNumber(counter);

    // Loop
    while (1) {
        if (isButtonPressed()) {
            counter++;
            if (counter > 99999999) counter = 0; // Reset after 9999
            OLED_printNumber(counter);
        }
    }
}