#include "display/eink_display_handler.h"
#include "display/constants.h"
#include "common/constants.h"
#include "esp_log.h"
#include "esp_heap_caps.h"
#include "esp_task_wdt.h"
#include <cstring>

#define TAG "EInkDisplayHandler"
#define BUSY_ACTIVE_LEVEL   0  // BUSY pin is active low
#define BUSY_INACTIVE_LEVEL 1

EInkDisplayHandler::EInkDisplayHandler(EventGroupHandle_t system_event_group)
  : DisplayHandler(system_event_group) {
  _refresh_mutex = xSemaphoreCreateMutex();
  if (_refresh_mutex == nullptr) {
    ESP_LOGE(TAG, "Failed to create refresh mutex");
  }
}

EInkDisplayHandler::~EInkDisplayHandler() {
  if (_touch_task_handle != nullptr) {
    vTaskDelete(_touch_task_handle);
  }
  if (_lvgl_display != nullptr) {
    lv_display_delete(_lvgl_display);
    _lvgl_display = nullptr;
    if (_lvgl_draw_buf != nullptr) {
      lv_draw_buf_destroy(_lvgl_draw_buf);
      _lvgl_draw_buf = nullptr;
    }
  }
  if (_lvgl_touch_indev != nullptr) {
    lvgl_port_remove_touch(_lvgl_touch_indev);
  }
  if (_framebuffer != nullptr) {
    heap_caps_free(_framebuffer);
  }
  if (_refresh_mutex != nullptr) {
    vSemaphoreDelete(_refresh_mutex);
  }
}

