#include "fast_jpeg.h"

#include "esphome/core/application.h"
#include "esphome/core/log.h"

#include <JPEGDEC.h>
#include "esp_heap_caps.h"
#include "esp_http_client.h"

#include <algorithm>
#include <cstring>
#include <new>
#include <strings.h>
#include <utility>

namespace esphome {
namespace fast_jpeg {

static const char *const TAG = "fast_jpeg";
static const int kMaxWidth = 720;
static const int kMaxHeight = 720;
static const size_t kJpegCap = 128 * 1024;

static int align16_(int value) { return (value + 15) & ~15; }

static void *psram_alloc_(size_t size) {
  void *ptr = heap_caps_malloc(size, MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT);
  if (ptr == nullptr)
    ptr = heap_caps_malloc(size, MALLOC_CAP_8BIT);
  return ptr;
}

int jpeg_draw_cb_(JPEGDRAW *draw) {
  auto *self = static_cast<FastJpeg *>(draw->pUser);
  if (self == nullptr || self->back_buffer() == nullptr)
    return 1;

  const int img_w = self->pending_width();
  const int img_h = self->pending_height();
  const int stride = self->pending_stride();
  if (img_w <= 0 || img_h <= 0 || stride <= 0)
    return 1;

  // JPEGDEC emits MCU tiles that can extend past the true image edge (270 is not
  // 16-pixel aligned). Clip or we write off the end of the framebuffer and reboot.
  const int x0 = draw->x;
  const int y0 = draw->y;
  if (x0 >= img_w || y0 >= img_h)
    return 1;

  const int copy_w = std::min(draw->iWidthUsed > 0 ? draw->iWidthUsed : draw->iWidth, img_w - x0);
  const int copy_h = std::min(draw->iHeight, img_h - y0);
  if (copy_w <= 0 || copy_h <= 0)
    return 1;

  auto *dst = reinterpret_cast<uint16_t *>(self->back_buffer());
  const auto *src = reinterpret_cast<const uint16_t *>(draw->pPixels);
  for (int row = 0; row < copy_h; row++) {
    memcpy(dst + (y0 + row) * stride + x0, src + row * draw->iWidth, static_cast<size_t>(copy_w) * sizeof(uint16_t));
  }
  return 1;
}

esp_err_t http_event_cb_(esp_http_client_event_t *evt) {
  auto *self = static_cast<FastJpeg *>(evt->user_data);
  if (self == nullptr)
    return ESP_OK;
  switch (evt->event_id) {
    case HTTP_EVENT_ON_HEADER:
      if (evt->header_key != nullptr && strcasecmp(evt->header_key, "ETag") == 0 && evt->header_value != nullptr) {
        strncpy(self->etag(), evt->header_value, 79);
        self->etag()[79] = '\0';
      }
      break;
    case HTTP_EVENT_ON_DATA:
      if (evt->data_len <= 0)
        break;
      if (*self->http_rx_len() + evt->data_len > self->jpeg_cap()) {
        ESP_LOGE(TAG, "JPEG larger than %u bytes", (unsigned) self->jpeg_cap());
        return ESP_FAIL;
      }
      memcpy(self->jpeg_buf() + *self->http_rx_len(), evt->data, evt->data_len);
      *self->http_rx_len() += evt->data_len;
      break;
    default:
      break;
  }
  return ESP_OK;
}

void FastJpeg::setup() {
  ESP_LOGI(TAG, "Setup");
  this->jpeg_cap_ = kJpegCap;
  this->jpeg_buf_ = static_cast<uint8_t *>(psram_alloc_(this->jpeg_cap_));
  this->decoder_ = new (std::nothrow) ::JPEGDEC();
  if (this->jpeg_buf_ == nullptr || this->decoder_ == nullptr) {
    this->mark_failed();
    ESP_LOGE(TAG, "Failed to allocate JPEG buffers");
    return;
  }

  BaseType_t ok = xTaskCreate(task_fn_, "fast_jpeg", 8192, this, 5, &this->task_);
  if (ok != pdPASS) {
    this->mark_failed();
    ESP_LOGE(TAG, "Failed to create fetch task");
    return;
  }
  this->last_fps_log_ = millis();
}

void FastJpeg::dump_config() {
  ESP_LOGCONFIG(TAG, "Fast JPEG stream");
  ESP_LOGCONFIG(TAG, "  URL: %s", this->url_.c_str());
  ESP_LOGCONFIG(TAG, "  Interval: %u ms", this->get_update_interval());
}

void FastJpeg::update() {
  if (this->is_failed() || this->task_ == nullptr || !this->enabled_.load())
    return;
  if (this->busy_.load() || this->decoded_ready_.load())
    return;
  xTaskNotifyGive(this->task_);
}

void FastJpeg::loop() {
  if (!this->decoded_ready_.exchange(false))
    return;
  std::swap(this->front_, this->back_);
  this->w_ = this->pending_w_;
  this->h_ = this->pending_h_;
  this->stride_ = this->pending_stride_;
  this->clear_bars_once_();
  this->paint_();
  this->frames_++;
  uint32_t now = millis();
  if (now - this->last_fps_log_ >= 5000) {
    ESP_LOGI(TAG, "%u frames in 5s (%.1f fps), last %dx%d", this->frames_, this->frames_ / 5.0f, this->w_, this->h_);
    this->frames_ = 0;
    this->last_fps_log_ = now;
  }
}

void FastJpeg::task_fn_(void *param) {
  auto *self = static_cast<FastJpeg *>(param);
  for (;;) {
    ulTaskNotifyTake(pdTRUE, portMAX_DELAY);
    self->busy_.store(true);
    self->fetch_once_();
    self->busy_.store(false);
  }
}

void FastJpeg::fetch_once_() {
  if (this->url_.empty() || this->jpeg_buf_ == nullptr)
    return;

  esp_http_client_config_t config = {};
  config.url = this->url_.c_str();
  config.timeout_ms = 2000;
  config.event_handler = http_event_cb_;
  config.user_data = this;
  config.keep_alive_enable = true;
  config.method = HTTP_METHOD_GET;

  if (this->http_ == nullptr) {
    this->http_ = esp_http_client_init(&config);
    if (this->http_ == nullptr) {
      ESP_LOGE(TAG, "HTTP init failed");
      return;
    }
  }
  auto *client = static_cast<esp_http_client_handle_t>(this->http_);
  esp_http_client_set_url(client, this->url_.c_str());
  esp_http_client_set_header(client, "Accept", "image/jpeg");
  if (this->etag_[0] != '\0')
    esp_http_client_set_header(client, "If-None-Match", this->etag_);

  this->http_rx_len_ = 0;
  uint32_t t0 = millis();
  esp_err_t err = esp_http_client_perform(client);
  if (err != ESP_OK) {
    ESP_LOGW(TAG, "HTTP failed: %s", esp_err_to_name(err));
    return;
  }
  int status = esp_http_client_get_status_code(client);
  if (status == 304)
    return;
  if (status != 200) {
    ESP_LOGW(TAG, "HTTP status %d", status);
    return;
  }
  uint32_t download_ms = millis() - t0;
  ESP_LOGD(TAG, "downloaded %u bytes in %ums", (unsigned) this->http_rx_len_, (unsigned) download_ms);
  if (!this->decode_jpeg_(this->http_rx_len_))
    return;
}

bool FastJpeg::ensure_framebuffers_(int width, int height) {
  if (width <= 0 || height <= 0 || width > kMaxWidth || height > kMaxHeight) {
    ESP_LOGW(TAG, "Unsupported JPEG size %dx%d", width, height);
    return false;
  }
  const int stride = align16_(width);
  const int rows = align16_(height);
  size_t need = static_cast<size_t>(stride) * static_cast<size_t>(rows) * 2;
  this->pending_stride_ = stride;
  if (need <= this->fb_cap_ && this->front_ != nullptr && this->back_ != nullptr)
    return true;
  heap_caps_free(this->front_);
  heap_caps_free(this->back_);
  this->front_ = nullptr;
  this->back_ = nullptr;
  this->front_ = static_cast<uint8_t *>(psram_alloc_(need));
  this->back_ = static_cast<uint8_t *>(psram_alloc_(need));
  this->fb_cap_ = need;
  if (this->front_ == nullptr || this->back_ == nullptr) {
    ESP_LOGE(TAG, "Framebuffer alloc failed (%u bytes)", (unsigned) need);
    return false;
  }
  memset(this->front_, 0, need);
  memset(this->back_, 0, need);
  return true;
}

bool FastJpeg::decode_jpeg_(size_t len) {
  if (len < 16 || this->decoder_ == nullptr)
    return false;
  ::JPEGDEC *jpeg = this->decoder_;
  if (!jpeg->openRAM(this->jpeg_buf_, static_cast<int>(len), jpeg_draw_cb_)) {
    ESP_LOGW(TAG, "JPEG open failed: %d", jpeg->getLastError());
    return false;
  }
  if (jpeg->getJPEGType() == JPEG_MODE_PROGRESSIVE) {
    ESP_LOGW(TAG, "Progressive JPEG is not supported");
    jpeg->close();
    return false;
  }
  int width = jpeg->getWidth();
  int height = jpeg->getHeight();
  jpeg->setUserPointer(this);
  if (!this->ensure_framebuffers_(width, height)) {
    jpeg->close();
    return false;
  }
  this->pending_w_ = width;
  this->pending_h_ = height;
  jpeg->setPixelType(this->big_endian_ ? ::RGB565_BIG_ENDIAN : ::RGB565_LITTLE_ENDIAN);
  uint32_t t0 = millis();
  if (!jpeg->decode(0, 0, 0)) {
    ESP_LOGW(TAG, "JPEG decode failed: %d", jpeg->getLastError());
    jpeg->close();
    return false;
  }
  jpeg->close();
  ESP_LOGI(TAG, "decoded %dx%d in %ums", width, height, (unsigned) (millis() - t0));
  this->decoded_ready_.store(true);
  return true;
}

void FastJpeg::clear_bars_once_() {
  if (this->bars_cleared_ || this->display_ == nullptr || this->display_->is_failed())
    return;
  const int canvas_w = this->display_->get_width();
  const int canvas_h = this->display_->get_height();
  if (canvas_w <= 0 || canvas_h <= 0)
    return;
  const int y = this->y_ >= 0 ? this->y_ : (canvas_h - this->h_) / 2;
  uint16_t line[kMaxWidth] = {};
  const int line_w = std::min(canvas_w, kMaxWidth);
  for (int row = 0; row < y; row++) {
    this->display_->draw_pixels_at(0, row, line_w, 1, reinterpret_cast<const uint8_t *>(line), display::COLOR_ORDER_RGB,
                                  display::COLOR_BITNESS_565, this->big_endian_);
  }
  for (int row = y + this->h_; row < canvas_h; row++) {
    this->display_->draw_pixels_at(0, row, line_w, 1, reinterpret_cast<const uint8_t *>(line), display::COLOR_ORDER_RGB,
                                  display::COLOR_BITNESS_565, this->big_endian_);
  }
  this->bars_cleared_ = true;
}

void FastJpeg::paint_() {
  if (this->display_ == nullptr || this->display_->is_failed() || this->front_ == nullptr || this->w_ <= 0)
    return;
  const int canvas_h = this->display_->get_height();
  const int y = this->y_ >= 0 ? this->y_ : (canvas_h - this->h_) / 2;
  const int x_pad = this->stride_ > this->w_ ? this->stride_ - this->w_ : 0;
  this->display_->draw_pixels_at(this->x_, y, this->w_, this->h_, this->front_, display::COLOR_ORDER_RGB,
                                display::COLOR_BITNESS_565, this->big_endian_, 0, 0, x_pad);
}

}  // namespace fast_jpeg
}  // namespace esphome