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

// ============================================================
// I2C (TWI) - hardware peripheral on PC4 (SDA) / PC5 (SCL)
// ============================================================

#define F_SCL 400000UL  // 400 kHz fast mode
#define TWBR_VAL ((F_CPU / F_SCL) - 16) / 2

void i2c_init(void) {
    TWSR = 0x00;          // prescaler = 1
    TWBR = (uint8_t)TWBR_VAL;
    TWCR = (1 << TWEN);   // enable TWI
}

void i2c_start(void) {
    TWCR = (1 << TWINT) | (1 << TWSTA) | (1 << TWEN);
    while (!(TWCR & (1 << TWINT)));
}

void i2c_stop(void) {
    TWCR = (1 << TWINT) | (1 << TWSTO) | (1 << TWEN);
    while (TWCR & (1 << TWSTO));
}

void i2c_write(uint8_t data) {
    TWDR = data;
    TWCR = (1 << TWINT) | (1 << TWEN);
    while (!(TWCR & (1 << TWINT)));
}

// ============================================================
// SSD1306
// ============================================================

#define SSD1306_ADDR  0x3C  // common; some modules are 0x3D

#define SSD1306_W 128
#define SSD1306_H 64
#define SSD1306_PAGES (SSD1306_H / 8)
#define SSD1306_BUF_SIZE (SSD1306_W * SSD1306_PAGES)  // 1024

#define CTRL_CMD     0x00
#define CTRL_DATA    0x40

static uint8_t framebuffer[SSD1306_BUF_SIZE];

void ssd1306_cmd(uint8_t cmd) {
    i2c_start();
    i2c_write(SSD1306_ADDR << 1);
    i2c_write(CTRL_CMD);
    i2c_write(cmd);
    i2c_stop();
}

void ssd1306_init(void) {
    _delay_ms(100);

    ssd1306_cmd(0xAE);                      // display off
    ssd1306_cmd(0xD5); ssd1306_cmd(0x80);  // clock divide
    ssd1306_cmd(0xA8); ssd1306_cmd(0x3F);  // multiplex = 64
    ssd1306_cmd(0xD3); ssd1306_cmd(0x00);  // display offset = 0
    ssd1306_cmd(0x40);                      // start line = 0
    ssd1306_cmd(0x8D); ssd1306_cmd(0x14);  // charge pump on
    ssd1306_cmd(0x20); ssd1306_cmd(0x00);  // horizontal addressing mode
    ssd1306_cmd(0xA1);                      // segment remap
    ssd1306_cmd(0xC8);                      // COM scan dir flipped
    ssd1306_cmd(0xDA); ssd1306_cmd(0x12);  // COM pins config
    ssd1306_cmd(0x81); ssd1306_cmd(0xCF);  // contrast
    ssd1306_cmd(0xD9); ssd1306_cmd(0xF1);  // pre-charge
    ssd1306_cmd(0xDB); ssd1306_cmd(0x40);  // VCOMH deselect
    ssd1306_cmd(0xA4);                      // resume from RAM
    ssd1306_cmd(0xA6);                      // normal (not inverted)
    ssd1306_cmd(0xAF);                      // display on
}

void ssd1306_display(void) {
    ssd1306_cmd(0x21); ssd1306_cmd(0); ssd1306_cmd(127);
    ssd1306_cmd(0x22); ssd1306_cmd(0); ssd1306_cmd(7);

    i2c_start();
    i2c_write(SSD1306_ADDR << 1);
    i2c_write(CTRL_DATA);
    for (uint16_t i = 0; i < SSD1306_BUF_SIZE; i++) {
        i2c_write(framebuffer[i]);
    }
    i2c_stop();
}

void ssd1306_clear(void) {
    memset(framebuffer, 0, SSD1306_BUF_SIZE);
}

void ssd1306_set_pixel(uint8_t x, uint8_t y, uint8_t on) {
    if (x >= SSD1306_W || y >= SSD1306_H) return;
    uint16_t idx = x + (y / 8) * SSD1306_W;
    uint8_t bit = 1 << (y & 7);
    if (on) framebuffer[idx] |= bit;
    else    framebuffer[idx] &= ~bit;
}

// ============================================================
// 5x7 font, ASCII 0x20-0x5A (' ' through 'Z')
// ============================================================

