#include "oled_display.h"
#include "font_5x7.h"
#include <string.h>
#include <stdlib.h>
#include "esp_log.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"

static const char *TAG = "OLED_DISPLAY";

/* Internal frame buffer: 72 columns x 5 pages = 360 bytes */
static uint8_t s_frame_buffer[OLED_BUFFER_SIZE];
static uint8_t s_column_offset = OLED_COLUMN_OFFSET;

/* Helper: Send command to OLED */
static esp_err_t oled_send_cmd(uint8_t cmd)
{
    i2c_cmd_handle_t link = i2c_cmd_link_create();
    i2c_master_start(link);
    i2c_master_write_byte(link, (OLED_I2C_ADDR << 1) | I2C_MASTER_WRITE, true);
    i2c_master_write_byte(link, 0x00, true); // Control byte: 0x00 = Command
    i2c_master_write_byte(link, cmd, true);
    i2c_master_stop(link);
    esp_err_t ret = i2c_master_cmd_begin(OLED_I2C_PORT, link, pdMS_TO_TICKS(50));
    i2c_cmd_link_delete(link);
    return ret;
}

/* Helper: Send multiple commands to OLED */
static esp_err_t oled_send_cmds(const uint8_t *cmds, size_t len)
{
    i2c_cmd_handle_t link = i2c_cmd_link_create();
    i2c_master_start(link);
    i2c_master_write_byte(link, (OLED_I2C_ADDR << 1) | I2C_MASTER_WRITE, true);
    i2c_master_write_byte(link, 0x00, true); // Control byte: 0x00 = Command
    i2c_master_write(link, cmds, len, true);
    i2c_master_stop(link);
    esp_err_t ret = i2c_master_cmd_begin(OLED_I2C_PORT, link, pdMS_TO_TICKS(100));
    i2c_cmd_link_delete(link);
    return ret;
}

esp_err_t oled_init(gpio_num_t sda, gpio_num_t scl, uint8_t col_offset)
{
    s_column_offset = col_offset;

    // 1. Configure I2C bus
    i2c_config_t conf = {
        .mode = I2C_MODE_MASTER,
        .sda_io_num = sda,
        .scl_io_num = scl,
        .sda_pullup_en = GPIO_PULLUP_ENABLE,
        .scl_pullup_en = GPIO_PULLUP_ENABLE,
        .master.clk_speed = OLED_I2C_FREQ_HZ,
    };

    esp_err_t err = i2c_param_config(OLED_I2C_PORT, &conf);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "i2c_param_config failed: %s", esp_err_to_name(err));
        return err;
    }

    err = i2c_driver_install(OLED_I2C_PORT, conf.mode, 0, 0, 0);
    if (err != ESP_OK && err != ESP_ERR_INVALID_STATE) {
        ESP_LOGE(TAG, "i2c_driver_install failed: %s", esp_err_to_name(err));
        return err;
    }

    vTaskDelay(pdMS_TO_TICKS(50)); // Allow display power to stabilize

    // 2. Display initialization sequence for 72x40 OLED (SH1106 / SSD1306 compatible)
    const uint8_t init_cmds[] = {
        0xAE,             // Display OFF
        0xD5, 0x80,       // Set Display Clock Divide Ratio / Osc Frequency
        0xA8, 0x27,       // Set Multiplex Ratio (40 - 1 = 39 = 0x27)
        0xD3, 0x00,       // Set Display Offset to 0
        0x40,             // Set Display Start Line to 0
        0x8D, 0x14,       // Enable Charge Pump
        0x20, 0x02,       // Set Memory Addressing Mode to Page mode (compatible with SH1106)
        0xA1,             // Set Segment Re-map (A0/A1)
        0xC8,             // Set COM Output Scan Direction (C0/C8)
        0xDA, 0x12,       // Set COM Pins hardware config
        0x81, 0xAF,       // Set Contrast Control (0xAF)
        0xD9, 0xF1,       // Set Pre-charge Period
        0xDB, 0x20,       // Set VCOMH Deselect Level
        0xA4,             // Resume to RAM content (Entire Display ON disabled)
        0xA6,             // Normal display (non-inverted)
        0xAF              // Display ON
    };

    err = oled_send_cmds(init_cmds, sizeof(init_cmds));
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "OLED init commands failed: %s (Check I2C connections)", esp_err_to_name(err));
        return err;
    }

    oled_clear();
    oled_update();

    ESP_LOGI(TAG, "OLED initialized (72x40, offset: %d, SDA: %d, SCL: %d)", 
             s_column_offset, (int)sda, (int)scl);
    return ESP_OK;
}

void oled_clear(void)
{
    memset(s_frame_buffer, 0x00, sizeof(s_frame_buffer));
}

void oled_fill(void)
{
    memset(s_frame_buffer, 0xFF, sizeof(s_frame_buffer));
}

