GxEPD2-esphome/esphome-zepd/components/zepd/zepd_display.h

#pragma once

#include "esphome/core/component.h"
#include "esphome/core/helpers.h"
#include "esphome/components/display/display_buffer.h"
#include "esphome/components/spi/spi.h"
#include "engine.h"
#include "engine_model.h"
#include "zepd_model_generated.h"

namespace esphome::zepd {

static const char *const TAG = "zepd";

// ESPHome display component (the engine adapter) for a zEPD model. The
// model arrives as one codegen'd zepd::EngineModel, selected at ESPHome
// compile time by display.py from the drivers/*.json registry (see
// claude/DESCRIPTOR_REGISTRY_PLAN.md) and written to
// zepd_model_generated.h - not templated over per-panel C++ classes the
// way GxEPD2 itself is. Driven by zepd::Engine, the generic replay state
// machine; this class is the ESPHome-side half of the engine adapter.
class ZepdDisplay : public display::DisplayBuffer {
 public:
  void set_cs_pin(GPIOPin *pin) { this->cs_ = pin; }
  void set_dc_pin(GPIOPin *pin) { this->dc_ = pin; }
  void set_reset_pin(GPIOPin *pin) { this->rst_ = pin; }
  void set_busy_pin(GPIOPin *pin) { this->busy_ = pin; }
  void set_spi_parent(spi::SPIComponent *parent) { this->spi_.set_spi_parent(parent); }
  void set_descriptor(const EngineModel *desc) { this->desc_ = desc; }
  void set_invert_mono(bool invert) { this->invert_mono_ = invert; }
  // The short model label test_pattern() can print, derived from the
  // registry's oem_part_numbers[0] (see display.py).
  void set_model_name(const char *name) { this->model_name_ = name; }
  // Whether this model's data is confirmed on real hardware (vs captured
  // but not yet bench-verified) - surfaced in dump_config(). From the
  // registry's `verified` field, via display.py.
  void set_verified(bool v) { this->verified_ = v; }

