#include "Receiver.h"
#include "Log.h"
#include <math.h>
#if !OPTICAL_USE_MCPWM_CAPTURE
#include <esp_cpu.h>
#include <esp32-hal-cpu.h>
#endif

uint32_t PulseReceiver::tickHz() const {
#if OPTICAL_USE_MCPWM_CAPTURE
  return captureResolutionHz_;
#else
  return cpuTickHz_;
#endif
}

uint32_t PulseReceiver::plannedTickHz(uint32_t expectedHz, float expectedDutyPct) const {
  if (!expectedHz || expectedDutyPct <= 0.0f || expectedDutyPct >= 100.0f) return 0;
#if OPTICAL_USE_MCPWM_CAPTURE
  return captureResolutionHz_ ? captureResolutionHz_ :
    MCPWM_CAPTURE_RESOLUTION_HZ;
#else
  return cpuTickHz_;
#endif
}

bool PulseReceiver::begin() {
  queue_ = xQueueCreate(512, sizeof(Edge));
  if (!queue_) return false;
#if OPTICAL_USE_MCPWM_CAPTURE
  // PWM generation uses MCPWM group 0. Group 1 is dedicated to input capture,
  // so RX cannot exhaust or conflict with the generator's resources.
  mcpwm_capture_timer_config_t timerConfig = {};
  timerConfig.group_id = 1;
  timerConfig.clk_src = MCPWM_CAPTURE_CLK_SRC_DEFAULT;
  timerConfig.resolution_hz = MCPWM_CAPTURE_RESOLUTION_HZ;
  if (mcpwm_new_capture_timer(&timerConfig, &captureTimer_) != ESP_OK) return false;
  if (mcpwm_capture_timer_get_resolution(captureTimer_, &captureResolutionHz_) != ESP_OK ||
      !captureResolutionHz_) return false;

  mcpwm_capture_channel_config_t channelConfig = {};
  channelConfig.gpio_num = GPIO_RX;
  // ACK pulses are sub-microsecond and consecutive acknowledgements can be
  // only 1 us apart. A low-priority capture interrupt can leave the channel
  // status pending long enough for the next timestamp to overwrite it.
  channelConfig.intr_priority = 3;
  channelConfig.prescale = 1;
  channelConfig.flags.pos_edge = true;
  channelConfig.flags.neg_edge = false;
  if (mcpwm_new_capture_channel(captureTimer_, &channelConfig, &risingChannel_) != ESP_OK)
    return false;
  mcpwm_capture_event_callbacks_t callbacks = {};
  callbacks.on_cap = onCapture;
  if (mcpwm_capture_channel_register_event_callbacks(
      risingChannel_, &callbacks, this) != ESP_OK) return false;

  channelConfig.flags.pos_edge = false;
  channelConfig.flags.neg_edge = true;
  if (mcpwm_new_capture_channel(captureTimer_, &channelConfig, &fallingChannel_) != ESP_OK)
    return false;
  if (mcpwm_capture_channel_register_event_callbacks(
      fallingChannel_, &callbacks, this) != ESP_OK) return false;

  // The TX channel is created immediately before a driver test. Only the end
  // of the active PWM pulse is armed; handling its start here would occupy the
  // shared MCPWM ISR during the RX acknowledgement only ~300 ns later.
  return true;
#else
  pinMode(GPIO_RX, INPUT);
  cpuTickHz_ = getCpuFrequencyMhz() * 1000000UL;
  attachInterruptArg(GPIO_RX, onGpio, this, CHANGE);
  return cpuTickHz_ != 0;
#endif
}

bool PulseReceiver::start(uint32_t expectedHz, float expectedDutyPct,
                          bool activeLightOn) {
  (void)activeLightOn;
  if (!plannedTickHz(expectedHz, expectedDutyPct)) return false;
  expectedHz_ = expectedHz;
  expectedDutyPct_ = expectedDutyPct;
#if !OPTICAL_USE_MCPWM_CAPTURE
  cpuTickHz_ = getCpuFrequencyMhz() * 1000000UL;
  if (!cpuTickHz_) return false;
#endif
  resetStream();
  Log::event("CAPTURE", "RX optical polarity will be detected automatically");
  return startCapture(false);
}