esp_err_t oled_update(void)
{
    for (uint8_t page = 0; page < OLED_PAGES; page++) {
        // Page address + column address setup
        uint8_t page_cmds[] = {
            (uint8_t)(0xB0 | page),
            (uint8_t)(0x00 | (s_column_offset & 0x0F)),
            (uint8_t)(0x10 | ((s_column_offset >> 4) & 0x0F))
        };

        esp_err_t err = oled_send_cmds(page_cmds, sizeof(page_cmds));
        if (err != ESP_OK) {
            return err;
        }

        // Push 72 bytes of pixel data for the current page
        i2c_cmd_handle_t link = i2c_cmd_link_create();
        i2c_master_start(link);
        i2c_master_write_byte(link, (OLED_I2C_ADDR << 1) | I2C_MASTER_WRITE, true);
        i2c_master_write_byte(link, 0x40, true); // Control byte: 0x40 = Data
        i2c_master_write(link, &s_frame_buffer[page * OLED_WIDTH], OLED_WIDTH, true);
        i2c_master_stop(link);

        err = i2c_master_cmd_begin(OLED_I2C_PORT, link, pdMS_TO_TICKS(50));
        i2c_cmd_link_delete(link);

        if (err != ESP_OK) {
            return err;
        }
    }
    return ESP_OK;
}

esp_err_t oled_set_contrast(uint8_t contrast)
{
    uint8_t cmds[] = {0x81, contrast};
    return oled_send_cmds(cmds, sizeof(cmds));
}

esp_err_t oled_invert(bool invert)
{
    return oled_send_cmd(invert ? 0xA7 : 0xA6);
}

void oled_draw_pixel(int16_t x, int16_t y, uint8_t color)
{
    if (x < 0 || x >= OLED_WIDTH || y < 0 || y >= OLED_HEIGHT) {
        return;
    }

    uint16_t index = (y / 8) * OLED_WIDTH + x;
    uint8_t bit = y % 8;

    switch (color) {
        case OLED_COLOR_WHITE:
            s_frame_buffer[index] |= (1 << bit);
            break;
        case OLED_COLOR_BLACK:
            s_frame_buffer[index] &= ~(1 << bit);
            break;
        case OLED_COLOR_INVERT:
            s_frame_buffer[index] ^= (1 << bit);
            break;
        default:
            break;
    }
}

void oled_draw_line(int16_t x0, int16_t y0, int16_t x1, int16_t y1, uint8_t color)
{
    int16_t dx = abs(x1 - x0);
    int16_t dy = abs(y1 - y0);
    int16_t sx = (x0 < x1) ? 1 : -1;
    int16_t sy = (y0 < y1) ? 1 : -1;
    int16_t err = dx - dy;

    while (1) {
        oled_draw_pixel(x0, y0, color);
        if (x0 == x1 && y0 == y1) break;
        int16_t e2 = 2 * err;
        if (e2 > -dy) {
            err -= dy;
            x0 += sx;
        }
        if (e2 < dx) {
            err += dx;
            y0 += sy;
        }
    }
}

void oled_draw_rect(int16_t x, int16_t y, int16_t w, int16_t h, uint8_t color)
{
    if (w <= 0 || h <= 0) return;
    oled_draw_line(x, y, x + w - 1, y, color);
    oled_draw_line(x, y + h - 1, x + w - 1, y + h - 1, color);
    oled_draw_line(x, y, x, y + h - 1, color);
    oled_draw_line(x + w - 1, y, x + w - 1, y + h - 1, color);
}

void oled_fill_rect(int16_t x, int16_t y, int16_t w, int16_t h, uint8_t color)
{
    for (int16_t i = x; i < x + w; i++) {
        for (int16_t j = y; j < y + h; j++) {
            oled_draw_pixel(i, j, color);
        }
    }
}

void oled_draw_char(int16_t x, int16_t y, char c, uint8_t color)
{
    if (c < 32 || c > 127) {
        c = '?';
    }

    const uint8_t *char_data = font_5x7[c - 32];
    for (uint8_t col = 0; col < 5; col++) {
        uint8_t line = char_data[col];
        for (uint8_t row = 0; row < 7; row++) {
            if (line & (1 << row)) {
                oled_draw_pixel(x + col, y + row, color);
            } else if (color == OLED_COLOR_WHITE) {
                oled_draw_pixel(x + col, y + row, OLED_COLOR_BLACK);
            }
        }
    }
}

void oled_draw_string(int16_t x, int16_t y, const char *str, uint8_t color)
{
    if (!str) return;
    int16_t cur_x = x;
    int16_t cur_y = y;

    while (*str) {
        if (*str == '\n') {
            cur_x = x;
            cur_y += 8;
        } else {
            oled_draw_char(cur_x, cur_y, *str, color);
            cur_x += 6; // 5 pixels width + 1 pixel inter-character spacing
        }
        str++;
    }
}