void EInkDisplayHandler::init() {
  ESP_LOGI(TAG, "Initializing E-Ink display handler...");

  // Initialize GPIO pins
  gpio_config_t io_conf = {};
  io_conf.pin_bit_mask = (1ULL << PIN_DC) | (1ULL << PIN_RST);
  io_conf.mode = GPIO_MODE_OUTPUT;
  io_conf.pull_up_en = GPIO_PULLUP_DISABLE;
  io_conf.pull_down_en = GPIO_PULLDOWN_DISABLE;
  io_conf.intr_type = GPIO_INTR_DISABLE;
  gpio_config(&io_conf);

  // Configure BUSY pin as input (no pull-up like sample code)
  io_conf.pin_bit_mask = (1ULL << PIN_BUSY);
  io_conf.mode = GPIO_MODE_INPUT;
  io_conf.pull_up_en = GPIO_PULLUP_DISABLE;
  gpio_config(&io_conf);

  // 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 = DISPLAY_BUFFER_SIZE;

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

  // Add SPI device
  spi_device_interface_config_t devcfg = {};
  devcfg.clock_speed_hz = 6 * 1000 * 1000;  // 6 MHz (reduced for reliability)
  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;

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

  // Initialize base display and touch devices
  init_devices(false);  // Don't set ready bit yet

  // Allocate framebuffer - try PSRAM first, fallback to internal RAM
  _framebuffer = (uint8_t*)heap_caps_malloc(DISPLAY_BUFFER_SIZE, MALLOC_CAP_SPIRAM);
  if (_framebuffer != nullptr) {
    _framebuffer_in_psram = true;
    ESP_LOGI(TAG, "Framebuffer allocated in PSRAM (%d bytes)", DISPLAY_BUFFER_SIZE);
  } else {
    ESP_LOGW(TAG, "PSRAM not available, allocating framebuffer in internal RAM");
    _framebuffer = (uint8_t*)heap_caps_malloc(DISPLAY_BUFFER_SIZE, MALLOC_CAP_INTERNAL);
    _framebuffer_in_psram = false;
    if (_framebuffer == nullptr) {
      ESP_LOGE(TAG, "Failed to allocate framebuffer");
      return;
    }
    ESP_LOGI(TAG, "Framebuffer allocated in internal RAM (%d bytes)", DISPLAY_BUFFER_SIZE);
  }
  memset(_framebuffer, 0xFF, DISPLAY_BUFFER_SIZE);  // Initialize to white

  // Perform initial full refresh to clear display BEFORE creating LVGL display
  // This prevents LVGL from trying to render during the initial clear
  ESP_LOGI(TAG, "Performing initial display clear...");
  _perform_full_refresh(_framebuffer);
  ESP_LOGI(TAG, "Initial display clear complete");

  // Create LVGL display manually (no esp_lcd panel for e-paper)
  lv_display_t* disp = lv_display_create(DISPLAY_WIDTH, DISPLAY_HEIGHT);
  if (disp == nullptr) {
    ESP_LOGE(TAG, "Failed to create LVGL display");
    return;
  }

  /* 1-bit e-paper display */
  lv_display_set_color_format(disp, LV_COLOR_FORMAT_I1);

  /* Create a draw buffer covering ~40 lines */
  _lvgl_draw_buf = lv_draw_buf_create(DISPLAY_WIDTH, DISPLAY_HEIGHT, LV_COLOR_FORMAT_I1, LV_STRIDE_AUTO);
  if (_lvgl_draw_buf == nullptr) {
    ESP_LOGE(TAG, "Failed to create LVGL draw buffer");
    lv_display_delete(disp);
    return;
  }

  lv_display_set_draw_buffers(disp, _lvgl_draw_buf, NULL);
  lv_display_set_render_mode(disp, LV_DISPLAY_RENDER_MODE_DIRECT);

  // Set custom flush callback and user data
  lv_display_set_flush_cb(disp, _lvgl_flush_cb);
  lv_display_set_user_data(disp, this);

  _lvgl_display = disp;

  ESP_LOGI(TAG, "LVGL display registered");

  // Register GT911 touch input with LVGL, only if touch handle is valid
  esp_lcd_touch_handle_t tp_handle = get_touch_handle();
  if (tp_handle == nullptr) {
    ESP_LOGE(TAG, "Touch handle is NULL — touch initialization failed; skipping LVGL touch registration");
  } else {
    const lvgl_port_touch_cfg_t touch_cfg = {
      .disp = _lvgl_display,
      .handle = tp_handle,
      .scale = {}, // Default scaling
    };

    _lvgl_touch_indev = lvgl_port_add_touch(&touch_cfg);
    if (_lvgl_touch_indev == nullptr) {
      ESP_LOGE(TAG, "Failed to register LVGL touch input");
      return;
    }

    // Override touch read callback to check BUSY pin
    lv_indev_set_read_cb(_lvgl_touch_indev, _lvgl_touch_read_cb);
    lv_indev_set_user_data(_lvgl_touch_indev, this);

    ESP_LOGI(TAG, "LVGL touch input registered");
  }

  // Set display ready bits
  xEventGroupSetBits(_system_event_group, DISPLAY_READY_BIT | TOUCH_CALIBRATED_BIT);
  ESP_LOGI(TAG, "E-Ink display handler initialized successfully");
}

void EInkDisplayHandler::start_touch_task() {
  // Note: With lvgl_port_add_touch, the ESP-IDF LVGL port handles touch reading internally
  // We don't need a separate touch task unless we want custom processing
  ESP_LOGI(TAG, "Touch input handled by LVGL port");
}

void EInkDisplayHandler::request_full_refresh() {
  if (xSemaphoreTake(_refresh_mutex, pdMS_TO_TICKS(100)) == pdTRUE) {
    _force_full_refresh = true;
    _partial_refresh_count = 0;
    xSemaphoreGive(_refresh_mutex);
    ESP_LOGI(TAG, "Full refresh requested");
  }
}

bool EInkDisplayHandler::is_busy() const {
  return gpio_get_level(PIN_BUSY) == BUSY_ACTIVE_LEVEL;  // BUSY is active LOW
}

