#include <string.h>

#include "driver/gpio.h"
#include "driver/spi_master.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "esp_check.h"
#include "esp_log.h"

#include "epd7in3e.h"

/* Command bytes and register values below are a line-for-line transcription
 * of Waveshare's official EPD_7in3e.c reference driver (RaspberryPi_JetsonNano/c
 * variant) -- this panel's controller has no public datasheet, so the
 * upstream driver is the source of truth. Unlike that driver (which toggles
 * CS around every single byte), this port holds CS low for the duration of
 * each logical command or data phase and DMAs data in chunks, since 192,000
 * one-byte SPI transactions would make a full refresh impractically slow. */

#define EPD_SPI_HOST SPI2_HOST
#define EPD_SPI_CHUNK_SIZE 4096

static const char *TAG = "epd7in3e";

#define EPD_CHECK(expr) ESP_RETURN_ON_ERROR((expr), TAG, #expr)

static spi_device_handle_t s_spi;

static void epd_delay_ms(uint32_t ms)
{
    vTaskDelay(pdMS_TO_TICKS(ms));
}

/* BUSY: LOW = busy, HIGH = idle. */
static void epd_wait_busy(void)
{
    /* A full refresh on this panel can hold BUSY low for 15-30+ seconds.
     * A 1ms poll delay rounds down to 0 FreeRTOS ticks at the default
     * 100Hz tick rate, so vTaskDelay() never actually blocks -- confirmed
     * on hardware as a tight busy-spin that starves the idle task long
     * enough to trip the 5s task watchdog. 20ms is safely >= 1 tick
     * regardless of tick rate and still imperceptibly responsive here. */
    while (gpio_get_level((gpio_num_t)CONFIG_EPD_PIN_BUSY) == 0) {
        epd_delay_ms(20);
    }
}

static esp_err_t epd_spi_write(const uint8_t *data, size_t len)
{
    while (len > 0) {
        size_t n = len > EPD_SPI_CHUNK_SIZE ? EPD_SPI_CHUNK_SIZE : len;
        spi_transaction_t t = {
            .length = n * 8,
            .tx_buffer = data,
        };
        EPD_CHECK(spi_device_polling_transmit(s_spi, &t));
        data += n;
        len -= n;
    }
    return ESP_OK;
}

static esp_err_t epd_send_command(uint8_t cmd)
{
    gpio_set_level((gpio_num_t)CONFIG_EPD_PIN_DC, 0);
    gpio_set_level((gpio_num_t)CONFIG_EPD_PIN_CS, 0);
    esp_err_t err = epd_spi_write(&cmd, 1);
    gpio_set_level((gpio_num_t)CONFIG_EPD_PIN_CS, 1);
    return err;
}

static esp_err_t epd_send_data(const uint8_t *data, size_t len)
{
    gpio_set_level((gpio_num_t)CONFIG_EPD_PIN_DC, 1);
    gpio_set_level((gpio_num_t)CONFIG_EPD_PIN_CS, 0);
    esp_err_t err = epd_spi_write(data, len);
    gpio_set_level((gpio_num_t)CONFIG_EPD_PIN_CS, 1);
    return err;
}

static esp_err_t epd_send_data_byte(uint8_t data)
{
    return epd_send_data(&data, 1);
}

static void epd_reset(void)
{
    gpio_set_level((gpio_num_t)CONFIG_EPD_PIN_RST, 1);
    epd_delay_ms(20);
    gpio_set_level((gpio_num_t)CONFIG_EPD_PIN_RST, 0);
    epd_delay_ms(2);
    gpio_set_level((gpio_num_t)CONFIG_EPD_PIN_RST, 1);
    epd_delay_ms(20);
}

/* Power on, "second setting" registers, refresh, power off -- mirrors
 * EPD_7IN3E_TurnOnDisplay() in the reference driver. */
static esp_err_t epd_turn_on_display(void)
{
    EPD_CHECK(epd_send_command(0x04)); // POWER_ON
    epd_wait_busy();

    EPD_CHECK(epd_send_command(0x06));
    EPD_CHECK(epd_send_data((uint8_t[]){ 0x6F, 0x1F, 0x17, 0x49 }, 4));

    EPD_CHECK(epd_send_command(0x12)); // DISPLAY_REFRESH
    EPD_CHECK(epd_send_data_byte(0x00));
    epd_wait_busy();

    EPD_CHECK(epd_send_command(0x02)); // POWER_OFF
    EPD_CHECK(epd_send_data_byte(0x00));
    epd_wait_busy();

    return ESP_OK;
}

