#include "display/epd_handler.h"
#include "esp_log.h"
#include "display/constants.h"
#include "common/constants.h"
#include "esp_lcd_touch_gt911.h"
#include <driver/i2c.h>
#define TAG "EPDHandler"

#define BUSY_ACTIVE_LEVEL   0  // BUSY pin is active low
#define BUSY_INACTIVE_LEVEL 1
#define DMA_TRANSFER_CHUNK_SIZE 4096  // 4KB chunk size for DMA transfers

EPDHandler::EPDHandler() {
  spi_mutex_ = xSemaphoreCreateMutex();
  if (spi_mutex_ == nullptr) {
    ESP_LOGE(TAG, "Failed to create SPI mutex");
  }
  spi_transaction_mutex_ = xSemaphoreCreateMutex();
  if (spi_transaction_mutex_ == nullptr) {
    ESP_LOGE(TAG, "Failed to create SPI transaction mutex");
  }
}

EPDHandler::~EPDHandler() {
  if (spi_mutex_ != nullptr) {
    vSemaphoreDelete(spi_mutex_);
  }
  if (spi_transaction_mutex_ != nullptr) {
    vSemaphoreDelete(spi_transaction_mutex_);
  }
  if (spi_ != nullptr) {
    spi_bus_remove_device(spi_);
  }
}

esp_err_t EPDHandler::init() {
  esp_err_t err;

  // Initialize SPI bus
  spi_bus_config_t buscfg = {};
  buscfg.mosi_io_num = 11;  // MOSI pin
  buscfg.miso_io_num = -1;  // No MISO for e-paper
  buscfg.sclk_io_num = 12;  // SCK pin
  buscfg.quadwp_io_num = -1;
  buscfg.quadhd_io_num = -1;
  buscfg.max_transfer_sz = DMA_TRANSFER_CHUNK_SIZE;

  err = spi_bus_initialize(SPI2_HOST, &buscfg, SPI_DMA_CH_AUTO);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "Failed to initialize SPI bus: %s", esp_err_to_name(err));
    return err;
  }

  // Add SPI device
  spi_device_interface_config_t devcfg = {};
  devcfg.clock_speed_hz = 10 * 1000 * 1000;  // 10 MHz
  devcfg.mode = 0;  // SPI mode 0
  devcfg.spics_io_num = PIN_CS;
  devcfg.queue_size = 7;  // Queue size for non-blocking transactions
  devcfg.pre_cb = nullptr;

  err = spi_bus_add_device(SPI2_HOST, &devcfg, &spi_);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "Failed to add SPI device: %s", esp_err_to_name(err));
    return err;
  }

  return ESP_OK;
}

// Check if display is busy (refreshing)
bool EPDHandler::is_busy(void) const {
  return gpio_get_level(PIN_BUSY) != BUSY_ACTIVE_LEVEL;  // BUSY is active LOW
}
void EPDHandler::wait_for_idle(void) const {
  ESP_LOGI(TAG, "Waiting for display ready (BUSY pin)...");
  int initial_level = gpio_get_level(PIN_BUSY);
  ESP_LOGI(TAG, "Initial BUSY pin level: %d (0=BUSY, 1=FREE)", initial_level);

  // If already free, no need to wait
  if (initial_level == BUSY_INACTIVE_LEVEL) {
    ESP_LOGI(TAG, "Display already ready (BUSY pin = 1)");
    return;
  }
  while (gpio_get_level(PIN_BUSY) != BUSY_INACTIVE_LEVEL) {
    vTaskDelay(pdMS_TO_TICKS(10));
  }
  ESP_LOGI(TAG, "Display is now ready (BUSY pin = 1)");
}

esp_err_t EPDHandler::epd_write_cmd(const uint8_t cmd, uint32_t transaction_id) {
  ESP_LOGI(TAG, "epd_write_cmd: waiting to send 0x%02X", cmd);

  SemaphoreGuard transaction_guard(spi_transaction_mutex_);
  esp_err_t err =
    wait_for_transaction_end_(pdMS_TO_TICKS(5000), transaction_id, transaction_guard);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "Failed to wait for previous transaction end before sending cmd 0x%02X: %s",
      cmd, esp_err_to_name(err));
    return err;
  }

  SemaphoreGuard guard(spi_mutex_);
  if (!guard.take(pdMS_TO_TICKS(5000))) {
    ESP_LOGE(TAG, "SPI mutex timeout for cmd 0x%02X", cmd);
    return ESP_ERR_TIMEOUT;
  }
  err = dangerous_epd_write_cmd_without_lock_(cmd);
  ESP_LOGI(TAG, "epd_write_cmd: 0x%02X done", cmd);
  return err;
}