void EInkDisplayHandler::_lvgl_flush_cb(lv_display_t* disp, const lv_area_t* area, uint8_t* px_map) {
  EInkDisplayHandler* handler = static_cast<EInkDisplayHandler*>(lv_display_get_user_data(disp));
  if (handler == nullptr) {
    ESP_LOGE(TAG, "Invalid handler in flush callback");
    lv_display_flush_ready(disp);
    return;
  }

  // Check if display is busy with detailed logging
  int busy_level = gpio_get_level(PIN_BUSY);
  ESP_LOGI(TAG, "Flush callback: BUSY pin = %d, is_busy() = %d", busy_level, handler->is_busy());

  if (handler->is_busy()) {
    ESP_LOGW(TAG, "Display busy (BUSY pin = 0), skipping flush");
    lv_display_flush_ready(disp);
    return;
  }

  // Wait for any ongoing refresh to complete
  handler->_wait_for_busy();

  bool perform_full_refresh = false;

  if (xSemaphoreTake(handler->_refresh_mutex, pdMS_TO_TICKS(100)) == pdTRUE) {
    // Check if full refresh is needed
    if (handler->_force_full_refresh) {
      perform_full_refresh = true;
      handler->_force_full_refresh = false;
      handler->_partial_refresh_count = 0;
    } else {
      handler->_partial_refresh_count++;
      if (handler->_partial_refresh_count >= PARTIAL_REFRESH_THRESHOLD) {
        perform_full_refresh = true;
        handler->_partial_refresh_count = 0;
      }
    }
    xSemaphoreGive(handler->_refresh_mutex);
  }

  // Copy LVGL buffer to framebuffer
  // For 1-bit mode, LVGL provides data in packed format (8 pixels per byte)
  int32_t w = lv_area_get_width(area);
  int32_t h = lv_area_get_height(area);

  ESP_LOGI(TAG, "Flushing area: x=%d, y=%d, w=%d, h=%d, full_refresh=%d",
    area->x1, area->y1, w, h, perform_full_refresh);

  // For simplicity with e-paper, we'll do full frame updates
  // Copy the entire buffer
  for (int32_t y = 0; y < h; y++) {
    int32_t fb_y = area->y1 + y;
    if (fb_y >= DISPLAY_HEIGHT) break;

    for (int32_t x = 0; x < w; x += 8) {
      int32_t fb_x = area->x1 + x;
      if (fb_x >= DISPLAY_WIDTH) break;

      // Calculate byte position in framebuffer (row-major, 1-bit packed)
      size_t fb_byte_idx = (fb_y * DISPLAY_WIDTH + fb_x) / 8;
      size_t px_byte_idx = (y * w + x) / 8;

      if (fb_byte_idx < DISPLAY_BUFFER_SIZE && px_byte_idx < (w * h / 8)) {
        handler->_framebuffer[fb_byte_idx] = px_map[px_byte_idx];
      }
    }
  }

  // Perform refresh
  if (perform_full_refresh) {
    ESP_LOGI(TAG, "Performing full refresh...");
    handler->_perform_full_refresh(handler->_framebuffer);
  } else {
    ESP_LOGI(TAG, "Performing partial refresh...");
    handler->_perform_partial_refresh(handler->_framebuffer);
  }

  lv_display_flush_ready(disp);
}