esp_err_t epd_init(void)
{
    gpio_config_t out_cfg = {
        .pin_bit_mask = (1ULL << CONFIG_EPD_PIN_DC) | (1ULL << CONFIG_EPD_PIN_RST) | (1ULL << CONFIG_EPD_PIN_CS),
        .mode = GPIO_MODE_OUTPUT,
    };
    EPD_CHECK(gpio_config(&out_cfg));

    gpio_config_t busy_cfg = {
        .pin_bit_mask = (1ULL << CONFIG_EPD_PIN_BUSY),
        .mode = GPIO_MODE_INPUT,
    };
    EPD_CHECK(gpio_config(&busy_cfg));

    gpio_set_level((gpio_num_t)CONFIG_EPD_PIN_CS, 1);

    spi_bus_config_t bus_cfg = {
        .mosi_io_num = CONFIG_EPD_PIN_MOSI,
        .miso_io_num = -1,
        .sclk_io_num = CONFIG_EPD_PIN_CLK,
        .quadwp_io_num = -1,
        .quadhd_io_num = -1,
        .max_transfer_sz = EPD_SPI_CHUNK_SIZE,
    };
    EPD_CHECK(spi_bus_initialize(EPD_SPI_HOST, &bus_cfg, SPI_DMA_CH_AUTO));

    spi_device_interface_config_t dev_cfg = {
        .clock_speed_hz = CONFIG_EPD_SPI_CLOCK_HZ,
        .mode = 0,
        /* CS is bit-banged around each command/data phase above rather than
         * hardware-driven, since a phase can span many chunked SPI
         * transactions and must stay asserted across all of them. */
        .spics_io_num = -1,
        .queue_size = 1,
    };
    EPD_CHECK(spi_bus_add_device(EPD_SPI_HOST, &dev_cfg, &s_spi));

    epd_reset();
    epd_wait_busy();
    epd_delay_ms(30);

    EPD_CHECK(epd_send_command(0xAA)); // CMDH
    EPD_CHECK(epd_send_data((uint8_t[]){ 0x49, 0x55, 0x20, 0x08, 0x09, 0x18 }, 6));

    EPD_CHECK(epd_send_command(0x01));
    EPD_CHECK(epd_send_data_byte(0x3F));

    EPD_CHECK(epd_send_command(0x00));
    EPD_CHECK(epd_send_data((uint8_t[]){ 0x5F, 0x69 }, 2));

    EPD_CHECK(epd_send_command(0x03));
    EPD_CHECK(epd_send_data((uint8_t[]){ 0x00, 0x54, 0x00, 0x44 }, 4));

    EPD_CHECK(epd_send_command(0x05));
    EPD_CHECK(epd_send_data((uint8_t[]){ 0x40, 0x1F, 0x1F, 0x2C }, 4));

    EPD_CHECK(epd_send_command(0x06));
    EPD_CHECK(epd_send_data((uint8_t[]){ 0x6F, 0x1F, 0x17, 0x49 }, 4));

    EPD_CHECK(epd_send_command(0x08));
    EPD_CHECK(epd_send_data((uint8_t[]){ 0x6F, 0x1F, 0x1F, 0x22 }, 4));

    EPD_CHECK(epd_send_command(0x30));
    EPD_CHECK(epd_send_data_byte(0x03));

    EPD_CHECK(epd_send_command(0x50));
    EPD_CHECK(epd_send_data_byte(0x3F));

    EPD_CHECK(epd_send_command(0x60));
    EPD_CHECK(epd_send_data((uint8_t[]){ 0x02, 0x00 }, 2));

    EPD_CHECK(epd_send_command(0x61));
    EPD_CHECK(epd_send_data((uint8_t[]){ 0x03, 0x20, 0x01, 0xE0 }, 4));

    EPD_CHECK(epd_send_command(0x84));
    EPD_CHECK(epd_send_data_byte(0x01));

    EPD_CHECK(epd_send_command(0xE3));
    EPD_CHECK(epd_send_data_byte(0x2F));

    EPD_CHECK(epd_send_command(0x04)); // POWER_ON; waits for the panel to release the idle signal
    epd_wait_busy();

    ESP_LOGI(TAG, "EPD initialized (CLK=%d MOSI=%d CS=%d DC=%d RST=%d BUSY=%d)",
             CONFIG_EPD_PIN_CLK, CONFIG_EPD_PIN_MOSI, CONFIG_EPD_PIN_CS,
             CONFIG_EPD_PIN_DC, CONFIG_EPD_PIN_RST, CONFIG_EPD_PIN_BUSY);

    return ESP_OK;
}

