#pragma once

#include <cstdint>
#include <cstdio>

#include "esphome/core/gpio.h"
#include "esphome/core/hal.h"
#include "esphome/core/log.h"
#include "esphome/components/spi/spi.h"
#include "epd_command.h"

namespace esphome::gxepd2 {

// Stock GxEPD2 vendor value (20MHz) for this panel. Was dropped to 10MHz
// (matching epaper_spi) while chasing the solid-black/solid-red symptom,
// which turned out to be unrelated - display.py never wired the configured
// lambda to set_writer() at all, so nothing but auto_clear's default fill
// ever rendered, regardless of SPI clock. Rolled back to test the stock
// GxEPD2 clock now that the real bug is fixed.
using EpdSPIDevice = spi::SPIDevice<spi::BIT_ORDER_MSB_FIRST, spi::CLOCK_POLARITY_LOW, spi::CLOCK_PHASE_LEADING,
                                     spi::DATA_RATE_20MHZ>;

// Replays one GxEPD2 driver's refresh cycle - reset pulse, init command table,
// RAM-window set, per-plane byte burst, refresh trigger, busy-pin wait - as a
// state machine stepped from loop(), instead of GxEPD2's own synchronous
// top-to-bottom call stack. step() never blocks: every state does a small
// bounded unit of work and returns, so nothing here can starve another task or
// need a watchdog subscription.
//
// request_power_off()/request_hibernate() are a separate, standalone entry
// point (no init/window/plane-write) replaying just
// power_off_sequence/hibernate_sequence - for a user's own lambda, not run
// automatically by request_full_update(). Called while a refresh is still
// in flight (e.g. from on_powerdown()'s teardown path), they force a reset
// pulse first to abort it rather than writing over a busy controller - see
// begin_power_down_(). See epd_command.h.
class CycleEngine {
 public:
  void begin(const EpdPanelDescriptor *desc, GPIOPin *cs, GPIOPin *dc, GPIOPin *rst, GPIOPin *busy,
             EpdSPIDevice *spi) {
    this->desc_ = desc;
    this->cs_ = cs;
    this->dc_ = dc;
    this->rst_ = rst;
    this->busy_ = busy;
    this->spi_ = spi;
    this->cs_->setup();
    this->dc_->setup();
    if (this->rst_ != nullptr)
      this->rst_->setup();
    if (this->busy_ != nullptr)
      this->busy_->setup();
    this->cs_->digital_write(true);
    this->dc_->digital_write(true);
  }

  // Arms the state machine; returns immediately. planes[1] is ignored when
  // desc->plane_count == 1. Undefined to call while !idle().
  void request_full_update(const uint8_t *plane0, const uint8_t *plane1) {
    this->arm_full_update_(plane0, plane1, this->rst_ != nullptr ? State::RESET_LOW : State::INIT_CMD);
  }

#ifdef USE_GXEPD2_DIAGNOSTICS
  // Dev-time diagnostic, compiled out unless the display's `diagnostics:
  // true` is set (display.py adds -DUSE_GXEPD2_DIAGNOSTICS): arms a full
  // update starting at INIT_CMD, skipping the reset pulse entirely - the
  // same path a panel with no reset pin wired already takes every cycle.
  // Lets a bench test use the screen as the oracle for "is the panel
  // actually non-responsive right now" (e.g. right after hibernate()): if a
  // distinct pattern drawn this way does NOT appear, the prior power call
  // really did take the panel off-line; a normal request_full_update() right
  // after (reset pulse included) should then wake it and render correctly.
  void diag_request_update_no_reset(const uint8_t *plane0, const uint8_t *plane1) {
    this->arm_full_update_(plane0, plane1, State::INIT_CMD);
  }
#endif

  // Arms a standalone powerOff() replay - no init/window/plane write, just
  // the captured power_off_sequence (if any) then DONE. Safe to call any
  // time, including mid-refresh: see begin_power_down_() below for why that
  // matters. Safe to call repeatedly: GxEPD2's own powerOff() is itself
  // guarded on "already off," which the capture already baked in (see
  // make_power_sequence_resolvers in emit_descriptor.h).
  void request_power_off() { this->begin_power_down_(false); }

