#include "png_decoder.h"
#ifdef USE_RUNTIME_IMAGE_PNG

#include "esphome/components/display/display_buffer.h"
#include "esphome/core/application.h"
#include "esphome/core/helpers.h"
#include "esphome/core/log.h"

static const char *const TAG = "image_decoder.png";

namespace esphome::runtime_image {

/**
 * @brief Callback method that will be called by the PNGLE engine when the basic
 * data of the image is received (i.e. width and height);
 *
 * @param pngle The PNGLE object, including the context data.
 * @param w The width of the image.
 * @param h The height of the image.
 */
static void init_callback(pngle_t *pngle, uint32_t w, uint32_t h) {
  PngDecoder *decoder = (PngDecoder *) pngle_get_user_data(pngle);
  decoder->set_size(w, h);
}

/**
 * @brief Callback method that will be called by the PNGLE engine when a chunk
 * of the image is decoded.
 *
 * @param pngle The PNGLE object, including the context data.
 * @param x The X coordinate to draw the rectangle on.
 * @param y The Y coordinate to draw the rectangle on.
 * @param w The width of the rectangle to draw.
 * @param h The height of the rectangle to draw.
 * @param rgba The color to paint the rectangle in.
 */
static void draw_callback(pngle_t *pngle, uint32_t x, uint32_t y, uint32_t w, uint32_t h, const uint8_t rgba[4]) {
  PngDecoder *decoder = (PngDecoder *) pngle_get_user_data(pngle);
  Color color(rgba[0], rgba[1], rgba[2], rgba[3]);
  decoder->draw(x, y, w, h, color);

  // Feed watchdog periodically to avoid triggering during long decode operations.
  // Feed every 1024 pixels to balance efficiency and responsiveness.
  uint32_t pixels = w * h;
  decoder->increment_pixels_decoded(pixels);
  if ((decoder->get_pixels_decoded() % 1024) < pixels) {
    App.feed_wdt();
  }
}

PngDecoder::PngDecoder(RuntimeImage *image) : ImageDecoder(image) {
  {
    pngle_t *pngle = this->allocator_.allocate(1, PNGLE_T_SIZE);
    if (!pngle) {
      ESP_LOGE(TAG, "Failed to allocate memory for PNGLE engine!");
      return;
    }
    memset(pngle, 0, PNGLE_T_SIZE);
    pngle_reset(pngle);
    this->pngle_ = pngle;
  }
}

PngDecoder::~PngDecoder() {
  if (this->pngle_) {
    pngle_reset(this->pngle_);
    this->allocator_.deallocate(this->pngle_, PNGLE_T_SIZE);
  }
}

int PngDecoder::prepare(size_t expected_size) {
  ImageDecoder::prepare(expected_size);
  if (!this->pngle_) {
    ESP_LOGE(TAG, "PNG decoder engine not initialized!");
    return DECODE_ERROR_OUT_OF_MEMORY;
  }
  pngle_set_user_data(this->pngle_, this);
  pngle_set_init_callback(this->pngle_, init_callback);
  pngle_set_draw_callback(this->pngle_, draw_callback);
  return 0;
}

int HOT PngDecoder::decode(uint8_t *buffer, size_t size) {
  if (!this->pngle_) {
    ESP_LOGE(TAG, "PNG decoder engine not initialized!");
    return DECODE_ERROR_OUT_OF_MEMORY;
  }
  // PNG can be decoded progressively, but wait for a reasonable chunk
  if (size < 256 && this->expected_size_ > 0 && size < this->expected_size_ - this->decoded_bytes_) {
    ESP_LOGD(TAG, "Waiting for more data");
    return 0;
  }
  auto fed = pngle_feed(this->pngle_, buffer, size);
  if (fed < 0) {
    ESP_LOGE(TAG, "Error decoding image: %s", pngle_error(this->pngle_));
  } else {
    this->decoded_bytes_ += fed;
  }
  return fed;
}

}  // namespace esphome::runtime_image

#endif  // USE_RUNTIME_IMAGE_PNG