#include "os_oled.h"
#include <string.h>
#include "esp_log.h"
#include "driver/i2c_master.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"

static const char *TAG = "OLED_DRV";

static uint8_t s_frame_buffer[OLED_BUFFER_SIZE];
static uint8_t s_column_offset = OLED_COLUMN_OFFSET;

static i2c_master_bus_handle_t s_i2c_bus = NULL;
static i2c_master_dev_handle_t s_i2c_dev = NULL;

/* Send a single command byte to the OLED */
static esp_err_t oled_send_cmd(uint8_t cmd)
{
    uint8_t buf[2] = {0x00, cmd};   /* 0x00 = Co=0, D/C#=0 (command) */
    return i2c_master_transmit(s_i2c_dev, buf, sizeof(buf), 50);
}

/* Send multiple command bytes in one command transaction */
static esp_err_t oled_send_cmds(const uint8_t *cmds, size_t len)
{
    uint8_t buf[1 + 64];
    if (len > sizeof(buf) - 1) return ESP_ERR_INVALID_SIZE;
    buf[0] = 0x00;   /* Co=0, D/C#=0: remaining bytes are commands */
    memcpy(&buf[1], cmds, len);
    return i2c_master_transmit(s_i2c_dev, buf, 1 + len, 100);
}

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

    /* --- I2C bus --- */
    i2c_master_bus_config_t bus_cfg = {
        .i2c_port        = OLED_I2C_PORT,
        .sda_io_num      = sda,
        .scl_io_num      = scl,
        .clk_source      = I2C_CLK_SRC_DEFAULT,
        .glitch_ignore_cnt = 7,
        .flags.enable_internal_pullup = true,
    };
    esp_err_t err = i2c_new_master_bus(&bus_cfg, &s_i2c_bus);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "i2c_new_master_bus failed: %s", esp_err_to_name(err));
        return err;
    }

    /* --- OLED device on the bus --- */
    i2c_device_config_t dev_cfg = {
        .dev_addr_length = I2C_ADDR_BIT_LEN_7,
        .device_address  = OLED_I2C_ADDR,
        .scl_speed_hz    = OLED_I2C_FREQ_HZ,
    };
    err = i2c_master_bus_add_device(s_i2c_bus, &dev_cfg, &s_i2c_dev);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "i2c_master_bus_add_device failed: %s", esp_err_to_name(err));
        return err;
    }

    vTaskDelay(pdMS_TO_TICKS(50));

    /* SH1106 / SSD1306 init sequence for 72x40 */
    const uint8_t init_cmds[] = {
        0xAE,             /* Display OFF */
        0xD5, 0x80,       /* Clock divide / Osc freq */
        0xA8, 0x27,       /* Multiplex Ratio (40 - 1 = 39 = 0x27) */
        0xD3, 0x00,       /* Display Offset = 0 */
        0x40,             /* Start Line = 0 */
        0x8D, 0x14,       /* Charge Pump enable */
        0x20, 0x02,       /* Memory Addressing Mode: Page mode */
        0xA1,             /* Segment Re-map */
        0xC8,             /* COM Output Scan Direction */
        0xDA, 0x12,       /* COM Pins hardware config */
        0x81, 0xAF,       /* Contrast */
        0xD9, 0xF1,       /* Pre-charge Period */
        0xDB, 0x20,       /* VCOMH Deselect Level */
        0xA4,             /* Resume to RAM */
        0xA6,             /* Normal display */
        0xAF              /* Display ON */
    };

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

    oled_driver_clear();
    oled_driver_flush();
    ESP_LOGI(TAG, "OLED driver initialized (72x40, col_offset=%d)", s_column_offset);
    return ESP_OK;
}

uint8_t *oled_driver_get_buffer(void)
{
    return s_frame_buffer;
}

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

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

esp_err_t oled_driver_flush(void)
{
    for (uint8_t page = 0; page < OLED_PAGES; page++) {
        /* Set page address + column start */
        const 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;

        /* Send pixel data: [0x40, pixel0, pixel1, ...] */
        uint8_t data_buf[1 + OLED_WIDTH];
        data_buf[0] = 0x40;   /* Co=0, D/C#=1 (data) */
        memcpy(&data_buf[1], &s_frame_buffer[page * OLED_WIDTH], OLED_WIDTH);

        err = i2c_master_transmit(s_i2c_dev, data_buf, sizeof(data_buf), 50);
        if (err != ESP_OK) return err;
    }
    return ESP_OK;
}

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

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

void oled_driver_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 COLOR_WHITE:  s_frame_buffer[index] |=  (1 << bit); break;
        case COLOR_BLACK:  s_frame_buffer[index] &= ~(1 << bit); break;
        case COLOR_INVERT: s_frame_buffer[index] ^=  (1 << bit); break;
        default: break;
    }
}