  // Arms a standalone hibernate() replay: power_off_sequence then
  // hibernate_sequence (the deep-sleep trigger only - see epd_command.h).
  // See begin_power_down_() below for the mid-refresh case. No paired "wake"
  // call exists or is needed: request_full_update()'s RESET_LOW is an
  // unconditional hardware reset pulse every cycle, which is exactly how
  // GxEPD2 wakes a hibernating (RST-wakeable) panel.
  void request_hibernate() { this->begin_power_down_(true); }

  bool idle() const { return this->state_ == State::IDLE; }

  // Called once per loop() tick. Does one bounded unit of work.
  void step() {
    switch (this->state_) {
      case State::IDLE:
        return;
      case State::RESET_LOW:
        this->rst_->digital_write(true);  // preset high, then pulse low - see GxEPD2_EPD::_reset()
        this->deadline_ = millis() + 10;
        this->set_state_(State::RESET_LOW_WAIT);
        return;
      case State::RESET_LOW_WAIT:
        if (millis() < this->deadline_)
          return;
        this->rst_->digital_write(false);
        this->deadline_ = millis() + 10;
        this->set_state_(State::RESET_HIGH_WAIT);
        return;
      case State::RESET_HIGH_WAIT:
        if (millis() < this->deadline_)
          return;
        this->rst_->digital_write(true);
        this->deadline_ = millis() + 10;
        this->set_state_(State::RESET_SETTLE);
        return;
      case State::RESET_SETTLE:
        if (millis() < this->deadline_)
          return;
        this->cursor_ = 0;
        this->cmd_phase_ = CmdPhase::BEFORE_DELAY;
        // Post-reset busy wait, matching every vendor _InitDisplay()'s
        // _waitWhileBusy() right after the reset pulse. Some controllers hold
        // BUSY while they load settings/waveforms from OTP after a hardware
        // reset (notably JD79676 / GDEY0213F52, whose init is otherwise just
        // 0xE9 + power-on) - issuing any command before that clears gets it
        // silently dropped, so the panel never refreshes. Panels that don't
        // assert BUSY post-reset read "not busy" on the first poll and fall
        // straight through to INIT_CMD, so this is safe for all of them.
        // begin_power_down_()'s abort path sets reset_then_wait_busy_ to
        // false instead - see there for why.
        if (this->reset_then_wait_busy_) {
          this->start_busy_wait_(this->reset_done_next_);
        } else {
          this->set_state_(this->reset_done_next_);
        }
        return;

      case State::INIT_CMD:
        this->step_sequence_(this->desc_->init_sequence, this->desc_->init_sequence_len, State::INIT_CMD,
                              State::SET_WINDOW);
        return;

      case State::SET_WINDOW:
        // window_sequence[this->plane_]: index 0 on the first entry
        // (plane_ reset to 0 in request_full_update()); when
        // window_per_plane is true, step_write_plane_() re-enters this
        // state with plane_ already advanced to 1 before the next plane's
        // byte burst, picking window_sequence[1] instead - see
        // panel_579c_gdey0579f51.h for why a panel needs that.
        this->step_sequence_(this->desc_->window_sequence[this->plane_], this->desc_->window_sequence_len[this->plane_],
                              State::SET_WINDOW, State::WRITE_PLANE);
        return;

      case State::WRITE_PLANE:
        this->step_write_plane_();
        return;

      case State::REFRESH_CMD:
        this->step_sequence_(this->desc_->refresh_sequence, this->desc_->refresh_sequence_len, State::REFRESH_CMD,
                              State::DONE);
        return;

      case State::POWER_OFF_CMD:
        // request_hibernate() routes through here first (hibernate() is
        // itself powerOff() + one more command on every GxEPD2 driver - see
        // epd_command.h) - done_state branches to HIBERNATE_CMD only then.
        this->step_sequence_(this->desc_->power_off_sequence, this->desc_->power_off_sequence_len,
                              State::POWER_OFF_CMD, this->hibernating_ ? State::HIBERNATE_CMD : State::DONE);
        return;

      case State::HIBERNATE_CMD:
        this->step_sequence_(this->desc_->hibernate_sequence, this->desc_->hibernate_sequence_len,
                              State::HIBERNATE_CMD, State::DONE);
        return;

      case State::POST_DELAY:
        if (millis() < this->deadline_)
          return;
        this->set_state_(this->delay_return_state_);
        return;

      case State::WAIT_BUSY:
        this->step_wait_busy_();
        return;

      case State::DONE:
        ESP_LOGD("gxepd2", "Display update finished");
        this->set_state_(State::IDLE);
        return;

      case State::ABORTED:
        // A busy timeout means the panel isn't behaving as wired/expected
        // (unplugged, dead, wrong busy polarity, ...) - pressing ahead
        // with the rest of the sequence (more command bytes, another
        // busy-wait) is very unlikely to "just work" and only obscures
        // which step actually failed. Bail the whole cycle here instead of
        // continuing to whatever state the timed-out wait was headed to.
        ESP_LOGE("gxepd2", "Aborting display update: busy timeout");
        this->cursor_ = 0;
        this->plane_ = 0;
        this->cmd_phase_ = CmdPhase::BEFORE_DELAY;
        this->set_state_(State::IDLE);
        return;
    }
  }