  // Bench visual-check pattern, bit-depth aware. Everything scales to what
  // THIS panel can actually show (see test_colors_): the color-bar count,
  // the ladder color cycle and the corner markers all follow the panel's
  // real color set instead of a hardcoded 4, so the drawn pattern doubles as
  // a statement of what the engine model expects of the panel. Adds a
  // dithered grey ramp so a 1bpp panel shows a graded sweep, not just two
  // flat blocks.
  //
  // Pass a font to label the panel with its model (and the panel can spare the
  // room); pass nothing to skip the label - handy when you just want a quick
  // pattern without defining a font. From the display lambda:
  //   id(<display_id>).test_pattern();                 // no label
  //   id(<display_id>).test_pattern(id(<font_id>));    // labelled
  void test_pattern(display::BaseFont *font = nullptr) {
    const int w = this->get_width();
    const int h = this->get_height();
    Color colors[8];
    const int n = this->test_colors_(colors, 8);

    this->fill(Color(0xFF, 0xFF, 0xFF));

    const int smin = w < h ? w : h;
    const int border = smin / 100 > 1 ? smin / 100 : 1;
    const int x0 = border;
    int y0 = border;
    const int iw = w - 2 * border;
    int ih = h - 2 * border;

    // Name band, above every test zone. Drawn only when a font was passed and
    // the panel is tall enough to give it a fifth of the interior without
    // crowding the zones below.
    if (font != nullptr && this->model_name_ != nullptr && this->model_name_[0] != '\0') {
      int tw, xoff, baseline, th;
      font->measure(this->model_name_, &tw, &xoff, &baseline, &th);
      const int band = th + 4;
      if (band <= ih / 5) {
        this->print(w / 2, y0 + band / 2, font, Color(0x00, 0x00, 0x00), display::TextAlign::CENTER, this->model_name_);
        y0 += band;
        ih -= band;
      }
    }

    // Vertical budget for the four zones below any name band.
    const int bars_h = ih * 34 / 100;
    int dither_h = ih * 12 / 100;
    if (dither_h < 8) dither_h = 8;
    const int row_h = ih * 18 / 100 > 2 ? ih * 18 / 100 : 2;
    const int dy = y0 + bars_h;
    const int ry = dy + dither_h;
    const int col_y = ry + row_h;
    const int col_h = (y0 + ih) - col_y;

    // Color bars: one solid block per available color, in palette/wire-code
    // order (black, white, then the accents). A wrong or missing block
    // points straight at the color - or the wire code - it stands for.
    const int bar_w = iw / n;
    for (int i = 0; i < n; i++) {
      const int bx = x0 + i * bar_w;
      const int bw = (i == n - 1) ? (x0 + iw - bx) : bar_w;
      this->filled_rectangle(bx, y0, bw, bars_h, colors[i]);
    }

    // Dithered grey ramp: brightness 255 (left) -> 0 (right) via an ordered
    // 4x4 Bayer threshold. On a 1bpp panel this is the only zone that shows
    // a graded sweep; on deeper panels it's a quick mono/uniformity check.
    static const uint8_t BAYER4[4][4] = {
        {0, 8, 2, 10}, {12, 4, 14, 6}, {3, 11, 1, 9}, {15, 7, 13, 5}};
    for (int yy = 0; yy < dither_h; yy++) {
      for (int xx = 0; xx < iw; xx++) {
        const int bright = (iw <= 1) ? 255 : 255 - xx * 255 / (iw - 1);
        const int thr = (BAYER4[yy & 3][xx & 3] * 2 + 1) * 255 / 32;
        if (bright <= thr)
          this->draw_pixel_at(x0 + xx, dy + yy, Color(0x00, 0x00, 0x00));
      }
    }

    // Row ladder: 1px rows cycling the full color set - a dropped, doubled
    // or reordered row breaks the N-row cadence.
    for (int yy = 0; yy < row_h; yy++)
      this->horizontal_line(x0, ry + yy, iw, colors[yy % n]);

    // Column ladder: the same rotated 90 degrees, catching nibble/column
    // order bugs the row ladder can't (packed formats put several pixels per
    // byte, so a column-order fault can leave every row looking fine).
    for (int xx = 0; xx < iw; xx++)
      this->vertical_line(x0 + xx, col_y, col_h, colors[xx % n]);

    // Corner markers for orientation, using the panel's ink colors (every
    // color but white). BL is always black; TL/TR take the next inks when
    // the panel has them; BR is left blank as the distinct "no marker"
    // state. A 90/180 rotation or mirror can't survive this even though the
    // bars and ladders are otherwise symmetric.
    Color ink[8];
    int ni = 0;
    for (int i = 0; i < n; i++) {
      if (colors[i].red == 0xFF && colors[i].green == 0xFF && colors[i].blue == 0xFF)
        continue;  // skip white - a white marker is invisible
      ink[ni++] = colors[i];
    }
    const int mk = smin / 12 > 3 ? smin / 12 : 3;
    if (ni > 0) this->filled_rectangle(x0, y0 + ih - mk, mk, mk, ink[0]);        // BL
    if (ni > 1) this->filled_rectangle(x0, y0, mk, mk, ink[1]);                  // TL
    if (ni > 2) this->filled_rectangle(x0 + iw - mk, y0, mk, mk, ink[2]);        // TR

    // Corner-to-corner diagonals over everything: the cheapest coordinate-
    // transform check - a clean line means no row/column reordering.
    this->line(0, 0, w - 1, h - 1, Color(0x00, 0x00, 0x00));
    this->line(w - 1, 0, 0, h - 1, Color(0x00, 0x00, 0x00));

    // Border frame last, so it stays an unbroken rectangle on all four edges
    // even where a zone fills to the edge - catches a clipped/offset origin.
    this->rectangle(0, 0, w, h, Color(0x00, 0x00, 0x00));
  }

