#include <assert.h>
#include <string.h>

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

#include <time.h>

#include "epd.hpp"

internal const char *kTagEPD = "EPD";

internal spi_device_handle_t g_spi_handle;
internal bool g_is_asleep = true;
internal constexpr size_t kTotal_bytes = (EPD_WIDTH * EPD_HEIGHT) / 8;
internal uint8 g_scratch_buffer[4096];

internal struct epd_refresh_stats{
  time_t last_full_refresh_bw;
  int fast_count_bw;
  time_t last_full_refresh_4g;
  int fast_count_4g;
} g_epd_stats = {0, 0, 0, 0};

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);

  int chunk_count = 0;
  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;

    // Yield every 16 chunks (~64KB) to prevent task watchdog timeout
    if (++chunk_count % 16 == 0)
    {
      vTaskDelay(1);
    }
  }
}

internal void epd_fill_layer(uint8 cmd, uint8 color_byte)
{
  epd_writecommand(cmd);

  memset(g_scratch_buffer, color_byte, sizeof(g_scratch_buffer));

  size_t remaining = kTotal_bytes;
  while (remaining > 0)
  {
    size_t to_write =
        (remaining > sizeof(g_scratch_buffer)) ? sizeof(g_scratch_buffer) : remaining;
    epd_write_buffer(g_scratch_buffer, to_write);
    remaining -= to_write;
  }
}

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);
}

internal void epd_init_display_full(void)
{
  ESP_LOGI(kTagEPD, "Performing FULL initialization");
  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);
}

internal void epd_init_display_fast(void)
{
  ESP_LOGI(kTagEPD, "Performing FAST initialization");
  epd_writecommand(0X00); // PANEL SETTING
  epd_writedata(0x1F);

  epd_writecommand(0x50); // VCOM AND DATA INTERVAL SETTING
  epd_writedata(0x29);
  epd_writedata(0x07);

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

  // Enhanced display drive (Booster values for fast mode)
  epd_writecommand(0x06); // Booster Soft Start
  epd_writedata(0x27);
  epd_writedata(0x27);
  epd_writedata(0x18);
  epd_writedata(0x17);

  epd_writecommand(0xE0);
  epd_writedata(0x02);
  epd_writecommand(0xE5);
  epd_writedata(0x5A);
}

internal void epd_init_display_4g_full(void)
{
  ESP_LOGI(kTagEPD, "Performing FULL 4-GRAY initialization");
  epd_writecommand(0X00); // PANEL SETTING
  epd_writedata(0x1F);

  epd_writecommand(0X50); // VCOM AND DATA INTERVAL SETTING
  epd_writedata(0x29);
  epd_writedata(0x07);

  epd_writecommand(0x04); // POWER ON
  vTaskDelay(pdMS_TO_TICKS(100)); // Standard delay for full
  epd_wait_until_idle();

  // Standard display drive (Full Booster values)
  epd_writecommand(0x06); // Booster Soft Start
  epd_writedata(0x17);
  epd_writedata(0x17);
  epd_writedata(0x28);
  epd_writedata(0x17);

  epd_writecommand(0xE0);
  epd_writedata(0x02);
  epd_writecommand(0xE5);
  epd_writedata(0x5F); // 0x5F -- 4 Gray
}

internal void epd_init_display_4g_fast(void)
{
  ESP_LOGI(kTagEPD, "Performing FAST 4-GRAY initialization");
  epd_writecommand(0X00); // PANEL SETTING
  epd_writedata(0x1F);

  epd_writecommand(0X50); // VCOM AND DATA INTERVAL SETTING
  epd_writedata(0x29);
  epd_writedata(0x07);

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

  // Enhanced display drive (Fast Booster values)
  epd_writecommand(0x06); // Booster Soft Start
  epd_writedata(0x27);
  epd_writedata(0x27);
  epd_writedata(0x18);
  epd_writedata(0x17);

  epd_writecommand(0xE0);
  epd_writedata(0x02);
  epd_writecommand(0xE5);
  epd_writedata(0x5F); // 0x5F -- 4 Gray
}