 private:
  // State/CmdPhase must come before arm_full_update_() below: a member
  // function's PARAMETER types (unlike its body) are resolved at the point
  // of declaration, not deferred to complete-class context, so `State` has
  // to already be a known type by the time arm_full_update_()'s signature is
  // parsed.
  enum class State {
    IDLE,
    RESET_LOW,
    RESET_LOW_WAIT,
    RESET_HIGH_WAIT,
    RESET_SETTLE,
    INIT_CMD,
    SET_WINDOW,
    WRITE_PLANE,
    REFRESH_CMD,
    POWER_OFF_CMD,
    HIBERNATE_CMD,
    POST_DELAY,
    WAIT_BUSY,
    DONE,
    ABORTED,
  };
  enum class CmdPhase { BEFORE_DELAY, ISSUE };

  // Shared by request_full_update() and (dev builds only) the no-reset
  // diagnostic path above - identical setup, different starting state.
  void arm_full_update_(const uint8_t *plane0, const uint8_t *plane1, State start_state) {
    this->planes_[0] = plane0;
    this->planes_[1] = plane1;
    uint32_t w = this->desc_->window_per_plane ? this->desc_->plane_width : this->desc_->width;
    // width*height*bpp/8: /8 at 1bpp, /4 at 2bpp, /2 at 4bpp.
    this->plane_bytes_ = (uint32_t) w * this->desc_->height * this->desc_->bits_per_pixel / 8;
    this->bytes_per_row_ = this->plane_bytes_ / this->desc_->height;
    // Wire length can exceed the buffer: PAIRED_EDGES sends each row twice
    // (2x-height RAM). plane_bytes_ stays the real buffer size; wire_bytes_
    // is what actually goes out on 0x10/0x24/... per plane.
    uint32_t wire_lines = this->desc_->height;
    if (this->desc_->line_order == LineOrder::PAIRED_EDGES)
      wire_lines *= 2;
    this->wire_bytes_ = wire_lines * this->bytes_per_row_;
    this->cursor_ = 0;
    this->plane_ = 0;
    this->cmd_phase_ = CmdPhase::BEFORE_DELAY;
    // Normal reset path: post-reset busy wait (see RESET_SETTLE below), then
    // the real init sequence - this->reset_done_next_/reset_then_wait_busy_
    // are only ever overridden by begin_power_down_()'s abort path.
    this->reset_done_next_ = State::INIT_CMD;
    this->reset_then_wait_busy_ = true;
    this->set_state_(start_state);
  }

  // Shared by request_power_off()/request_hibernate(). Most controllers
  // ignore new commands outright while BUSY is asserted, so writing
  // power_off_sequence/hibernate_sequence straight over a cycle that's still
  // mid-refresh would silently do nothing - the panel never actually powers
  // down, it just keeps running the update it was already doing. If a reset
  // pin is wired, force a hardware reset pulse first (the same one
  // request_full_update() issues every normal cycle) to abort whatever's in
  // flight and bring the controller back to a known, command-accepting
  // state before sending the power-down sequence - skipping the post-reset
  // busy wait that a real update takes, since this is an emergency path
  // (e.g. on_powerdown()'s tight teardown budget) that can't afford to
  // stall on busy_timeout_ms. Without a reset pin there's no way to
  // interrupt a busy controller, so this is best-effort in that case:
  // jamming the state machine straight to POWER_OFF_CMD, same as calling
  // this while idle.
  void begin_power_down_(bool hibernating) {
    this->cursor_ = 0;
    this->plane_ = 0;
    this->cmd_phase_ = CmdPhase::BEFORE_DELAY;
    this->hibernating_ = hibernating;
    if (this->state_ != State::IDLE && this->rst_ != nullptr) {
      this->reset_done_next_ = State::POWER_OFF_CMD;
      this->reset_then_wait_busy_ = false;
      this->set_state_(State::RESET_LOW);
    } else {
      this->set_state_(State::POWER_OFF_CMD);
    }
  }

