#include "pico/stdlib.h"
#include "hardware/i2c.h"
#include "hardware/gpio.h"
#include <string.h>
#include <stdlib.h>
#include <stdio.h>

// OLED and I2C Configuration
#define OLED_WIDTH 128
#define OLED_HEIGHT 64
#define OLED_ADDR 0x3C
#define I2C_PORT i2c0
#define SDA_PIN 4
#define SCL_PIN 5

// Keypad Configuration
#define ROWS 4
#define COLS 4
const uint row_pins[ROWS] = {6, 7, 8, 9};    // R1-R4
const uint col_pins[COLS] = {10, 11, 12, 13}; // C1-C4

// Physical key layout (what's on your keypad)
const char physical_keymap[ROWS][COLS] = {
    {'1', '2', '3', 'A'},
    {'4', '5', '6', 'B'},
    {'7', '8', '9', 'C'},
    {'*', '0', '#', 'D'}
};

// Logical mapping (how we interpret them)
const char logical_keymap[ROWS][COLS] = {
    {'1', '2', '3', '/'},  // A → /
    {'4', '5', '6', '*'},  // B → *
    {'7', '8', '9', '-'},  // C → -
    {'C', '0', '=', '+'}   // *→C, #→=, D→+
};

// 5x7 Font (ASCII 32-127)
static const uint8_t font_5x7[] = {
    0x00,0x00,0x00,0x00,0x00, // space
    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, // )
    0x08,0x2A,0x1C,0x2A,0x08, // *
    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
    0x00,0x7F,0x41,0x41,0x00, // [
    0x02,0x04,0x08,0x10,0x20, // backslash
    0x00,0x41,0x41,0x7F,0x00, // ]
    0x04,0x02,0x01,0x02,0x04, // ^
    0x40,0x40,0x40,0x40,0x40, // _
    0x00,0x01,0x02,0x04,0x00, // `
    0x20,0x54,0x54,0x54,0x78, // a
    0x7F,0x48,0x44,0x44,0x38, // b
    0x38,0x44,0x44,0x44,0x20, // c
    0x38,0x44,0x44,0x48,0x7F, // d
    0x38,0x54,0x54,0x54,0x18, // e
    0x08,0x7E,0x09,0x01,0x02, // f
    0x08,0x14,0x54,0x54,0x3C, // g
    0x7F,0x08,0x04,0x04,0x78, // h
    0x00,0x44,0x7D,0x40,0x00, // i
    0x20,0x40,0x44,0x3D,0x00, // j
    0x00,0x7F,0x10,0x28,0x44, // k
    0x00,0x41,0x7F,0x40,0x00, // l
    0x7C,0x04,0x18,0x04,0x78, // m
    0x7C,0x08,0x04,0x04,0x78, // n
    0x38,0x44,0x44,0x44,0x38, // o
    0x7C,0x14,0x14,0x14,0x08, // p
    0x08,0x14,0x14,0x18,0x7C, // q
    0x7C,0x08,0x04,0x04,0x08, // r
    0x48,0x54,0x54,0x54,0x20, // s
    0x04,0x3F,0x44,0x40,0x20, // t
    0x3C,0x40,0x40,0x20,0x7C, // u
    0x1C,0x20,0x40,0x20,0x1C, // v
    0x3C,0x40,0x30,0x40,0x3C, // w
    0x44,0x28,0x10,0x28,0x44, // x
    0x0C,0x50,0x50,0x50,0x3C, // y
    0x44,0x64,0x54,0x4C,0x44, // z
    0x00,0x08,0x36,0x41,0x00, // {
    0x00,0x00,0x7F,0x00,0x00, // |
    0x00,0x41,0x36,0x08,0x00, // }
    0x08,0x08,0x2A,0x1C,0x08, // →
    0x08,0x1C,0x2A,0x08,0x08  // ←
};

// Display State
static uint8_t oled_buffer[OLED_WIDTH * OLED_HEIGHT / 8];
static uint8_t current_line = 0;

// Calculator State
static char input_buffer[32] = {0};
static uint8_t input_pos = 0;
static float result = 0;
static bool new_input = true;

void oled_command(uint8_t cmd) {
    uint8_t buf[2] = {0x00, cmd};
    i2c_write_blocking(I2C_PORT, OLED_ADDR, buf, 2, false);
}

void oled_init() {
    sleep_ms(100);
    oled_command(0xAE); // Display off
    
    oled_command(0x20); oled_command(0x00); // Horizontal addressing
    oled_command(0xB0); // Page addressing
    
    oled_command(0xC8); // COM scan direction
    oled_command(0x00); oled_command(0x10); // Column addresses
    oled_command(0x40); // Start line
    oled_command(0x81); oled_command(0xFF); // Contrast
    
    oled_command(0xA1); // Segment remap
    oled_command(0xA6); // Normal display
    oled_command(0xA8); oled_command(0x3F); // Multiplex ratio
    oled_command(0xA4); // Entire display on
    
    oled_command(0xD3); oled_command(0x00); // Display offset
    oled_command(0xD5); oled_command(0xF0); // Display clock
    oled_command(0xD9); oled_command(0x22); // Pre-charge
    oled_command(0xDA); oled_command(0x12); // COM pins
    oled_command(0xDB); oled_command(0x20); // VCOMH
    oled_command(0x8D); oled_command(0x14); // Charge pump
    
    oled_command(0xAF); // Display on
}