#if OPTICAL_USE_MCPWM_CAPTURE
bool PulseReceiver::configureDriverTxCapture(bool risingEdge) {
  if (running_) return false;
  if (txChannel_) {
    if (mcpwm_del_capture_channel(txChannel_) != ESP_OK) return false;
    txChannel_ = nullptr;
  }
  mcpwm_capture_channel_config_t config = {};
  config.gpio_num = GPIO_PWM;
  config.intr_priority = 3;
  config.prescale = 1;
  config.flags.pos_edge = risingEdge;
  config.flags.neg_edge = !risingEdge;
  config.flags.io_loop_back = true;
  if (mcpwm_new_capture_channel(captureTimer_, &config, &txChannel_) != ESP_OK)
    return false;
  mcpwm_capture_event_callbacks_t callbacks = {};
  callbacks.on_cap = onCapture;
  return mcpwm_capture_channel_register_event_callbacks(
    txChannel_, &callbacks, this) == ESP_OK;
}
#endif

bool PulseReceiver::startDriver(uint32_t frequencyHz, uint32_t pulseNs,
                                bool activeTxLightOn) {
#if OPTICAL_USE_MCPWM_CAPTURE
  if (!frequencyHz || !pulseNs || !tickHz() || running_) return false;
  resetStream();
  const uint64_t pulseTicks =
    (static_cast<uint64_t>(pulseNs) * tickHz() + 500000000ULL) /
    1000000000ULL;
  const uint32_t periodTicks = tickHz() / frequencyHz;
  if (!pulseTicks || pulseTicks >= periodTicks || pulseTicks > UINT32_MAX)
    return false;
  driverPulseTicks_ = static_cast<uint32_t>(pulseTicks);
  driverReleaseSlackTicks_ = static_cast<uint32_t>(
    (static_cast<uint64_t>(DRIVER_RESPONSE_TIMEOUT_NS) * tickHz() +
     999999999ULL) / 1000000000ULL);
  const uint8_t activeRawLevel = activeTxLightOn ?
    TX_LIGHT_ON_GPIO_LEVEL : TX_LIGHT_OFF_GPIO_LEVEL;
  const bool pulseEndIsRising = activeRawLevel == LOW;
  if (!driverReleaseSlackTicks_ ||
      !configureDriverTxCapture(pulseEndIsRising)) return false;
  return startCapture(true);
#else
  return false;
#endif
}

bool PulseReceiver::startCapture(bool withTx) {
#if OPTICAL_USE_MCPWM_CAPTURE
  // Progress updates keep one capture session alive. Pulse-width stages stop
  // capture only after PWM is quiet, so resetStream never races the ISR.
  if (running_) return true;
  if (mcpwm_capture_timer_enable(captureTimer_) != ESP_OK) return false;
  if (mcpwm_capture_channel_enable(risingChannel_) != ESP_OK) {
    mcpwm_capture_timer_disable(captureTimer_);
    return false;
  }
  if (mcpwm_capture_channel_enable(fallingChannel_) != ESP_OK) {
    mcpwm_capture_channel_disable(risingChannel_);
    mcpwm_capture_timer_disable(captureTimer_);
    return false;
  }
  if (withTx && mcpwm_capture_channel_enable(txChannel_) != ESP_OK) {
    mcpwm_capture_channel_disable(fallingChannel_);
    mcpwm_capture_channel_disable(risingChannel_);
    mcpwm_capture_timer_disable(captureTimer_);
    return false;
  }
  txCaptureEnabled_ = withTx;
  running_ = true;
  if (mcpwm_capture_timer_start(captureTimer_) != ESP_OK) {
    running_ = false;
    if (txCaptureEnabled_) mcpwm_capture_channel_disable(txChannel_);
    txCaptureEnabled_ = false;
    mcpwm_capture_channel_disable(fallingChannel_);
    mcpwm_capture_channel_disable(risingChannel_);
    mcpwm_capture_timer_disable(captureTimer_);
    return false;
  }
#else
  if (withTx) return false;
  running_ = true;
#endif
  return true;
}

void PulseReceiver::stop() {
  const bool wasRunning = running_;
  running_ = false;
#if OPTICAL_USE_MCPWM_CAPTURE
  if (wasRunning) {
    // Mask capture interrupts before stopping the shared capture timer.
    if (txCaptureEnabled_) mcpwm_capture_channel_disable(txChannel_);
    txCaptureEnabled_ = false;
    mcpwm_capture_channel_disable(fallingChannel_);
    mcpwm_capture_channel_disable(risingChannel_);
    mcpwm_capture_timer_stop(captureTimer_);
    mcpwm_capture_timer_disable(captureTimer_);
  }
#else
  (void)wasRunning;
#endif
}