  // Every command byte and every data byte is its own SPI transaction, CS
  // toggled around each one - this looks wasteful, but it's mechanically
  // what GxEPD2_EPD::_writeCommand()/_writeData(uint8_t) do (see
  // third_party/GxEPD2/src/GxEPD2_EPD.cpp), and matching the vendor cycle byte-for-byte
  // is the point of this engine.
  void write_byte_(uint8_t b, bool dc_high) {
    this->dc_->digital_write(dc_high);
    this->spi_->enable();
    this->cs_->digital_write(false);
    this->spi_->write_byte(b);
    this->cs_->digital_write(true);
    this->spi_->disable();
  }

  void write_command_(uint8_t cmd, const uint8_t *data, uint8_t data_len) {
#if ESPHOME_LOG_LEVEL >= ESPHOME_LOG_LEVEL_DEBUG
    // Command-level skeleton trace - deliberately left in (not gated to
    // VERY_VERBOSE) while diagnosing a real-hardware content bug: this is
    // the exact byte-for-byte counterpart to compare against a captured
    // epaper_spi trace on the same panel. Bulk plane data isn't logged
    // here (see step_write_plane_) - only the command skeleton matters for
    // that comparison.
    char buf[32] = {0};
    int n = 0;
    for (uint8_t i = 0; i < data_len && n < (int) sizeof(buf) - 3; i++)
      n += snprintf(buf + n, sizeof(buf) - n, "%02X ", data[i]);
    ESP_LOGD("gxepd2", "cmd 0x%02X data[%u]: %s", cmd, data_len, buf);
#endif
    this->write_byte_(cmd, false);
    for (uint8_t i = 0; i < data_len; i++)
      this->write_byte_(data[i], true);
  }

  // Steps through one EpdCommand table one entry (and one delay/busy-wait,
  // if that entry has one) at a time, resuming in self_state after any
  // wait so the same call re-enters here for the next entry. Shared by
  // INIT_CMD and REFRESH_CMD - which table is being stepped is entirely
  // determined by which state called this.
  void step_sequence_(const EpdCommand *seq, uint8_t len, State self_state, State done_state) {
    if (this->cursor_ >= len) {
      this->cursor_ = 0;
      this->cmd_phase_ = CmdPhase::BEFORE_DELAY;
      this->set_state_(done_state);
      return;
    }
    const EpdCommand &c = seq[this->cursor_];
    if (this->cmd_phase_ == CmdPhase::BEFORE_DELAY) {
      this->cmd_phase_ = CmdPhase::ISSUE;
      if (c.delay_before_ms > 0) {
        this->start_delay_(c.delay_before_ms, self_state);
        return;
      }
    }
    this->write_command_(c.cmd, c.data, c.data_len);
    this->cursor_++;
    this->cmd_phase_ = CmdPhase::BEFORE_DELAY;
    if (c.wait_busy_after) {
      this->start_busy_wait_(self_state);
    } else if (c.delay_after_ms > 0) {
      this->start_delay_(c.delay_after_ms, self_state);
    }
    // else: stays in self_state; the next step() call processes the next entry.
  }