void epd_init_display(bool is_4gray)
{
  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));

  time_t now = time(NULL);

  if (is_4gray)
  {
    double diff = difftime(now, g_epd_stats.last_full_refresh_4g);
    bool force_full = (g_epd_stats.last_full_refresh_4g == 0) || (diff > 86400.0) ||
                      (g_epd_stats.fast_count_4g >= 5);

    if (force_full)
    {
      epd_init_display_4g_full();
      g_epd_stats.last_full_refresh_4g = now;
      g_epd_stats.fast_count_4g = 0;
    }
    else
    {
      epd_init_display_4g_fast();
      g_epd_stats.fast_count_4g++;
    }
    ESP_LOGI(kTagEPD, "4G Refresh stats: (Fast count: %d/5, Age: %.1f hours)",
             g_epd_stats.fast_count_4g, diff / 3600.0);
  }
  else
  {
    double diff = difftime(now, g_epd_stats.last_full_refresh_bw);
    bool force_full = (g_epd_stats.last_full_refresh_bw == 0) || (diff > 86400.0) ||
                      (g_epd_stats.fast_count_bw >= 5);

    if (force_full)
    {
      epd_init_display_full();
      g_epd_stats.last_full_refresh_bw = now;
      g_epd_stats.fast_count_bw = 0;
    }
    else
    {
      epd_init_display_fast();
      g_epd_stats.fast_count_bw++;
    }
    ESP_LOGI(kTagEPD, "BW Refresh stats: (Fast count: %d/5, Age: %.1f hours)",
             g_epd_stats.fast_count_bw, diff / 3600.0);
  }
}

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 level)
{
  uint8 color_byte = static_cast<uint8>(level);
  ESP_LOGI(kTagEPD, "Clearing display (byte=0x%02X)", color_byte);

  epd_fill_layer(0x10,
                 0xFF); // Old data layer (0xFF is mapped to White)
  epd_fill_layer(0x13, color_byte); // New data layer

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

void epd_draw_bitmap(epd_color clearColor, const uint8 *bitmap)
{
  uint8 color_byte = static_cast<uint8>(clearColor);
  ESP_LOGI(kTagEPD, "Drawing bitmap (clearColor byte=0x%02X)", color_byte);

  epd_fill_layer(0x10, color_byte); // Send clear color to "Old" layer

  epd_writecommand(0x13); // "New" data layer
  epd_write_buffer(bitmap, kTotal_bytes);

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

void epd_draw_bitmap_grayscale(epd_color clearColor, const uint8 *bitmap)
{
  ESP_LOGI(kTagEPD, "Drawing 4-GRAY bitmap (with unpacking)");

  // THE LOGIC:
  // The input bitmap is 2-bits per pixel packed (4 pixels per byte).
  // Total pixels: 800 * 480 = 384,000.
  // Input size: 384,000 / 4 = 96,000 bytes.
  // Output: Two frames of 800 * 480 / 8 = 48,000 bytes each.

  auto process_layer = [&](uint8 cmd, bool is_new_layer)
  {
    epd_writecommand(cmd);
    size_t scratch_idx = 0;

    // Process 48,000 output bytes (each covers 8 pixels)
    for (size_t i = 0; i < 48000; i++)
    {
      uint8 output_byte = 0;
      // Each output byte comes from 2 input bytes (j=0, j=1)
      for (int j = 0; j < 2; j++)
      {
        uint8 input_byte = bitmap[i * 2 + j];
        // Extract 4 pixels from each input byte
        for (int k = 0; k < 4; k++)
        {
          uint8 pixel_bits = (input_byte >> (6 - k * 2)) & 0x03; // Correct bit order
          bool bit_val = false;

          if (is_new_layer)
          {
            // NEW Layer (0x13) Mapping
            if (pixel_bits == 0x03) bit_val = true;  // White (11)
            else if (pixel_bits == 0x00) bit_val = false; // Black (00)
            else if (pixel_bits == 0x02) bit_val = true;  // Gray1 (10)
            else if (pixel_bits == 0x01) bit_val = false; // Gray2 (01)
          }
          else
          {
            // OLD Layer (0x10) Mapping
            if (pixel_bits == 0x03) bit_val = true;  // White (11)
            else if (pixel_bits == 0x00) bit_val = false; // Black (00)
            else if (pixel_bits == 0x02) bit_val = false; // Gray1 (10)
            else if (pixel_bits == 0x01) bit_val = true;  // Gray2 (01)
          }

          if (bit_val) output_byte |= (1 << (7 - (j * 4 + k)));
        }
      }

      // The demo code applies a bitwise NOT to the result
      g_scratch_buffer[scratch_idx++] = ~output_byte;

      if (scratch_idx >= sizeof(g_scratch_buffer))
      {
        epd_write_buffer(g_scratch_buffer, scratch_idx);
        scratch_idx = 0;
      }
    }

    if (scratch_idx > 0)
    {
      epd_write_buffer(g_scratch_buffer, scratch_idx);
    }
  };

  process_layer(0x10, false); // Old data
  process_layer(0x13, true);  // New data

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