  void setup() override {
    const auto *desc = this->desc_;
    // window_per_plane panels (e.g. a master+slave dual-region panel) size
    // each plane's own buffer from plane_width (each region's own local
    // width), not the descriptor's overall width (the full logical canvas
    // the adapter draws into) - both regions are the same size here, so
    // one plane_bytes_ value covers both buffers.
    uint32_t w = desc->window_per_plane ? desc->plane_width : desc->width;
    // width*height*bpp/8: /8 at 1bpp, /4 at 2bpp (4-color), /2 at 4bpp (7-color).
    this->plane_bytes_ = (uint32_t) w * desc->height * desc->bits_per_pixel / 8;
    RAMAllocator<uint8_t> allocator;
    this->plane0_ = allocator.allocate(this->plane_bytes_);
    bool need_plane1 = desc->plane_count > 1;
    this->plane1_ = need_plane1 ? allocator.allocate(this->plane_bytes_) : nullptr;
    if (this->plane0_ == nullptr || (need_plane1 && this->plane1_ == nullptr)) {
      ESP_LOGE(TAG, "Failed to allocate %u-byte plane buffer(s)", (unsigned) this->plane_bytes_);
      this->mark_failed();
      return;
    }
    if (desc->bits_per_pixel > 1) {
      // Packed modes: clear every pixel to the panel's white/blank code,
      // replicated across the byte - e.g. 4-color white=0x01 -> 0x55, 7-color
      // white=0x01 -> 0x11, matching GxEPD2's own writeScreenBuffer() fill
      // (color4()*0x55 / color7()|<<4). Both regions of a window_per_plane
      // panel (master/slave) get the same blank fill.
      uint8_t fill = replicate_code_(desc->blank_code, desc->bits_per_pixel);
      memset(this->plane0_, fill, this->plane_bytes_);
      if (this->plane1_ != nullptr)
        memset(this->plane1_, fill, this->plane_bytes_);
    } else {
      // 1bpp plane model: plane0 white = all bits set (0xFF); plane1 (the
      // accent/red plane, when present) = 0 = not-red.
      memset(this->plane0_, 0xFF, this->plane_bytes_);
      if (this->plane1_ != nullptr)
        memset(this->plane1_, 0x00, this->plane_bytes_);
    }

    this->spi_.spi_setup();
    this->engine_.begin(desc, this->cs_, this->dc_, this->rst_, this->busy_, &this->spi_);
  }

  // Arms the state machine and returns immediately - the real several-
  // seconds-long write/refresh/busy-wait cycle happens across many
  // subsequent loop() ticks, not before this call returns.
  void update() override {
    if (!this->engine_.idle()) {
      ESP_LOGW(TAG, "Skipping update: previous display refresh is still running");
      return;
    }
    this->do_update_();
    this->engine_.request_full_update(this->plane0_, this->plane1_);
  }

  void loop() override { this->engine_.step(); }

  // True once the engine has returned to idle after the last update()/
  // power_off()/hibernate() call. Lets a more complex application (a lambda
  // chaining several updates, a deep_sleep automation deciding whether to
  // wait) check "is the panel still mid-refresh" instead of guessing from
  // timing - the same thing power_off()/hibernate()'s own guards below check
  // internally.
  bool idle() const { return this->engine_.idle(); }

  // Turns off the panel's driving-voltage generation (reduces fading over
  // time between refreshes). Call from a lambda, e.g.
  // id(<display_id>).power_off(); - not called automatically by update(),
  // since this engine does a full hardware reset every cycle regardless of
  // prior power state, so correctness never depends on it. (It IS called
  // automatically from on_powerdown() below.) Safe to call mid-refresh: see
  // zepd::Engine::request_power_off().
  void power_off() {
    if (!this->engine_.idle())
      ESP_LOGW(TAG, "power_off() called mid-refresh; aborting the in-progress update");
    this->engine_.request_power_off();
  }