  // Writes up to BYTE_BUDGET bytes of the current plane per call, then
  // returns - this bound is what keeps a full-frame write (thousands of
  // bytes) from ever making one loop() tick run long, replacing the old
  // per-256-byte esp_task_wdt_reset()+delay(1) throttle with simply not
  // doing enough work in one call to need it.
  //
  // 64 was overly conservative: at 20MHz SPI that's ~25us of actual bus
  // time, so a 426T82 full-plane write (48000 bytes) was bottlenecked
  // almost entirely on loop() scheduling overhead (~17ms/step observed),
  // not SPI - ~12.7s for one plane. Raised to 16384 (up to ~25ms of
  // blocking per step is fine here), cutting that to 3 steps instead of
  // 750 - measured ~1.07s for the plane write on hardware. The remaining
  // time is dominated by write_byte_()'s per-byte SPI enable()/disable()
  // (ESP-IDF transaction setup/teardown, ~21.7us/byte vs ~0.4us raw
  // transfer), which is intentionally NOT batched away: it lets other
  // devices on a shared spi: bus interleave between our bytes instead of
  // being locked out for the whole burst. 1s is well under half this
  // panel's own multi-second refresh time anyway, so it wasn't worth
  // trading that fairness for more throughput.
  static constexpr uint32_t BYTE_BUDGET = 16384;

  void step_write_plane_() {
    if (this->cursor_ == 0) {
      const uint8_t *buf = this->planes_[this->plane_];
      uint32_t histo[3] = {0, 0, 0};  // [0]=0x00, [1]=0xFF, [2]=other
      if (buf != nullptr) {
        for (uint32_t i = 0; i < this->plane_bytes_; i++) {
          if (buf[i] == 0x00)
            histo[0]++;
          else if (buf[i] == 0xFF)
            histo[1]++;
          else
            histo[2]++;
        }
      }
      ESP_LOGD("gxepd2", "plane %u: cmd 0x%02X, %u wire bytes (buf %u), invert=%d, buf[0]=0x%02X buf[last]=0x%02X, histo(00/FF/other)=%u/%u/%u",
               this->plane_, this->desc_->plane_write_cmd[this->plane_], (unsigned) this->wire_bytes_,
               (unsigned) this->plane_bytes_, this->desc_->plane_invert[this->plane_], buf != nullptr ? buf[0] : 0,
               buf != nullptr ? buf[this->plane_bytes_ - 1] : 0, (unsigned) histo[0], (unsigned) histo[1],
               (unsigned) histo[2]);
      this->write_byte_(this->desc_->plane_write_cmd[this->plane_], false);
    }
    const uint8_t *buf = this->planes_[this->plane_];
    bool invert = this->desc_->plane_invert[this->plane_];
    uint32_t end = this->cursor_ + BYTE_BUDGET;
    if (end > this->wire_bytes_)
      end = this->wire_bytes_;
    for (; this->cursor_ < end; this->cursor_++) {
      uint32_t src = this->line_src_index_(this->cursor_);
      uint8_t b = buf != nullptr ? buf[src] : 0xFF;
      this->write_byte_(invert ? static_cast<uint8_t>(~b) : b, true);
    }
    if (this->cursor_ >= this->wire_bytes_) {
      this->cursor_ = 0;
      this->plane_++;
      this->cmd_phase_ = CmdPhase::BEFORE_DELAY;
      if (this->plane_ >= this->desc_->plane_count) {
        this->plane_ = 0;
        this->set_state_(State::REFRESH_CMD);
      } else if (this->desc_->window_per_plane) {
        // Another plane to go and it needs its own window-set first (a
        // genuinely separate physical RAM region, not just another color
        // plane sharing the window already set) - re-enter SET_WINDOW,
        // which will pick window_sequence[this->plane_] now that plane_
        // has already advanced.
        this->set_state_(State::SET_WINDOW);
      }
      // else: more planes, sharing one window already set - stay in
      // WRITE_PLANE, next step() call starts the next plane's byte burst.
    }
  }

  // Maps a wire-transmission-order byte position (0..wire_bytes_-1) to the
  // buffer byte it reads. The wire is a sequence of lines, each bytes_per_row_
  // bytes from one buffer row; line_order picks which buffer row each wire
  // line carries (and, for PAIRED_EDGES, there are 2x as many wire lines as
  // buffer rows). SEQUENTIAL is the identity (wire line i -> row i).
  uint32_t line_src_index_(uint32_t wire_pos) const {
    uint32_t wire_line = wire_pos / this->bytes_per_row_;
    uint32_t byte_in_line = wire_pos % this->bytes_per_row_;
    uint32_t buf_row;
    switch (this->desc_->line_order) {
      case LineOrder::PAIRED_EDGES: {
        // wire lines come in pairs: for i in [0,H), line 2i = row (H-1-i),
        // line 2i+1 = row i. Matches GxEPD2 writeNative's emission exactly.
        uint32_t i = wire_line / 2;
        buf_row = (wire_line % 2 == 0) ? (this->desc_->height - 1 - i) : i;
        break;
      }
      case LineOrder::SEQUENTIAL:
      default:
        buf_row = wire_line;
        break;
    }
    return buf_row * this->bytes_per_row_ + byte_in_line;
  }