static const uint8_t font5x7[][5] PROGMEM = {
    {0x00,0x00,0x00,0x00,0x00}, // ' '
    {0x00,0x00,0x5F,0x00,0x00}, // '!'
    {0x00,0x07,0x00,0x07,0x00}, // '"'
    {0x14,0x7F,0x14,0x7F,0x14}, // '#'
    {0x24,0x2A,0x7F,0x2A,0x12}, // '$'
    {0x23,0x13,0x08,0x64,0x62}, // '%'
    {0x36,0x49,0x55,0x22,0x50}, // '&'
    {0x00,0x05,0x03,0x00,0x00}, // '''
    {0x00,0x1C,0x22,0x41,0x00}, // '('
    {0x00,0x41,0x22,0x1C,0x00}, // ')'
    {0x14,0x08,0x3E,0x08,0x14}, // '*'
    {0x08,0x08,0x3E,0x08,0x08}, // '+'
    {0x00,0x50,0x30,0x00,0x00}, // ','
    {0x08,0x08,0x08,0x08,0x08}, // '-'
    {0x00,0x60,0x60,0x00,0x00}, // '.'
    {0x20,0x10,0x08,0x04,0x02}, // '/'
    {0x3E,0x51,0x49,0x45,0x3E}, // '0'
    {0x00,0x42,0x7F,0x40,0x00}, // '1'
    {0x42,0x61,0x51,0x49,0x46}, // '2'
    {0x21,0x41,0x45,0x4B,0x31}, // '3'
    {0x18,0x14,0x12,0x7F,0x10}, // '4'
    {0x27,0x45,0x45,0x45,0x39}, // '5'
    {0x3C,0x4A,0x49,0x49,0x30}, // '6'
    {0x01,0x71,0x09,0x05,0x03}, // '7'
    {0x36,0x49,0x49,0x49,0x36}, // '8'
    {0x06,0x49,0x49,0x29,0x1E}, // '9'
    {0x00,0x36,0x36,0x00,0x00}, // ':'
    {0x00,0x56,0x36,0x00,0x00}, // ';'
    {0x08,0x14,0x22,0x41,0x00}, // '<'
    {0x14,0x14,0x14,0x14,0x14}, // '='
    {0x00,0x41,0x22,0x14,0x08}, // '>'
    {0x02,0x01,0x51,0x09,0x06}, // '?'
    {0x32,0x49,0x79,0x41,0x3E}, // '@'
    {0x7E,0x11,0x11,0x11,0x7E}, // 'A'
    {0x7F,0x49,0x49,0x49,0x36}, // 'B'
    {0x3E,0x41,0x41,0x41,0x22}, // 'C'
    {0x7F,0x41,0x41,0x22,0x1C}, // 'D'
    {0x7F,0x49,0x49,0x49,0x41}, // 'E'
    {0x7F,0x09,0x09,0x09,0x01}, // 'F'
    {0x3E,0x41,0x49,0x49,0x7A}, // 'G'
    {0x7F,0x08,0x08,0x08,0x7F}, // 'H'
    {0x00,0x41,0x7F,0x41,0x00}, // 'I'
    {0x20,0x40,0x41,0x3F,0x01}, // 'J'
    {0x7F,0x08,0x14,0x22,0x41}, // 'K'
    {0x7F,0x40,0x40,0x40,0x40}, // 'L'
    {0x7F,0x02,0x0C,0x02,0x7F}, // 'M'
    {0x7F,0x04,0x08,0x10,0x7F}, // 'N'
    {0x3E,0x41,0x41,0x41,0x3E}, // 'O'
    {0x7F,0x09,0x09,0x09,0x06}, // 'P'
    {0x3E,0x41,0x51,0x21,0x5E}, // 'Q'
    {0x7F,0x09,0x19,0x29,0x46}, // 'R'
    {0x46,0x49,0x49,0x49,0x31}, // 'S'
    {0x01,0x01,0x7F,0x01,0x01}, // 'T'
    {0x3F,0x40,0x40,0x40,0x3F}, // 'U'
    {0x1F,0x20,0x40,0x20,0x1F}, // 'V'
    {0x7F,0x20,0x18,0x20,0x7F}, // 'W'
    {0x63,0x14,0x08,0x14,0x63}, // 'X'
    {0x03,0x04,0x78,0x04,0x03}, // 'Y'
    {0x61,0x51,0x49,0x45,0x43}, // 'Z'
};

void ssd1306_draw_char(uint8_t x, uint8_t y, char c) {
    if (c < ' ' || c > 'Z') c = '?';
    uint8_t idx = c - ' ';
    for (uint8_t col = 0; col < 5; col++) {
        uint8_t bits = pgm_read_byte(&font5x7[idx][col]);
        for (uint8_t row = 0; row < 7; row++) {
            if (bits & (1 << row)) {
                ssd1306_set_pixel(x + col, y + row, 1);
            }
        }
    }
}