void EInkDisplayHandler::_lvgl_touch_read_cb(lv_indev_t* indev, lv_indev_data_t* data) {
  EInkDisplayHandler* handler = static_cast<EInkDisplayHandler*>(lv_indev_get_user_data(indev));

  // Disable touch input during display refresh (BUSY)
  if (handler->is_busy()) {
    data->state = LV_INDEV_STATE_RELEASED;
    data->continue_reading = false;
    return;
  }

  esp_lcd_touch_handle_t tp_handle = handler->get_touch_handle();
  if (tp_handle == nullptr) {
    data->state = LV_INDEV_STATE_RELEASED;
    return;
  }


  // Read touch data from GT911
  esp_err_t ret = esp_lcd_touch_read_data(tp_handle);
  if (ret == ESP_OK) {
    uint8_t touch_cnt = 0;
    // Get touch data using new API
    esp_lcd_touch_point_data_t point_data[1];
    esp_lcd_touch_get_data(tp_handle, point_data, &touch_cnt, 1);

    if (touch_cnt > 0) {
      ESP_LOGI(TAG, "Touch data read successfully: x=%d, y=%d", point_data[0].x, point_data[0].y);
      data->point.x = point_data[0].x;
      data->point.y = point_data[0].y;
      data->state = LV_INDEV_STATE_PRESSED;
    } else {
      data->state = LV_INDEV_STATE_RELEASED;
    }
  } else {
    data->state = LV_INDEV_STATE_RELEASED;
  }

  data->continue_reading = false;
}

void EInkDisplayHandler::_perform_full_refresh(const uint8_t* framebuffer) {
  ESP_LOGI(TAG, "Starting full refresh (3 seconds)...");

  _wait_for_busy();

  spi_transaction_t* rtrans;  // Declare once for entire function

  // Step 1: Write old data (0x10) - typically all zeros for full refresh
  epd_write_cmd(0x10);

  if (xSemaphoreTake(_spi_mutex, pdMS_TO_TICKS(5000)) != pdTRUE) {
    ESP_LOGE(TAG, "SPI mutex timeout in full refresh step 1");
    return;
  }
  gpio_set_level(PIN_DC, 1);  // Data mode

  ESP_LOGI(TAG, "Starting SPI data transmission for old data (0x10)...");

  // Use simpler polling transmission instead of queued to avoid complexity
  static uint8_t zero_byte = 0x00;  // Static to persist
  esp_err_t ret;
  for (size_t i = 0; i < DISPLAY_BUFFER_SIZE; i++) {
    spi_transaction_t t = {};
    t.length = 8;
    t.tx_buffer = &zero_byte;

    ret = spi_device_polling_transmit(_spi, &t);
    if (ret != ESP_OK) {
      ESP_LOGE(TAG, "Failed to send SPI byte %zu: %s", i, esp_err_to_name(ret));
      break;
    }

    // Yield every 1000 bytes to prevent watchdog timeout
    if (i % 1000 == 999) {
      xSemaphoreGive(_spi_mutex);
      vTaskDelay(pdMS_TO_TICKS(1));  // Small delay
      if (xSemaphoreTake(_spi_mutex, pdMS_TO_TICKS(5000)) != pdTRUE) {
        ESP_LOGE(TAG, "SPI mutex timeout during yield at byte %zu", i);
        return;
      }
      gpio_set_level(PIN_DC, 1);  // Re-set data mode after yield
      ESP_LOGI(TAG, "Transmitted %zu/%zu bytes (%.1f%%)", i + 1, DISPLAY_BUFFER_SIZE,
        (float)(i + 1) * 100.0f / DISPLAY_BUFFER_SIZE);
    }
  }
  ESP_LOGI(TAG, "Completed SPI data transmission for old data");
  xSemaphoreGive(_spi_mutex);

  // Step 2: Write new data (0x13)
  epd_write_cmd(0x13);

  if (xSemaphoreTake(_spi_mutex, pdMS_TO_TICKS(5000)) != pdTRUE) {
    ESP_LOGE(TAG, "SPI mutex timeout in full refresh step 2");
    return;
  }
  gpio_set_level(PIN_DC, 1);  // Data mode

  ESP_LOGI(TAG, "Starting SPI data transmission for new data (0x13)...");

  // Use polling transmission for simplicity and reliability
  for (size_t i = 0; i < DISPLAY_BUFFER_SIZE; i++) {
    uint8_t data_byte = framebuffer[i];  // Write data directly (no inversion)

    spi_transaction_t t = {};
    t.length = 8;
    t.tx_buffer = &data_byte;

    esp_err_t ret = spi_device_polling_transmit(_spi, &t);
    if (ret != ESP_OK) {
      ESP_LOGE(TAG, "Failed to send SPI byte %zu: %s", i, esp_err_to_name(ret));
      break;
    }

    // Yield every 100 bytes to prevent watchdog timeout
    if (i % 100 == 99) {
      xSemaphoreGive(_spi_mutex);
      vTaskDelay(pdMS_TO_TICKS(5));  // Increased delay for better yielding
      if (xSemaphoreTake(_spi_mutex, pdMS_TO_TICKS(5000)) != pdTRUE) {
        ESP_LOGE(TAG, "SPI mutex timeout during yield at byte %zu", i);
        return;
      }
      gpio_set_level(PIN_DC, 1);  // Re-set data mode after yield
      ESP_LOGI(TAG, "Transmitted %zu/%zu bytes (%.1f%%)", i + 1, DISPLAY_BUFFER_SIZE,
        (float)(i + 1) * 100.0f / DISPLAY_BUFFER_SIZE);
    }
  }

  ESP_LOGI(TAG, "Completed SPI data transmission for new data");
  xSemaphoreGive(_spi_mutex);

  // Step 3: Trigger display refresh (DRF)
  epd_write_cmd(0x12);
  // Critical delay - sample code says "!!!The delay here is necessary, 200uS at least!!!"
  vTaskDelay(pdMS_TO_TICKS(10));
  ESP_LOGI(TAG, "Display refresh triggered, BUSY pin: %d", gpio_get_level(PIN_BUSY));

  // Wait for refresh to complete
  _wait_for_busy();

  ESP_LOGI(TAG, "Full refresh complete");
}