void oled_clear() {
    memset(oled_buffer, 0, sizeof(oled_buffer));
    current_line = 0;
}

void oled_update() {
    for (uint8_t page = 0; page < 8; page++) {
        oled_command(0xB0 | page);
        oled_command(0x00);
        oled_command(0x10);
        
        uint8_t buf[OLED_WIDTH + 1];
        buf[0] = 0x40;
        memcpy(&buf[1], &oled_buffer[page * OLED_WIDTH], OLED_WIDTH);
        i2c_write_blocking(I2C_PORT, OLED_ADDR, buf, OLED_WIDTH + 1, false);
    }
}

void oled_draw_char(uint8_t x, uint8_t y, char c) {
    if (c < 32 || c > 127) c = ' ';
    const uint8_t *char_data = &font_5x7[(c - 32) * 5];
    
    for (uint8_t col = 0; col < 5; col++) {
        uint8_t col_data = char_data[col];
        for (uint8_t bit = 0; bit < 7; bit++) {
            if (col_data & (1 << bit)) {
                uint16_t index = x + col + ((y + bit) / 8) * OLED_WIDTH;
                if (index < sizeof(oled_buffer)) {
                    oled_buffer[index] |= (1 << ((y + bit) % 8));
                }
            }
        }
    }
}

void oled_draw_string(uint8_t x, uint8_t y, const char *str) {
    while (*str) {
        oled_draw_char(x, y, *str++);
        x += 6;
        if (x + 5 >= OLED_WIDTH) {
            x = 0;
            y += 8;
        }
    }
}

void oled_new_line() {
    if (current_line < 7) {
        current_line++;
    } else {
        memmove(oled_buffer, oled_buffer + OLED_WIDTH, OLED_WIDTH * 7);
        memset(oled_buffer + OLED_WIDTH * 7, 0, OLED_WIDTH);
    }
}

void keypad_init() {
    for (int r = 0; r < ROWS; r++) {
        gpio_init(row_pins[r]);
        gpio_set_dir(row_pins[r], GPIO_OUT);
        gpio_put(row_pins[r], 1);
    }
    for (int c = 0; c < COLS; c++) {
        gpio_init(col_pins[c]);
        gpio_set_dir(col_pins[c], GPIO_IN);
        gpio_pull_up(col_pins[c]);
    }
}

char keypad_get_key() {
    for (int r = 0; r < ROWS; r++) {
        gpio_put(row_pins[r], 0);
        for (int c = 0; c < COLS; c++) {
            if (!gpio_get(col_pins[c])) {
                // Wait for key release with debouncing
                while(!gpio_get(col_pins[c])) {
                    sleep_ms(10);
                }
                gpio_put(row_pins[r], 1);
                return logical_keymap[r][c];
            }
        }
        gpio_put(row_pins[r], 1);
        sleep_ms(1);
    }
    return 0;
}

void process_input(char key) {
    if (key == 'C') {  // Physical '*' key
        // Clear everything
        input_pos = 0;
        input_buffer[0] = '\0';
        result = 0;
        new_input = true;
        oled_clear();
        oled_update();
        return;
    } 
    else if (key == '=' && input_pos > 0) {
        // Evaluate expression
        char* endptr;
        float a = strtof(input_buffer, &endptr);
        char op = *endptr;
        float b = strtof(endptr + 1, NULL);
        
        switch (op) {
            case '+': result = a + b; break;
            case '-': result = a - b; break;
            case '*': result = a * b; break;
            case '/': result = a / b; break;
            default: result = a; break;
        }
        
        // Display result
        oled_new_line();
        char result_str[16];
        snprintf(result_str, sizeof(result_str), "= %.2f", result);
        oled_draw_string(0, current_line * 8, result_str);
        oled_update();
        
        // Reset for new input
        input_pos = 0;
        input_buffer[0] = '\0';
        new_input = true;
    }
    else if (input_pos < sizeof(input_buffer) - 1) {
        // Add to input buffer
        input_buffer[input_pos++] = key;
        input_buffer[input_pos] = '\0';
        
        // Display input
        if (new_input) {
            oled_new_line();
            new_input = false;
        }
        oled_draw_char((input_pos - 1) * 6, current_line * 8, key);
        oled_update();
    }
}

int main() {
    // Initialize hardware
    i2c_init(I2C_PORT, 400 * 1000);
    gpio_set_function(SDA_PIN, GPIO_FUNC_I2C);
    gpio_set_function(SCL_PIN, GPIO_FUNC_I2C);
    gpio_pull_up(SDA_PIN);
    gpio_pull_up(SCL_PIN);
    
    keypad_init();
    oled_init();
    oled_clear();
    
    // Initial display
    oled_draw_string(0, 0, "Calculator");
    oled_update();
    sleep_ms(1000);
    oled_clear();

    // Main loop
    while (1) {
        char key = keypad_get_key();
        if (key) {
            process_input(key);
        }
        sleep_ms(10);
    }
}