void ssd1306_draw_string(uint8_t x, uint8_t y, const char *s) {
    while (*s) {
        ssd1306_draw_char(x, y, *s++);
        x += 6;
        if (x + 5 >= SSD1306_W) break;
    }
}

// ============================================================
// UART with interrupt-driven RX ring buffer
// ============================================================

#define BAUD 9600
#define UBRR_VALUE ((F_CPU / (16UL * BAUD)) - 1)

#define RX_BUF_SIZE 64  // must be power of 2
static volatile uint8_t rx_buf[RX_BUF_SIZE];
static volatile uint8_t rx_head = 0;
static volatile uint8_t rx_tail = 0;

void uart_init(void) {
    UBRR0H = (uint8_t)(UBRR_VALUE >> 8);
    UBRR0L = (uint8_t)(UBRR_VALUE);
    // Enable RX, TX, and RX Complete interrupt
    UCSR0B = (1 << RXEN0) | (1 << TXEN0) | (1 << RXCIE0);
    UCSR0C = (1 << UCSZ01) | (1 << UCSZ00);
}

ISR(USART_RX_vect) {
    uint8_t c = UDR0;  // reading clears the flag
    uint8_t next = (rx_head + 1) & (RX_BUF_SIZE - 1);
    if (next != rx_tail) {  // drop if buffer full
        rx_buf[rx_head] = c;
        rx_head = next;
    }
}

void uart_tx(char c) {
    while (!(UCSR0A & (1 << UDRE0)));
    UDR0 = c;
}

void uart_print(const char *s) {
    while (*s) uart_tx(*s++);
}

uint8_t uart_available(void) {
    return rx_head != rx_tail;
}

char uart_rx(void) {
    while (rx_head == rx_tail);  // wait for data
    char c = rx_buf[rx_tail];
    rx_tail = (rx_tail + 1) & (RX_BUF_SIZE - 1);
    return c;
}

// ============================================================
// Text buffer for the OLED
// ============================================================

#define MAX_LINES 8        // 8 lines * 8px = 64px tall
#define MAX_COLS  21       // 128px / 6px per char

static char lines[MAX_LINES][MAX_COLS + 1];
static uint8_t cur_line = 0;
static uint8_t cur_col  = 0;

void redraw_oled(void) {
    ssd1306_clear();
    for (uint8_t i = 0; i < MAX_LINES; i++) {
        ssd1306_draw_string(0, i * 8, lines[i]);
    }
    // Cursor: vertical bar
    for (uint8_t y = 0; y < 7; y++) {
        ssd1306_set_pixel(cur_col * 6, cur_line * 8 + y, 1);
    }
    ssd1306_display();
}

void scroll_up(void) {
    for (uint8_t i = 0; i < MAX_LINES - 1; i++) {
        memcpy(lines[i], lines[i + 1], MAX_COLS + 1);
    }
    memset(lines[MAX_LINES - 1], 0, MAX_COLS + 1);
}

void handle_char(char c) {
    if (c == '\r' || c == '\n') {
        cur_col = 0;
        if (cur_line < MAX_LINES - 1) {
            cur_line++;
        } else {
            scroll_up();
        }
    } else if (c == 0x08 || c == 0x7F) {
        if (cur_col > 0) {
            cur_col--;
            lines[cur_line][cur_col] = ' ';
        }
    } else if (c >= 'a' && c <= 'z') {
        // Convert lowercase to uppercase (font is uppercase only)
        handle_char(c - 32);
        return;
    } else if (c >= ' ' && c <= 'Z') {
        if (cur_col < MAX_COLS) {
            lines[cur_line][cur_col] = c;
            cur_col++;
        }
        if (cur_col >= MAX_COLS) {
            cur_col = 0;
            if (cur_line < MAX_LINES - 1) {
                cur_line++;
            } else {
                scroll_up();
            }
        }
    }
    // Other characters: ignore
}

// Process one character from UART (echo + buffer update)
void process_char(char c) {
    if (c == 0x7F || c == 0x08) {
        uart_tx(0x08); uart_tx(' '); uart_tx(0x08);
    } else {
        uart_tx(c);
        if (c == '\r') uart_tx('\n');
    }
    handle_char(c);
}

// ============================================================
// Main
// ============================================================

int main(void) {
    i2c_init();
    ssd1306_init();
    ssd1306_clear();
    uart_init();

    memset(lines, 0, sizeof(lines));

    sei();  // enable global interrupts

    uart_print("OLED terminal ready. Type away!\r\n");
    redraw_oled();

    while (1) {
        if (uart_available()) {
            // Drain ALL pending characters before doing the expensive redraw
            while (uart_available()) {
                process_char(uart_rx());
            }
            redraw_oled();
        }
    }
}