  // power_off() + deep sleep, wakeable only by a hardware reset pulse - which
  // the next update() always issues (request_full_update()'s RESET_LOW), so
  // no paired "wake" call is needed; just call update() as usual. Safe to
  // call mid-refresh: see zepd::Engine::request_hibernate().
  void hibernate() {
    if (!this->engine_.idle())
      ESP_LOGW(TAG, "hibernate() called mid-refresh; aborting the in-progress update");
    this->engine_.request_hibernate();
  }

  // ESPHome calls this on every component right before deep_sleep actually
  // suspends the MCU, and on safe_reboot() - see
  // Application::run_powerdown_hooks() (fires after teardown_components()
  // completes or times out, deep_sleep's should_teardown_() is unconditional
  // on every supported platform). Powers the panel down with no YAML/lambda
  // needed for the common case: hibernate() when a reset pin is wired (the
  // next update()'s unconditional RESET_LOW wakes it - see hibernate()
  // above), else power_off() (needs no reset pulse to recover, so it's safe
  // with no rst pin at all).
  //
  // Deliberately NOT paired with a teardown() override to wait out a
  // mid-flight refresh: ESPHome's teardown budget (1s on reboot, 5s on deep
  // sleep) is far shorter than this engine's slower busy_timeout_ms classes
  // (up to 45-60s for 4bpp/7-color panels), so stalling teardown here would
  // just lose that race every time and get force-advanced anyway - the
  // budget is sized for a network flush (see api_server::teardown()), not a
  // hardware operation. If a refresh is still running when this fires,
  // power_off()/hibernate() abort it with a reset pulse first rather than
  // writing a power-down command over a controller that's still asserting
  // BUSY and would just ignore it (see zepd::Engine::request_power_off()). A
  // caller who cares about that race can check idle() (above) and hold off
  // triggering deep sleep themselves.
  void on_powerdown() override {
    if (this->rst_ != nullptr) {
      this->hibernate();
    } else {
      this->power_off();
    }
  }

#ifdef USE_ZEPD_DIAGNOSTICS
  // Dev-time diagnostics - compiled out entirely unless this display's YAML
  // sets `diagnostics: true` (display.py then adds -DUSE_ZEPD_DIAGNOSTICS),
  // so a production build pays nothing for them. A stash for throwaway bench
  // tests, not a stable API - add more here as needed, don't feel obliged to
  // keep any of it around once a test has done its job.
  //
  // Draws the current lambda/pages content and pushes it WITHOUT the reset
  // pulse update() normally does first - see
  // zepd::Engine::diag_request_update_no_reset(). Call from a lambda, e.g. to
  // check whether power_off()/hibernate() actually took: call one of those,
  // then this with a pattern distinct from what's already on screen - no
  // change (or odd BUSY behavior) means the panel really is off-line; follow
  // with a normal update() to confirm the next real reset wakes it and to
  // leave the panel in a known state either way.
  void diag_update_no_reset() {
    if (!this->engine_.idle()) {
      ESP_LOGW(TAG, "Skipping diag_update_no_reset: display refresh still in progress");
      return;
    }
    this->do_update_();
    this->engine_.diag_request_update_no_reset(this->plane0_, this->plane1_);
  }
#endif