esp_err_t epd_display_stream(epd_read_fn_t read_fn, void *ctx)
{
    ESP_RETURN_ON_FALSE(read_fn != NULL, ESP_ERR_INVALID_ARG, TAG, "read_fn required");

    EPD_CHECK(epd_send_command(0x10));

    gpio_set_level((gpio_num_t)CONFIG_EPD_PIN_DC, 1);
    gpio_set_level((gpio_num_t)CONFIG_EPD_PIN_CS, 0);

    /* Static rather than a stack local: the default main task stack
     * (CONFIG_ESP_MAIN_TASK_STACK_SIZE, 3584 bytes) is smaller than this
     * 4KB chunk buffer alone, let alone with the rest of the call chain --
     * confirmed on hardware as a stack protection fault/crash. Not
     * reentrant, but this driver only ever runs from one task at a time. */
    static uint8_t chunk[EPD_SPI_CHUNK_SIZE];
    size_t total = 0;
    size_t n;
    esp_err_t err = ESP_OK;
    while ((n = read_fn(chunk, sizeof(chunk), ctx)) > 0) {
        err = epd_spi_write(chunk, n);
        if (err != ESP_OK) {
            break;
        }
        total += n;
    }

    gpio_set_level((gpio_num_t)CONFIG_EPD_PIN_CS, 1);
    EPD_CHECK(err);

    if (total != EPD_FRAME_BYTES) {
        /* Whatever was received has already been clocked into the panel's
         * internal RAM over SPI, but epd_turn_on_display() (the actual
         * physical refresh trigger) hasn't been called yet -- returning
         * here instead leaves the visible screen exactly as it was, rather
         * than refreshing onto a mostly-garbage buffer. Confirmed on
         * hardware: a misdirected fetch that returned a ~10KB error page
         * instead of a 192,000-byte frame still triggered a refresh before
         * this check existed, painting garbage over a previously-good image. */
        ESP_LOGE(TAG, "Stream supplied %u bytes, expected %u -- aborting refresh",
                 (unsigned)total, (unsigned)EPD_FRAME_BYTES);
        return ESP_ERR_INVALID_SIZE;
    }

    return epd_turn_on_display();
}

typedef struct {
    const uint8_t *data;
    size_t len;
    size_t pos;
} epd_buf_ctx_t;

static size_t epd_buf_read(uint8_t *chunk, size_t chunk_size, void *ctx_)
{
    epd_buf_ctx_t *c = (epd_buf_ctx_t *)ctx_;
    size_t remaining = c->len - c->pos;
    size_t n = remaining < chunk_size ? remaining : chunk_size;
    if (n == 0) {
        return 0;
    }
    memcpy(chunk, c->data + c->pos, n);
    c->pos += n;
    return n;
}

esp_err_t epd_display_buffer(const uint8_t *frame, size_t len)
{
    epd_buf_ctx_t buf_ctx = { .data = frame, .len = len, .pos = 0 };
    return epd_display_stream(epd_buf_read, &buf_ctx);
}

typedef struct {
    uint8_t fill_byte;
    size_t remaining;
} epd_fill_ctx_t;

static size_t epd_fill_read(uint8_t *chunk, size_t chunk_size, void *ctx_)
{
    epd_fill_ctx_t *c = (epd_fill_ctx_t *)ctx_;
    size_t n = c->remaining < chunk_size ? c->remaining : chunk_size;
    if (n == 0) {
        return 0;
    }
    memset(chunk, c->fill_byte, n);
    c->remaining -= n;
    return n;
}

esp_err_t epd_clear(epd_color_t color)
{
    epd_fill_ctx_t fill_ctx = {
        .fill_byte = (uint8_t)((color << 4) | color),
        .remaining = EPD_FRAME_BYTES,
    };
    return epd_display_stream(epd_fill_read, &fill_ctx);
}

esp_err_t epd_sleep(void)
{
    EPD_CHECK(epd_send_command(0x02)); // POWER_OFF
    EPD_CHECK(epd_send_data_byte(0x00));
    epd_wait_busy();

    EPD_CHECK(epd_send_command(0x07)); // DEEP_SLEEP
    EPD_CHECK(epd_send_data_byte(0xA5));

    return ESP_OK;
}