#include "epd.hpp"

#include "driver/gpio.h"
#include "driver/spi_master.h"
#include "esp_log.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include <assert.h>
#include <string.h>

internal const char *kTagEPD = "EPD";

internal spi_device_handle_t g_spi_handle;
internal bool g_is_asleep = true;

internal void epd_spi_init(void)
{
  // 1. Initialize the SPI Bus
  spi_bus_config_t buscfg = {};
  buscfg.miso_io_num =
      TFT_MISO; // Initialize MISO as input, matching `spi.begin(TFT_SCLK,
                // TFT_MISO, TFT_MOSI, -1)`
  buscfg.mosi_io_num = TFT_MOSI;
  buscfg.sclk_io_num = TFT_SCLK;
  buscfg.quadwp_io_num = -1;
  buscfg.quadhd_io_num = -1;
  buscfg.max_transfer_sz = 4096;

  // SPI2_HOST is the general-purpose SPI host on ESP32-S3
  ESP_ERROR_CHECK(spi_bus_initialize(SPI2_HOST, &buscfg, SPI_DMA_CH_AUTO));

  // 2. Add the EPD device to the bus
  spi_device_interface_config_t devcfg = {};
  devcfg.clock_speed_hz = SPI_FREQUENCY;
  devcfg.mode = 0;              // Standard EPD SPI mode
  devcfg.spics_io_num = TFT_CS; // The driver handles CS automatically
  devcfg.queue_size = 7;
  devcfg.flags = 0;

  ESP_ERROR_CHECK(spi_bus_add_device(SPI2_HOST, &devcfg, &g_spi_handle));
}

internal void epd_gpio_init()
{
  gpio_reset_pin((gpio_num_t)TFT_BUSY);
  gpio_set_direction((gpio_num_t)TFT_BUSY, GPIO_MODE_INPUT);
  gpio_reset_pin((gpio_num_t)TFT_RST);
  gpio_set_direction((gpio_num_t)TFT_RST, GPIO_MODE_OUTPUT);
  gpio_set_level((gpio_num_t)TFT_RST,
                 1); // Set high, do not share pin with another SPI device
  gpio_reset_pin((gpio_num_t)TFT_DC);
  gpio_set_direction((gpio_num_t)TFT_DC, GPIO_MODE_OUTPUT);
  gpio_set_level((gpio_num_t)TFT_DC, 1); // Data/Command high = data mode
  gpio_reset_pin((gpio_num_t)TFT_CS);
  gpio_set_direction((gpio_num_t)TFT_CS, GPIO_MODE_OUTPUT);
  gpio_set_level((gpio_num_t)TFT_CS, 1); // Chip select high (inactive)
}

internal void epd_writecommand(unsigned char command)
{
  assert(!g_is_asleep);
  // Pull DC Low for Command
  gpio_set_level((gpio_num_t)TFT_DC, GPIO_LOW);

  spi_transaction_t t = {};
  t.length = 8; // length in bits
  t.tx_buffer = &command;

  spi_device_transmit(g_spi_handle, &t);
}

internal void epd_write_buffer(const uint8 *data, size_t len)
{
  assert(!g_is_asleep);
  if (len == 0)
    return;

  // Pull DC High for Data
  gpio_set_level((gpio_num_t)TFT_DC, GPIO_HIGH);

  while (len > 0)
  {
    // max_transfer_sz is typically 4096 which is the limit for a single SPI
    // transaction
    size_t chunk = (len > 4096) ? 4096 : len;
    spi_transaction_t t = {};
    t.length = chunk * 8; // length in bits
    t.tx_buffer = data;

    spi_device_transmit(g_spi_handle, &t);
    data += chunk;
    len -= chunk;
  }
}

void epd_writedata(unsigned char data) { epd_write_buffer(&data, 1); }

internal void epd_wait_until_idle()
{
  // BUSY pin on this board: LOW = busy, HIGH = idle
  // (opposite of GDEY075T7 reference driver)
  while (gpio_get_level((gpio_num_t)TFT_BUSY) == 0)
  {
    vTaskDelay(pdMS_TO_TICKS(5));
  }
}

void epd_init(void)
{
  ESP_LOGI(kTagEPD, "Initializing EPaper Driver");

  epd_gpio_init();
  epd_spi_init();

  ESP_LOGI(kTagEPD, "EPaper Driver initialized successfully");
}