  void dump_config() override {
    ESP_LOGCONFIG(TAG, "zEPD Display:");
    ESP_LOGCONFIG(TAG, "  Model: %s", this->model_name_ != nullptr ? this->model_name_ : "(unknown)");
    ESP_LOGCONFIG(TAG, "  Dimensions: %ux%u", (unsigned) this->desc_->width, (unsigned) this->desc_->height);
    ESP_LOGCONFIG(TAG, "  Bits/pixel: %u (%s)", (unsigned) this->desc_->bits_per_pixel,
                  this->desc_->bits_per_pixel == 1 ? "mono/plane" : "packed palette");
    ESP_LOGCONFIG(TAG, "  Planes/regions: %u", (unsigned) this->desc_->plane_count);
    ESP_LOGCONFIG(TAG, "  Busy active level: %s", this->desc_->busy_active_level ? "HIGH" : "LOW");
    if (this->desc_->bits_per_pixel == 1)
      ESP_LOGCONFIG(TAG, "  invert_mono: %s", YESNO(this->invert_mono_));
    ESP_LOGCONFIG(TAG, "  Hardware-verified: %s", YESNO(this->verified_));
    if (!this->verified_)
      ESP_LOGW(TAG, "Model %s is not yet hardware-verified - the captured data may be wrong; "
                    "if it renders correctly, please report it so it can be marked verified",
               this->model_name_ != nullptr ? this->model_name_ : "(unknown)");
    LOG_UPDATE_INTERVAL(this);
  }

  display::DisplayType get_display_type() override { return display::DisplayType::DISPLAY_TYPE_COLOR; }

 protected:
  // Replicates a packed wire code across a whole byte (one code per pixel
  // slot), e.g. 2bpp code 0x01 -> 0x55, 4bpp code 0x01 -> 0x11.
  static uint8_t replicate_code_(uint8_t code, uint8_t bpp) {
    uint8_t mask = static_cast<uint8_t>((1u << bpp) - 1);
    uint8_t out = 0;
    for (uint8_t i = 0; i < 8 / bpp; i++)
      out = static_cast<uint8_t>((out << bpp) | (code & mask));
    return out;
  }

  // Maps a drawn RGB color to its panel wire code via the descriptor's
  // palette: exact matches win (so named colors are deterministic, like
  // GxEPD2's color4()/color7() switch), otherwise nearest in RGB space.
  static uint8_t classify_(Color c, const EpdPaletteEntry *pal, uint8_t len) {
    if (pal == nullptr || len == 0)
      return 0;
    uint8_t best = 0;
    long best_d = -1;
    for (uint8_t i = 0; i < len; i++) {
      long dr = static_cast<long>(c.red) - pal[i].r;
      long dg = static_cast<long>(c.green) - pal[i].g;
      long db = static_cast<long>(c.blue) - pal[i].b;
      long d = dr * dr + dg * dg + db * db;
      if (d == 0)
        return pal[i].code;
      if (best_d < 0 || d < best_d) {
        best_d = d;
        best = i;
      }
    }
    return pal[best].code;
  }

  // Packed code-per-pixel model (bits_per_pixel > 1): N bits/pixel, MSB
  // first, matching GxEPD2's own out_data packing in writeImage (4-color
  // 4px/byte, 7-color 2px/byte). Routes to the correct region for a
  // window_per_plane (master/slave) panel the same way the mono path does.
  void draw_packed_(int x, int y, Color color) {
    const uint8_t bpp = this->desc_->bits_per_pixel;
    uint8_t code = classify_(color, this->desc_->palette, this->desc_->palette_len);
    uint16_t stride = this->desc_->width;
    uint8_t *buf = this->plane0_;
    int local_x = x;
    if (this->desc_->window_per_plane && this->desc_->plane_count > 1) {
      stride = this->desc_->plane_width;
      if (x >= stride) {
        buf = this->plane1_;
        local_x = x - stride;
      }
    }
    uint32_t px = static_cast<uint32_t>(y) * stride + local_x;
    const uint8_t ppb = 8 / bpp;
    uint32_t idx = px / ppb;
    uint8_t shift = static_cast<uint8_t>((ppb - 1 - (px % ppb)) * bpp);
    uint8_t mask = static_cast<uint8_t>((1u << bpp) - 1);
    buf[idx] = static_cast<uint8_t>((buf[idx] & ~(mask << shift)) | ((code & mask) << shift));
  }