  void start_delay_(uint16_t delay_ms, State return_to) {
    this->deadline_ = millis() + delay_ms;
    this->delay_return_state_ = return_to;
    this->set_state_(State::POST_DELAY);
  }

  void start_busy_wait_(State next) {
    this->busy_wait_next_ = next;
    this->busy_wait_deadline_ = millis() + this->desc_->busy_timeout_ms;
    this->set_state_(State::WAIT_BUSY);
  }

  void step_wait_busy_() {
    if (this->busy_ == nullptr) {
      this->set_state_(this->busy_wait_next_);
      return;
    }
    bool busy = this->busy_->digital_read() == this->desc_->busy_active_level;
    if (!busy) {
      this->set_state_(this->busy_wait_next_);
      return;
    }
    if (millis() >= this->busy_wait_deadline_) {
      ESP_LOGW("gxepd2", "Busy timeout waiting for panel");
      this->set_state_(State::ABORTED);
    }
  }

  // Logs every state transition (only when the state actually changes -
  // WAIT_BUSY and POST_DELAY re-enter their own case on every step() call
  // while polling, which would otherwise spam this at loop() rate) so a
  // first hardware run can be checked against the vendor cycle it's meant
  // to replicate before this gets dialed back to WARN-only.
  void set_state_(State s) {
    if (s != this->state_)
      ESP_LOGD("gxepd2", "state: %s -> %s", state_name_(this->state_), state_name_(s));
    this->state_ = s;
  }

  static const char *state_name_(State s) {
    switch (s) {
      case State::IDLE:
        return "IDLE";
      case State::RESET_LOW:
        return "RESET_LOW";
      case State::RESET_LOW_WAIT:
        return "RESET_LOW_WAIT";
      case State::RESET_HIGH_WAIT:
        return "RESET_HIGH_WAIT";
      case State::RESET_SETTLE:
        return "RESET_SETTLE";
      case State::INIT_CMD:
        return "INIT_CMD";
      case State::SET_WINDOW:
        return "SET_WINDOW";
      case State::WRITE_PLANE:
        return "WRITE_PLANE";
      case State::REFRESH_CMD:
        return "REFRESH_CMD";
      case State::POWER_OFF_CMD:
        return "POWER_OFF_CMD";
      case State::HIBERNATE_CMD:
        return "HIBERNATE_CMD";
      case State::POST_DELAY:
        return "POST_DELAY";
      case State::WAIT_BUSY:
        return "WAIT_BUSY";
      case State::DONE:
        return "DONE";
      case State::ABORTED:
        return "ABORTED";
    }
    return "?";
  }

  const EpdPanelDescriptor *desc_{nullptr};
  GPIOPin *cs_{nullptr};
  GPIOPin *dc_{nullptr};
  GPIOPin *rst_{nullptr};
  GPIOPin *busy_{nullptr};
  EpdSPIDevice *spi_{nullptr};

  State state_{State::IDLE};
  uint32_t cursor_{0};
  CmdPhase cmd_phase_{CmdPhase::BEFORE_DELAY};
  uint8_t plane_{0};
  const uint8_t *planes_[2]{nullptr, nullptr};
  uint32_t plane_bytes_{0};   // real buffer size per plane
  uint32_t wire_bytes_{0};    // bytes actually sent per plane (>= plane_bytes_)
  uint32_t bytes_per_row_{0};
  uint32_t deadline_{0};
  State reset_done_next_{State::INIT_CMD};  // RESET_SETTLE's continuation - see arm_full_update_()/begin_power_down_()
  bool reset_then_wait_busy_{true};         // false only for begin_power_down_()'s abort path
  State delay_return_state_{State::DONE};
  uint32_t busy_wait_deadline_{0};
  State busy_wait_next_{State::DONE};
  bool hibernating_{false};  // POWER_OFF_CMD's done_state: HIBERNATE_CMD if set, else DONE
};

}  // namespace esphome::gxepd2