esp_err_t EPDHandler::epd_write_data(const uint8_t data, uint32_t transaction_id) {
  ESP_LOGI(TAG, "epd_write_data: waiting to send 0x%02X", data);
  SemaphoreGuard transaction_guard(spi_transaction_mutex_);
  esp_err_t err =
    wait_for_transaction_end_(pdMS_TO_TICKS(5000), transaction_id, transaction_guard);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "Failed to wait for previous transaction end before sending data 0x%02X: %s",
      data, esp_err_to_name(err));
    return err;
  }
  SemaphoreGuard guard(spi_mutex_);
  if (!guard.take(pdMS_TO_TICKS(5000))) {
    ESP_LOGE(TAG, "SPI mutex timeout for data 0x%02X", data);
    return ESP_ERR_TIMEOUT;
  }
  err = dangerous_epd_write_data_without_lock_(data);
  ESP_LOGI(TAG, "epd_write_data: 0x%02X done", data);
  return err;
}

esp_err_t EPDHandler::epd_write_cmd_with_data(const uint8_t cmd, std::vector<uint8_t>& data, uint32_t transaction_id) {
  const size_t data_len = data.size();
  ESP_LOGI(TAG, "epd_write_cmd_with_data: waiting to send cmd 0x%02X with %u bytes of data", cmd, data_len);

  SemaphoreGuard transaction_guard(spi_transaction_mutex_);
  esp_err_t err =
    wait_for_transaction_end_(pdMS_TO_TICKS(5000), transaction_id, transaction_guard);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "Failed to wait for previous transaction end before sending cmd 0x%02X: %s, with data",
      cmd, esp_err_to_name(err));
    return err;
  }

  SemaphoreGuard guard(spi_mutex_);
  if (!guard.take(pdMS_TO_TICKS(5000))) {
    ESP_LOGE(TAG, "SPI mutex timeout for cmd with data 0x%02X", cmd);
    return ESP_ERR_TIMEOUT;
  }
  err = dangerous_epd_write_cmd_without_lock_(cmd);
  if (err != ESP_OK) {
    return err;
  };
  for (size_t i = 0; i < data_len; ++i) {
    err = dangerous_epd_write_data_without_lock_(data[i]);
    if (err != ESP_OK) {
      return err;
    }
  }
  ESP_LOGI(TAG, "epd_write_cmd_with_data: cmd 0x%02X with %u bytes of data done", cmd, data_len);
  return ESP_OK;
}


esp_err_t EPDHandler::dangerous_epd_write_cmd_without_lock_(const uint8_t cmd) {
  ESP_LOGI(TAG, "dangerous_epd_write_cmd_without_lock_: sending 0x%02X", cmd);
  gpio_set_level(PIN_DC, 0); // Command mode
  spi_transaction_t t {};
  t.length = 8;t.tx_buffer = &cmd;
  esp_err_t err = spi_device_polling_transmit(spi_, &t);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "Failed to send data 0x%02X", cmd);
  } else {
    ESP_LOGI(TAG, "dangerous_epd_write_cmd_without_lock_: 0x%02X sent", cmd);
  }
  return err;
}

esp_err_t EPDHandler::dangerous_epd_write_data_without_lock_(const uint8_t data) {
  ESP_LOGI(TAG, "dangerous_epd_write_data_without_lock_: sending 0x%02X", data);
  gpio_set_level(PIN_DC, 1); // Data mode
  spi_transaction_t t = { };
  t.length = 8; t.tx_buffer = &data;
  esp_err_t err = spi_device_polling_transmit(spi_, &t);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "Failed to send data 0x%02X", data);
  } else {
    ESP_LOGI(TAG, "dangerous_epd_write_data_without_lock_: 0x%02X sent", data);
  }
  return err;
}