void PulseReceiver::resetStream() {
  if (queue_) xQueueReset(queue_);
  haveReorderEdge_ = false;
  __atomic_store_n(&driverRingWrite_, 0U, __ATOMIC_RELEASE);
  __atomic_store_n(&driverRingRead_, 0U, __ATOMIC_RELEASE);
  haveLastDriverEdge_ = false;
  lastDriverEdge_ = {};
  driverPulseTicks_ = 0;
  driverReleaseSlackTicks_ = 0;
  droppedItems_ = 0;
  polarityKnown_ = false;
  activeStartRising_ = false;
  polarityEdgeCount_ = 0;
  memset(polarityEdges_, 0, sizeof(polarityEdges_));
  waitingForActiveEnd_ = true;
  activeStart_ = activeEnd_ = 0;
  haveRawTick_ = false;
  lastRawTick_ = 0;
  tickEpoch_ = 0;
}

PulseReceiver::TimedEdge PulseReceiver::extendEdge(const Edge &e) {
  if (haveRawTick_ && e.tick < lastRawTick_ && lastRawTick_ - e.tick > 0x80000000UL)
    tickEpoch_ += 0x100000000ULL;
  haveRawTick_ = true; lastRawTick_ = e.tick;
  return {tickEpoch_ + e.tick, e.rising != 0};
}

bool PulseReceiver::consumeEdge(const Edge &rawEdge, PulsePeriod &out) {
  const TimedEdge edge = extendEdge(rawEdge);
  if (polarityKnown_) {
    // Deliberately ignore edge type after synchronization. A PWM waveform is
    // just alternating intervals: active, inactive, active, inactive. An
    // extra or missing edge therefore becomes a concrete wrong pulse/period
    // instead of an ambiguous GLITCH state.
    if (waitingForActiveEnd_) {
      activeEnd_ = edge.tick;
      waitingForActiveEnd_ = false;
      return false;
    }
    const uint64_t periodTicks = edge.tick - activeStart_;
    const uint64_t activeTicks = activeEnd_ - activeStart_;
    out = {activeStart_, static_cast<uint32_t>(periodTicks),
      static_cast<uint32_t>(activeTicks), tickHz()};
    activeStart_ = edge.tick;
    waitingForActiveEnd_ = true;
    return true;
  }

  // Optical mode does not use the DRIVER level setting. Compare the first two
  // alternating intervals with the configured active duration and select the
  // level that is actually present at RX. Three edges are enough to determine
  // polarity and, when the first interval is active, produce the first period.
  polarityEdges_[polarityEdgeCount_++] = edge;
  if (polarityEdgeCount_ < 3U) return false;

  const uint64_t firstInterval =
    polarityEdges_[1].tick - polarityEdges_[0].tick;
  const uint64_t secondInterval =
    polarityEdges_[2].tick - polarityEdges_[1].tick;
  const uint64_t expectedPeriod = tickHz() / expectedHz_;
  const uint64_t expectedActive = static_cast<uint64_t>(
    expectedPeriod * expectedDutyPct_ / 100.0f + 0.5f);
  const uint64_t firstError = firstInterval > expectedActive ?
    firstInterval - expectedActive : expectedActive - firstInterval;
  const uint64_t secondError = secondInterval > expectedActive ?
    secondInterval - expectedActive : expectedActive - secondInterval;
  const bool firstIntervalIsActive = firstError <= secondError;
  activeStartRising_ = firstIntervalIsActive ? polarityEdges_[0].rising :
                                               polarityEdges_[1].rising;
  polarityKnown_ = true;
  polarityEdgeCount_ = 0;

  if (firstIntervalIsActive) {
    out = {polarityEdges_[0].tick,
      static_cast<uint32_t>(polarityEdges_[2].tick - polarityEdges_[0].tick),
      static_cast<uint32_t>(firstInterval), tickHz()};
    activeStart_ = polarityEdges_[2].tick;
    waitingForActiveEnd_ = true;
    return true;
  }

  activeStart_ = polarityEdges_[1].tick;
  activeEnd_ = polarityEdges_[2].tick;
  waitingForActiveEnd_ = false;
  return false;
}

uint32_t PulseReceiver::takeDroppedItems() {
  return __atomic_exchange_n(&droppedItems_, 0, __ATOMIC_RELAXED);
}