void epd_shutdown(void)
{
  spi_bus_remove_device(g_spi_handle);
  spi_bus_free(SPI2_HOST);
}

void epd_init_display()
{
  g_is_asleep = false;

  // Module reset
  gpio_set_level((gpio_num_t)TFT_RST, 0);
  vTaskDelay(pdMS_TO_TICKS(10));
  gpio_set_level((gpio_num_t)TFT_RST, 1);
  vTaskDelay(pdMS_TO_TICKS(10));

  epd_writecommand(0x01); // POWER SETTING
  epd_writedata(0x07);
  epd_writedata(0x07);
  epd_writedata(0x3f);
  epd_writedata(0x3f);

  epd_writecommand(0x06); // Booster Soft Start
  epd_writedata(0x17);
  epd_writedata(0x17);
  epd_writedata(0x28);
  epd_writedata(0x17);

  epd_writecommand(0x04); // POWER ON
  vTaskDelay(pdMS_TO_TICKS(100));
  epd_wait_until_idle();

  epd_writecommand(0X00); // PANEL SETTING
  epd_writedata(0x1F);

  epd_writecommand(0x61); // TRES
  epd_writedata(EPD_WIDTH >> 8);
  epd_writedata(EPD_WIDTH & 0xFF);
  epd_writedata(EPD_HEIGHT >> 8);
  epd_writedata(EPD_HEIGHT & 0xFF);

  epd_writecommand(0x15);
  epd_writedata(0x00);

  epd_writecommand(0x50); // VCOM AND DATA INTERVAL SETTING
  epd_writedata(0x29); // BDV=10 (White Border), N2OCP=1 (Auto-copy NEW to OLD), DDX=01 (0=Black, 1=White)
  epd_writedata(0x07);

  epd_writecommand(0x60); // TCON SETTING
  epd_writedata(0x22);
}

void epd_shutdown_display(void)
{
  assert(!g_is_asleep);
  ESP_LOGI(kTagEPD, "Shutting down display (Power Off)");

  epd_writecommand(0x50); // VCOM AND DATA INTERVAL SETTING (pre-sleep)
  epd_writedata(0xF7);

  epd_writecommand(0x02); // POWER OFF
  epd_wait_until_idle();

  epd_writecommand(0x07); // DEEP SLEEP
  epd_writedata(0xA5);    // Deep sleep check code

  vTaskDelay(pdMS_TO_TICKS(10));
  g_is_asleep = true;
}

void epd_refresh(void)
{
  assert(!g_is_asleep);
  ESP_LOGI(kTagEPD, "Refreshing display...");
  epd_writecommand(0x12); // REFRESH
  vTaskDelay(pdMS_TO_TICKS(1));
  epd_wait_until_idle();
}

bool epd_is_asleep(void) { return g_is_asleep; }

void epd_clear(epd_color_t level)
{
  assert(!g_is_asleep);

  // Directly pass the color preference to hardware (hardware polarity is
  // configured in DDX)
  uint8 color_byte = static_cast<uint8>(level);
  ESP_LOGI(kTagEPD, "Clearing display (byte=0x%02X)", color_byte);

  constexpr size_t total_bytes = (EPD_WIDTH * EPD_HEIGHT) / 8;

  auto write_layer = [&](uint8 cmd, uint8 fill_byte)
  {
    uint8 chunk[256];
    memset(chunk, fill_byte, sizeof(chunk));

    epd_writecommand(cmd);
    size_t remaining = total_bytes;
    int chunk_count = 0;
    while (remaining > 0)
    {
      size_t to_write = (remaining > sizeof(chunk)) ? sizeof(chunk) : remaining;
      epd_write_buffer(chunk, to_write);
      remaining -= to_write;
      // Yield every 16 chunks (~4KB) to prevent task watchdog timeout
      if (++chunk_count % 16 == 0)
        vTaskDelay(1);
    }
  };

  write_layer(
      0x10,
      0xFF); // Old data layer (0xFF is mapped to White under new polarity)
  write_layer(0x13, color_byte); // New data layer

  ESP_LOGI(kTagEPD, "Data transmission complete (Refresh required)");
}