  void draw_absolute_pixel_internal(int x, int y, Color color) override {
    if (x < 0 || y < 0 || x >= this->desc_->width || y >= this->desc_->height)
      return;
    if (this->desc_->bits_per_pixel > 1) {
      this->draw_packed_(x, y, color);
      return;
    }
    bool red = color.red > color.green && color.red > color.blue && color.red > 0x7f;
    bool bright = color.red > 0x7f || color.green > 0x7f || color.blue > 0x7f;
    uint32_t idx = (static_cast<uint32_t>(y) * this->desc_->width + x) / 8;
    uint8_t mask = 0x80 >> (x % 8);
    // Red/accent pixels go to the color plane (when the panel has one),
    // everything else is black/white on the black plane. get_display_type()
    // reports DISPLAY_TYPE_COLOR, so callers draw literal RGB colors, not
    // the on/off is_on() convention, even on a plane_count==1 (BW-only)
    // panel - there just is no color plane for a "red" classification to
    // land on.
    //
    // Buffer convention here is 1=red/0=not-red, transformed to the wire
    // value by plane_invert[1].
    if (red && this->plane1_ != nullptr) {
      this->plane1_[idx] |= mask;
    } else {
      if (this->plane1_ != nullptr)
        this->plane1_[idx] &= ~mask;
      // plane0 bit set = white/blank. invert_mono_ flips the mono sense so a
      // bright color (e.g. ESPHome's default COLOR_ON) becomes ink, not blank
      // - see display.py's CONF_INVERT_MONO. Red classification is unaffected.
      if (bright ^ this->invert_mono_) {
        this->plane0_[idx] |= mask;
      } else {
        this->plane0_[idx] &= ~mask;
      }
    }
  }
  int get_width_internal() override { return this->desc_->width; }
  int get_height_internal() override { return this->desc_->height; }

  // The colors this panel can actually show, each as a representative RGB, in
  // a stable order (black, white, then accents). Packed panels (>1bpp): one
  // entry per distinct palette wire code, first RGB seen for each. 1bpp:
  // black + white, plus red when a second plane exists for the accent (only
  // red-classified pixels reach it - see draw_absolute_pixel_internal). This
  // is what makes test_pattern() bit-depth aware. Returns the count.
  int test_colors_(Color *out, int max) {
    int n = 0;
    if (this->desc_->bits_per_pixel > 1 && this->desc_->palette != nullptr) {
      uint8_t seen[16];
      int seen_n = 0;
      for (int i = 0; i < this->desc_->palette_len && n < max; i++) {
        const uint8_t code = this->desc_->palette[i].code;
        bool dup = false;
        for (int j = 0; j < seen_n; j++)
          if (seen[j] == code) { dup = true; break; }
        if (dup) continue;
        if (seen_n < (int) sizeof(seen)) seen[seen_n++] = code;
        const auto &e = this->desc_->palette[i];
        out[n++] = Color(e.r, e.g, e.b);
      }
    } else {
      if (n < max) out[n++] = Color(0x00, 0x00, 0x00);  // black
      if (n < max) out[n++] = Color(0xFF, 0xFF, 0xFF);  // white
      if (this->desc_->plane_count >= 2 && n < max) out[n++] = Color(0xFF, 0x00, 0x00);  // accent
    }
    return n;
  }

  const EngineModel *desc_{nullptr};
  const char *model_name_{nullptr};
  bool verified_{false};
  GPIOPin *cs_{nullptr};
  GPIOPin *dc_{nullptr};
  GPIOPin *rst_{nullptr};
  GPIOPin *busy_{nullptr};
  EpdSPIDevice spi_;
  Engine engine_;
  uint8_t *plane0_{nullptr};
  uint8_t *plane1_{nullptr};
  uint32_t plane_bytes_{0};
  bool invert_mono_{false};
};

}  // namespace esphome::zepd