#if OPTICAL_USE_MCPWM_CAPTURE
bool IRAM_ATTR PulseReceiver::onCapture(mcpwm_cap_channel_handle_t channel,
                                        const mcpwm_capture_event_data_t *data,
                                        void *ctx) {
  PulseReceiver *self = static_cast<PulseReceiver *>(ctx);
  if (!self->running_) return false;
  const bool rawRising = data->cap_edge == MCPWM_CAP_EDGE_POS;
  const Edge edge = {data->cap_value, static_cast<uint8_t>(rawRising),
    static_cast<uint8_t>(channel == self->txChannel_ ? CaptureSource::TX :
      CaptureSource::RX)};
  if (self->txCaptureEnabled_) {
    // All channels in one MCPWM group are dispatched serially by the same
    // group ISR. Keep this callback shorter than the minimum interval between
    // equal RX edges (about 2.1 us at W=2 us): a spinlock and several atomic
    // RMW operations here can leave the channel pending until its capture
    // register is overwritten by the next edge.
    if (self->haveLastDriverEdge_ &&
        self->lastDriverEdge_.tick == edge.tick &&
        self->lastDriverEdge_.rising == edge.rising &&
        self->lastDriverEdge_.source == edge.source) {
      // The same channel callback can be delivered twice while several MCPWM
      // capture status bits are pending. Two physical edges cannot have the
      // same source, direction and 12.5 ns hardware timestamp.
      return false;
    }
    const uint16_t write = self->driverRingWrite_;
    const uint16_t next = static_cast<uint16_t>(
      (write + 1U) & (DRIVER_RING_CAPACITY - 1U));
    if (next == self->driverRingRead_) {
      ++self->droppedItems_;
    } else {
      self->driverRing_[write] = edge;
      self->lastDriverEdge_ = edge;
      self->haveLastDriverEdge_ = true;
      asm volatile("memw" ::: "memory");
      self->driverRingWrite_ = next;
    }
    return false;
  }
  BaseType_t wake = pdFALSE;
  if (xQueueSendFromISR(self->queue_, &edge, &wake) != pdTRUE)
    __atomic_fetch_add(&self->droppedItems_, 1U, __ATOMIC_RELAXED);
  return wake == pdTRUE;
}
#else
void IRAM_ATTR PulseReceiver::onGpio(void *ctx) {
  PulseReceiver *self = static_cast<PulseReceiver *>(ctx);
  if (!self->running_) return;
  const bool level = gpio_get_level(static_cast<gpio_num_t>(GPIO_RX));
  const Edge edge = {esp_cpu_get_cycle_count(), static_cast<uint8_t>(level),
    static_cast<uint8_t>(CaptureSource::RX)};
  BaseType_t wake = pdFALSE;
  if (xQueueSendFromISR(self->queue_, &edge, &wake) != pdTRUE)
    __atomic_fetch_add(&self->droppedItems_, 1U, __ATOMIC_RELAXED);
  if (wake) portYIELD_FROM_ISR();
}
#endif

bool PulseReceiver::nextOrderedEdge(Edge &edge, TickType_t waitTicks) {
  if (!haveReorderEdge_) {
    if (xQueueReceive(queue_, &reorderEdge_, waitTicks) != pdTRUE) return false;
    haveReorderEdge_ = true;
  }
  Edge next = {};
  // Keep one-event look-ahead. If both channel interrupts were pending while
  // OLED/I2C ran, the MCPWM driver may dispatch them by channel number rather
  // than timestamp. The signed modular comparison restores their real order.
  if (xQueueReceive(queue_, &next, waitTicks) != pdTRUE) return false;
  if (static_cast<int32_t>(next.tick - reorderEdge_.tick) < 0) {
    edge = next;
  } else {
    edge = reorderEdge_;
    reorderEdge_ = next;
  }
  return true;
}

size_t PulseReceiver::readPeriods(PulsePeriod *periods, size_t capacity,
                                  TickType_t waitTicks) {
  size_t count = 0;
  Edge edge = {};
  while (count < capacity && nextOrderedEdge(edge, count ? 0 : waitTicks)) {
    if (consumeEdge(edge, periods[count])) {
      periods[count].activeTickHz = tickHz();
      ++count;
    }
  }
  return count;
}