void EInkDisplayHandler::_perform_partial_refresh(const uint8_t* framebuffer) {
  ESP_LOGI(TAG, "Starting partial refresh (0.3 seconds)...");

  _wait_for_busy();

  // Step 1: Configure VCOM for partial refresh
  const uint8_t vcom_data[] = { 0xA9, 0x07 };
  epd_write_cmd_with_data(0x50, vcom_data, 2);

  // Step 2: Enter partial refresh mode
  epd_write_cmd(0x91);

  // Step 3: Define partial window (full screen for now)
  // Format: 0x90 + 9 bytes (x_start_H, x_start_L, x_end_H, x_end_L, y_start_H, y_start_L, y_end_H, y_end_L, 0x01)
  // For full screen: x=0 to 799 (0x031F), y=0 to 479 (0x01DF)
  const uint8_t window_data[] = {
    0x00, 0x00,  // x_start = 0
    0x03, 0x1F,  // x_end = 799 (0x31F)
    0x00, 0x00,  // y_start = 0
    0x01, 0xDF,  // y_end = 479 (0x1DF)
    0x01         // PT_SCAN
  };
  epd_write_cmd_with_data(0x90, window_data, 9);

  // Step 4: Write new data (0x13 command)
  epd_write_cmd(0x13);

  if (xSemaphoreTake(_spi_mutex, pdMS_TO_TICKS(5000)) != pdTRUE) {
    ESP_LOGE(TAG, "SPI mutex timeout in partial refresh");
    return;
  }
  gpio_set_level(PIN_DC, 1);  // Data mode

  ESP_LOGI(TAG, "Starting SPI data transmission for partial refresh...");

  // Send data in chunks with task yields to prevent blocking
  static uint8_t tx_byte;
  const size_t CHUNK_SIZE = 1000;  // Send 1000 bytes between yields
  for (size_t i = 0; i < DISPLAY_BUFFER_SIZE; i++) {
    tx_byte = framebuffer[i];  // Write data directly (no inversion)
    spi_transaction_t t = {};
    t.length = 8;
    t.tx_buffer = &tx_byte;
    spi_device_polling_transmit(_spi, &t);

    // Yield to other tasks every CHUNK_SIZE bytes
    if (i % CHUNK_SIZE == (CHUNK_SIZE - 1)) {
      xSemaphoreGive(_spi_mutex);
      vTaskDelay(pdMS_TO_TICKS(1));  // Allow other tasks to run
      if (xSemaphoreTake(_spi_mutex, pdMS_TO_TICKS(5000)) != pdTRUE) {
        ESP_LOGE(TAG, "SPI mutex timeout during partial refresh yield");
        return;
      }
      gpio_set_level(PIN_DC, 1);  // Re-set data mode after yield
      ESP_LOGI(TAG, "Partial refresh progress: %zu/%zu bytes (%.1f%%)", i + 1, DISPLAY_BUFFER_SIZE,
        (float)(i + 1) * 100.0f / DISPLAY_BUFFER_SIZE);
    }
  }
  ESP_LOGI(TAG, "Completed SPI data transmission for partial refresh");
  xSemaphoreGive(_spi_mutex);

  // Step 5: Trigger partial display refresh (DRF)
  epd_write_cmd(0x12);
  // Critical delay - sample code says "!!!The delay here is necessary, 200uS at least!!!"
  vTaskDelay(pdMS_TO_TICKS(10));
  ESP_LOGI(TAG, "Partial refresh triggered, BUSY pin: %d", gpio_get_level(PIN_BUSY));

  // Wait for refresh to complete
  _wait_for_busy();

  // Step 6: Exit partial refresh mode
  epd_write_cmd(0x92);

  ESP_LOGI(TAG, "Partial refresh complete");
}