esp_err_t EPDHandler::transfer_spi_data(const uint8_t* data, const size_t& length, uint32_t transaction_id) {
  ESP_LOGI(TAG, "transfer_spi_data: waiting to send %zu bytes of data", length);

  SemaphoreGuard transaction_guard(spi_transaction_mutex_);
  esp_err_t err =
    wait_for_transaction_end_(pdMS_TO_TICKS(5000), transaction_id, transaction_guard);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "Failed to wait for previous transaction end before sending data of %zu bytes: %s",
      length, esp_err_to_name(err));
    return err;
  }
  SemaphoreGuard guard(spi_mutex_);
  if (!guard.take(pdMS_TO_TICKS(5000))) {
    ESP_LOGE(TAG, "SPI mutex timeout for data transfer of %zu bytes", length);
    return ESP_ERR_TIMEOUT;
  }
  ESP_LOGI(TAG, "transfer_spi_data: starting to send %zu bytes of data", length);

  size_t offset = 0;
  size_t remaining = length;
  gpio_set_level(PIN_DC, 1); // Data mode
  while (remaining > 0) {
    size_t transfer_size = (remaining < DMA_TRANSFER_CHUNK_SIZE) ? remaining : DMA_TRANSFER_CHUNK_SIZE;

    spi_transaction_t t = {};
    t.length = transfer_size * 8;  // Length in bits
    t.tx_buffer = data + offset;

    esp_err_t ret = spi_device_polling_transmit(spi_, &t);
    if (ret != ESP_OK) {
      ESP_LOGE(TAG, "Failed to send SPI chunk at offset %zu: %s", offset, esp_err_to_name(ret));
      if (ret == ESP_ERR_NO_MEM) {
        ESP_LOGE(TAG, "Current free heap size: %u bytes", esp_get_free_heap_size());
        ESP_LOGE(TAG, "Current free DMA-capable memory size: %u bytes",
          heap_caps_get_free_size(MALLOC_CAP_DMA));
      }
      return ret;
    }

    remaining -= transfer_size;
    offset += transfer_size;

    // Yield every 16KB to prevent watchdog timeout
    if (offset % (16 * 1024) == 0) {
      ESP_LOGI(TAG, "New data progress: %zu/%zu bytes sent, yielding...", offset, length);
      vTaskDelay(pdMS_TO_TICKS(1));
    }
  }

  ESP_LOGI(TAG, "transfer_spi_data: completed sending %zu bytes of data", length);
  return ESP_OK;
}

esp_err_t EPDHandler::begin_transaction_(TickType_t timeout, uint32_t& out_id) {
  ESP_LOGI(TAG, "begin_transaction_: waiting to obtain transaction mutex");
  if (xSemaphoreTake(spi_transaction_mutex_, timeout) != pdTRUE) {
    ESP_LOGE(TAG, "begin_transaction_: transaction mutex timeout");
    return ESP_ERR_TIMEOUT;
  }

  out_id = ++spi_transaction_id;
  ESP_LOGI(TAG, "begin_transaction_: transaction mutex obtained");
  return ESP_OK;
}

esp_err_t EPDHandler::end_transaction_(void) {
  ESP_LOGI(TAG, "end_transaction_: releasing transaction mutex");
  if (xSemaphoreGive(spi_transaction_mutex_) != pdTRUE) {
    ESP_LOGE(TAG, "end_transaction_: failed to release transaction mutex");
    return ESP_FAIL;
  }

  ESP_LOGI(TAG, "end_transaction_: transaction mutex released");
  return ESP_OK;
}

esp_err_t EPDHandler::wait_for_transaction_end_(TickType_t timeout, uint32_t awaiting_transaction_id, SemaphoreGuard& out_transaction_guard) {
  // Validate transaction ID if provided
  if (awaiting_transaction_id != 0 && awaiting_transaction_id != spi_transaction_id) {
    // Invalid transaction ID
    ESP_LOGE(TAG, "Invalid transaction ID 0x%08X while waiting, current transaction ID: 0x%08X",
      awaiting_transaction_id, spi_transaction_id);
    return ESP_ERR_INVALID_ARG;
  }
  SemaphoreGuard transaction_guard(spi_transaction_mutex_);
  if (awaiting_transaction_id == 0) {
    // wait for current transaction to complete
    ESP_LOGV(TAG, "Waiting for current transaction 0x%08X to complete",
      spi_transaction_id);
    // take the mutex to ensure no transaction is active
    if (!transaction_guard.take(timeout)) {
      ESP_LOGE(TAG, "SPI transaction mutex timeout while waiting for transaction end");
      return ESP_ERR_TIMEOUT;
    }
  }
  // awaited_transaction_id is valid and matches current transaction ID or 0
  out_transaction_guard = std::move(transaction_guard);
  return ESP_OK;
}