size_t PulseReceiver::readEvents(CaptureEvent *events, size_t capacity,
                                 TickType_t waitTicks) {
  if (!events || capacity < 3U || !txCaptureEnabled_ || !driverPulseTicks_ ||
      !driverReleaseSlackTicks_) return 0;

  constexpr size_t MAX_BATCH = 64;
  const size_t limit = capacity < MAX_BATCH ? capacity : MAX_BATCH;
  Edge ordered[MAX_BATCH] = {};

  uint16_t read = __atomic_load_n(&driverRingRead_, __ATOMIC_RELAXED);
  if (read == __atomic_load_n(&driverRingWrite_, __ATOMIC_ACQUIRE) && waitTicks) {
    vTaskDelay(waitTicks);
    read = __atomic_load_n(&driverRingRead_, __ATOMIC_RELAXED);
  }
  // Work on one immutable producer snapshot. RX belonging to a pulse start is
  // deliberately retained until that pulse's captured end arrives: only then
  // can the missing start interrupt be reconstructed and sorted before RX.
  const uint16_t write = __atomic_load_n(&driverRingWrite_, __ATOMIC_ACQUIRE);
  uint16_t scan = read;
  bool haveLatestTxEnd = false;
  bool haveReleaseTxEnd = false;
  uint32_t latestTxEnd = 0;
  uint32_t releaseTxEnd = 0;
  size_t projectedCount = 0;
  while (scan != write) {
    const Edge &edge = driverRing_[scan];
    const size_t needed = edge.source == static_cast<uint8_t>(CaptureSource::TX)
      ? 2U : 1U;
    // Reserve one output slot for WINDOW_END.
    if (projectedCount + needed + 1U > limit) break;
    projectedCount += needed;
    if (edge.source == static_cast<uint8_t>(CaptureSource::TX)) {
      if (haveLatestTxEnd) {
        releaseTxEnd = latestTxEnd;
        haveReleaseTxEnd = true;
      }
      latestTxEnd = edge.tick;
      haveLatestTxEnd = true;
    }
    scan = static_cast<uint16_t>((scan + 1U) % DRIVER_RING_CAPACITY);
  }
  // Keep the newest TX period in the ring. Arrival of the following TX end
  // proves that the previous end's ACK/fault window has completely elapsed.
  if (!haveReleaseTxEnd) {
    if (waitTicks) vTaskDelay(waitTicks);
    return 0;
  }

  const uint32_t releaseThrough = releaseTxEnd + driverReleaseSlackTicks_;
  size_t count = 0;
  while (read != write) {
    const Edge edge = driverRing_[read];
    if (edge.source == static_cast<uint8_t>(CaptureSource::TX) &&
        static_cast<int32_t>(edge.tick - releaseTxEnd) > 0)
      break;
    if (edge.source == static_cast<uint8_t>(CaptureSource::RX) &&
        static_cast<int32_t>(edge.tick - releaseThrough) > 0)
      break;
    const size_t needed = edge.source == static_cast<uint8_t>(CaptureSource::TX)
      ? 2U : 1U;
    if (count + needed > limit) break;
    read = static_cast<uint16_t>((read + 1U) % DRIVER_RING_CAPACITY);
    if (edge.source == static_cast<uint8_t>(CaptureSource::TX)) {
      Edge pulseStart = edge;
      pulseStart.tick -= driverPulseTicks_;
      pulseStart.rising = !edge.rising;
      ordered[count++] = pulseStart;
    }
    ordered[count++] = edge;
  }
  __atomic_store_n(&driverRingRead_, read, __ATOMIC_RELEASE);

  // MCPWM channels share one timer but their callbacks can be dispatched in
  // channel order when several interrupts are pending. Restore the hardware
  // order inside the captured batch using the common timestamp.
  for (size_t i = 1; i < count; ++i) {
    const Edge key = ordered[i];
    size_t j = i;
    while (j && static_cast<int32_t>(ordered[j - 1].tick - key.tick) > 0) {
      ordered[j] = ordered[j - 1];
      --j;
    }
    ordered[j] = key;
  }

  for (size_t i = 0; i < count; ++i) {
    const TimedEdge timed = extendEdge(ordered[i]);
    events[i] = {timed.tick, timed.rising,
      static_cast<CaptureSource>(ordered[i].source)};
  }
  const Edge marker = {releaseThrough, 0,
    static_cast<uint8_t>(CaptureSource::WINDOW_END)};
  const TimedEdge timedMarker = extendEdge(marker);
  events[count++] = {timedMarker.tick, false, CaptureSource::WINDOW_END};
  return count;
}