void EInkDisplayHandler::_wait_for_busy() {
  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;
  }

  int timeout = 0;
  while (gpio_get_level(PIN_BUSY) == BUSY_ACTIVE_LEVEL) {  // 0=BUSY, 1=FREE
    vTaskDelay(pdMS_TO_TICKS(100));
    timeout++;
    if (timeout > 100) {  // 10 second timeout
      ESP_LOGE(TAG, "Display BUSY timeout! Pin level: %d", gpio_get_level(PIN_BUSY));
      ESP_LOGW(TAG, "Attempting hardware reset...");

      // Hardware reset sequence
      gpio_set_level(PIN_RST, 0);
      vTaskDelay(pdMS_TO_TICKS(10));
      gpio_set_level(PIN_RST, 1);
      vTaskDelay(pdMS_TO_TICKS(100));

      // Re-initialize display
      ESP_LOGI(TAG, "Re-initializing display after reset...");
      _epd_init();

      // Check if reset worked
      int reset_timeout = 0;
      while (gpio_get_level(PIN_BUSY) == BUSY_ACTIVE_LEVEL) {
        vTaskDelay(pdMS_TO_TICKS(100));
        reset_timeout++;
        if (reset_timeout > 50) {  // 5 second timeout after reset
          ESP_LOGE(TAG, "Display reset failed! Still busy after reset.");
          break;
        }
      }

      if (gpio_get_level(PIN_BUSY) != BUSY_ACTIVE_LEVEL) {
        ESP_LOGI(TAG, "Display reset successful after %d tenths of a second", reset_timeout);
      }
      break;
    }

    // Log every 2 seconds to track progress
    if (timeout % 20 == 0) {
      ESP_LOGW(TAG, "Still waiting for BUSY pin, timeout: %d/100, level: %d",
        timeout, gpio_get_level(PIN_BUSY));
    }
  }
  ESP_LOGI(TAG, "Display ready after %d tenths of a second", timeout);
}

void EInkDisplayHandler::_convert_buffer_to_epaper(const uint8_t* lvgl_buf, uint8_t* epd_buf, size_t size) {
  // LVGL 1-bit format is already compatible with e-paper
  // Just copy directly
  memcpy(epd_buf, lvgl_buf, size);
}