#include "App.h"
#include "Config.h"
#include <WiFi.h>
#include <esp_mac.h>
#include <esp_system.h>
#include <string.h>

App::App() : startButton_(GPIO_BUTTON_START), modeButton_(GPIO_BUTTON_MODE), measurement_(receiver_) {}

void App::begin() {
  Serial.begin(SERIAL_BAUD);
  startButton_.begin(); modeButton_.begin(); pwm_.begin();
  bootCheckStartedMs_ = millis();
  bootResetCandidate_ = startButton_.pressed() && modeButton_.pressed();
  if (!bootResetCandidate_) finishInitialization(false);
}

void App::finishInitialization(bool factoryReset) {
  if (initialized_) return;
  if (factoryReset) {
    store_.defaults(settings_); store_.save(settings_); Serial.println("FACTORY DEFAULTS RESTORED");
  } else if (!store_.load(settings_)) {
    store_.save(settings_); Serial.println("NVS invalid/missing: defaults loaded");
  }
  params_ = store_.params(settings_);
  if (!display_.begin()) Serial.println("OLED unavailable; continuing through Serial");
  initialized_ = true;
  if (!receiver_.begin()) { Serial.println("FATAL: capture peripheral init failed"); finish(false, FailReason::UNSUPPORTED); return; }
  printConfiguration(); showIdle();
}

void App::update() {
  const uint32_t now = millis();
  const ButtonEvent startEvent = startButton_.update(now);
  const ButtonEvent modeEvent = modeButton_.update(now);
  if (!initialized_) {
    if (!startButton_.pressed() || !modeButton_.pressed()) finishInitialization(false);
    else if (now - bootCheckStartedMs_ >= FACTORY_RESET_HOLD_MS) finishInitialization(true);
    return;
  }
  if (state_ != AppState::IDLE && state_ != AppState::MENU && state_ != AppState::FINISHED &&
      startEvent == ButtonEvent::LONG) { abortTest(); return; }

  if (state_ == AppState::IDLE || state_ == AppState::FINISHED) {
    if (modeEvent == ButtonEvent::SHORT) {
      settings_.role = (settings_.role + 1U) % 3U; store_.save(settings_); params_ = store_.params(settings_); showIdle();
      Serial.printf("MODE: %s\n", roleName(static_cast<Role>(settings_.role)));
    } else if (modeEvent == ButtonEvent::LONG) {
      state_ = AppState::MENU; menuItem_ = 0; showMenu();
    } else if (startEvent == ButtonEvent::SHORT) startTest();
    return;
  }
  if (state_ == AppState::MENU) {
    if (modeEvent == ButtonEvent::SHORT) { menuItem_ = (menuItem_ + 1U) % 7U; showMenu(); }
    else if (modeEvent == ButtonEvent::LONG) {
      sanitizeRange(); store_.save(settings_); params_ = store_.params(settings_);
      state_ = AppState::IDLE; printConfiguration(); showIdle();
    } else if (startEvent == ButtonEvent::SHORT) changeMenu(+1);
    else if (startEvent == ButtonEvent::LONG || startEvent == ButtonEvent::REPEAT) changeMenu(-1);
    return;
  }
  if (state_ == AppState::SOLO_MEASURE) {
    const MeasureState ms = measurement_.update();
    if (ms == MeasureState::FAIL) { printStageStats(measurement_.stats(), requestedHz_); finish(false, measurement_.reason()); }
    else if (ms == MeasureState::PASS) { printStageStats(measurement_.stats(), requestedHz_); stagePassed(); }
  } else if (state_ == AppState::MASTER_DISCOVER || state_ == AppState::MASTER_WAIT_READY ||
             state_ == AppState::MASTER_WAIT_RESULT) {
    handleRadio(); updateMaster();
  } else {
    handleRadio(); updateSlave();
  }
}

void App::showIdle() {
  char one[24]; snprintf(one, sizeof(one), "MODE: %s", roleName(static_cast<Role>(settings_.role)));
  display_.show(one, "START=RUN");
}

void App::sanitizeRange() {
  settings_.startIndex %= countOf(START_FREQ_OPTIONS_HZ); settings_.endIndex %= countOf(END_FREQ_OPTIONS_HZ);
  if (END_FREQ_OPTIONS_HZ[settings_.endIndex] <= START_FREQ_OPTIONS_HZ[settings_.startIndex]) {
    size_t i = 0;
    while (i < countOf(END_FREQ_OPTIONS_HZ) && END_FREQ_OPTIONS_HZ[i] <= START_FREQ_OPTIONS_HZ[settings_.startIndex]) ++i;
    if (i == countOf(END_FREQ_OPTIONS_HZ)) { settings_.startIndex = 0; i = countOf(END_FREQ_OPTIONS_HZ) - 1; }
    settings_.endIndex = i;
  }
}

void App::changeMenu(int d) {
  uint8_t *value = nullptr; size_t count = 0;
  switch (menuItem_) {
    case 0: value = &settings_.startIndex; count = countOf(START_FREQ_OPTIONS_HZ); break;
    case 1: value = &settings_.endIndex; count = countOf(END_FREQ_OPTIONS_HZ); break;
    case 2: value = &settings_.stepIndex; count = countOf(STEP_OPTIONS_HZ); break;
    case 3: value = &settings_.accuracyIndex; count = countOf(ACCURACY_OPTIONS_PCT); break;
    case 4: value = &settings_.timeIndex; count = countOf(TEST_TIME_OPTIONS_MS); break;
    case 5: value = &settings_.repeatIndex; count = countOf(REPEAT_OPTIONS); break;
    default: value = &settings_.dutyIndex; count = countOf(DUTY_OPTIONS_PCT); break;
  }
  *value = static_cast<uint8_t>((*value + count + d) % count); sanitizeRange(); params_ = store_.params(settings_); showMenu();
}

void App::showMenu() {
  char one[22], two[22], all[12];
  Display::formatDuration(actualNominalTotalUs(), all, sizeof(all));
  switch (menuItem_) {
    case 0: snprintf(one, sizeof(one), "START FREQ"); Display::formatFrequency(params_.startHz, two, sizeof(two)); break;
    case 1: snprintf(one, sizeof(one), "END FREQ"); Display::formatFrequency(params_.endHz, two, sizeof(two)); break;
    case 2: snprintf(one, sizeof(one), "FREQ STEP"); Display::formatFrequency(params_.stepHz, two, sizeof(two)); break;
    case 3: snprintf(one, sizeof(one), "ACCURACY"); snprintf(two, sizeof(two), "+/-%g%%", params_.accuracyPct); break;
    case 4: snprintf(one, sizeof(one), "TEST TIME"); snprintf(two, sizeof(two), "%.1fs", params_.testTimeMs / 1000.0f); break;
    case 5: snprintf(one, sizeof(one), "REPEATS"); snprintf(two, sizeof(two), "%ux", params_.repeats); break;
    default: snprintf(one, sizeof(one), "PWM DUTY"); snprintf(two, sizeof(two), "%u%%", params_.dutyPct); break;
  }
  const size_t used = strlen(two); snprintf(two + used, sizeof(two) - used, " ALL %s", all); display_.show(one, two);
}

void App::startTest() {
  params_ = store_.params(settings_); stageCount_ = frequencyPointCount(params_.startHz, params_.endHz, params_.stepHz);
  stageIndex_ = 0; pendingReason_ = FailReason::NONE;
  if (!stageCount_) { finish(false, FailReason::UNSUPPORTED); return; }
  printConfiguration();
  const Role role = static_cast<Role>(settings_.role);
  if (role == Role::SOLO) {
    if (!prepareStage()) return;
    state_ = AppState::SOLO_MEASURE;
  } else if (!radio_.begin()) finish(false, FailReason::LINK_LOST);
  else if (role == Role::MASTER) startMasterDiscovery();
  else { state_ = AppState::SLAVE_READY; display_.show("SLAVE READY", "WAIT MASTER"); Serial.println("SLAVE READY"); }
}

bool App::prepareStage() {
  requestedHz_ = frequencyAt(params_.startHz, params_.endHz, params_.stepHz, stageIndex_);
  const uint32_t maxHz = TARGET_IS_C3 ? C3_STRICT_MAX_HZ :
    (receiver_.highRateBackend() ? S3_STRICT_MAX_HZ : C3_STRICT_MAX_HZ);
  if (requestedHz_ > maxHz) { finish(false, FailReason::UNSUPPORTED); return false; }
  if (!pwm_.start(requestedHz_, params_.dutyPct, actual_)) { finish(false, FailReason::RESOLUTION); return false; }
  const FailReason resolution = validateResolution(actual_.actualHz, actual_.actualDutyPct, params_.accuracyPct,
                                                   receiver_.tickHz(), actual_.bits);
  if (resolution != FailReason::NONE) { finish(false, resolution); return false; }
  Serial.printf("STAGE %lu/%lu requested=%luHz actual=%luHz duty=%.2f%% bits=%u\n",
    stageIndex_ + 1, stageCount_, requestedHz_, actual_.actualHz, actual_.actualDutyPct, actual_.bits);
  char f[12], one[24], two[24]; Display::formatFrequency(actual_.actualHz, f, sizeof(f));
  snprintf(one, sizeof(one), "F %s D %.1f%%", f, actual_.actualDutyPct);
  snprintf(two, sizeof(two), "%lu/%lu RUN", stageIndex_ + 1, stageCount_); display_.show(one, two);
  if (static_cast<Role>(settings_.role) == Role::SOLO && !startLocalMeasurement(actual_.actualHz, actual_.actualDutyPct)) {
    finish(false, FailReason::UNSUPPORTED); return false;
  }
  return true;
}

bool App::startLocalMeasurement(float hz, float duty) {
  return measurement_.start(hz, duty, params_.accuracyPct, params_.testTimeMs, params_.repeats, PWM_SETTLE_CYCLES);
}

void App::stagePassed() {
  pwm_.stop();
  if (++stageIndex_ >= stageCount_) { finish(true, FailReason::NONE); return; }
  if (static_cast<Role>(settings_.role) == Role::SOLO) { if (prepareStage()) state_ = AppState::SOLO_MEASURE; }
  else if (static_cast<Role>(settings_.role) == Role::MASTER) {
    requestedHz_ = frequencyAt(params_.startHz, params_.endHz, params_.stepHz, stageIndex_);
    pendingPacket_ = makePacket(MessageType::PREPARE); sendCurrent(MessageType::PREPARE);
    state_ = AppState::MASTER_WAIT_READY; retries_ = 0; deadlineMs_ = millis() + LINK_REPLY_TIMEOUT_MS;
  }
}

void App::startMasterDiscovery() {
  session_ = esp_random(); if (!session_) session_ = 1; sequence_ = 1; havePeer_ = false; radio_.flush();
  pendingPacket_ = makePacket(MessageType::DISCOVER); radio_.sendBroadcast(pendingPacket_);
  lastSendMs_ = millis(); deadlineMs_ = millis() + LINK_DISCOVERY_TIMEOUT_MS; retries_ = 0;
  state_ = AppState::MASTER_DISCOVER; display_.show("MASTER SEARCH", "WAIT SLAVE"); Serial.println("ESP-NOW DISCOVER");
}

ProtocolPacket App::makePacket(MessageType type) const {
  ProtocolPacket p = {};
  p.type = static_cast<uint8_t>(type); p.session = session_; p.stage = stageIndex_; p.sequence = sequence_;
  p.requestedHz = requestedHz_; p.actualHz = actual_.actualHz;
  p.actualDutyX100 = static_cast<uint16_t>(actual_.actualDutyPct * 100.0f + 0.5f);
  p.testTimeMs = params_.testTimeMs; p.repeats = params_.repeats;
  p.accuracyX100 = static_cast<uint16_t>(params_.accuracyPct * 100.0f + 0.5f); p.settleCycles = PWM_SETTLE_CYCLES;
  return p;
}

void App::sendCurrent(MessageType type) {
  ++sequence_; pendingPacket_ = makePacket(type); radio_.sendTo(peer_, pendingPacket_); lastSendMs_ = millis();
}

bool App::packetForCurrent(const ProtocolPacket &p) const {
  return p.session == session_ && p.stage == stageIndex_;
}

void App::handleRadio() {
  ReceivedPacket r;
  while (radio_.receive(r)) {
    const MessageType type = static_cast<MessageType>(r.packet.type);
    if (state_ != AppState::SLAVE_MEASURE)
      Serial.printf("ESP-NOW RX type=%u session=%08lX stage=%u seq=%u\n", r.packet.type, r.packet.session, r.packet.stage, r.packet.sequence);
    if (state_ == AppState::SLAVE_READY && type == MessageType::DISCOVER) {
      memcpy(peer_, r.mac, 6); havePeer_ = true; session_ = r.packet.session; stageIndex_ = 0; sequence_ = r.packet.sequence;
      ProtocolPacket ack = makePacket(MessageType::DISCOVER_ACK); ack.sequence = r.packet.sequence; radio_.sendTo(peer_, ack);
      state_ = AppState::SLAVE_WAIT_START; display_.show("SLAVE LINKED", "WAIT PREPARE"); continue;
    }
    if (state_ == AppState::MASTER_DISCOVER && type == MessageType::DISCOVER_ACK && r.packet.session == session_) {
      memcpy(peer_, r.mac, 6); havePeer_ = true; requestedHz_ = frequencyAt(params_.startHz, params_.endHz, params_.stepHz, stageIndex_);
      sendCurrent(MessageType::PREPARE); state_ = AppState::MASTER_WAIT_READY; retries_ = 0; deadlineMs_ = millis() + LINK_REPLY_TIMEOUT_MS;
      char mac[20]; Radio::macText(peer_, mac, sizeof(mac)); Serial.printf("SLAVE SELECTED %s\n", mac); continue;
    }
    if (state_ == AppState::SLAVE_WAIT_START && type == MessageType::DISCOVER &&
        r.packet.session == session_ && !memcmp(peer_, r.mac, 6)) {
      ProtocolPacket ack = makePacket(MessageType::DISCOVER_ACK);
      ack.sequence = r.packet.sequence; radio_.sendTo(peer_, ack); continue;
    }
    if (havePeer_ && !memcmp(peer_, r.mac, 6) && type == MessageType::RESULT &&
        r.packet.session == session_ && r.packet.stage < stageIndex_) {
      ProtocolPacket ack = {}; ack.type = static_cast<uint8_t>(MessageType::ACK);
      ack.session = session_; ack.stage = r.packet.stage; ack.sequence = r.packet.sequence;
      radio_.sendTo(peer_, ack); continue; // idempotent ACK for a retried old result
    }
    if (!havePeer_ || memcmp(peer_, r.mac, 6) || !packetForCurrent(r.packet)) continue;
    if (type == MessageType::ABORT) { finish(false, FailReason::ABORTED); continue; }
    if (state_ == AppState::MASTER_WAIT_READY && type == MessageType::READY) {
      if (!prepareStage()) continue;
      sendCurrent(MessageType::START_STAGE); state_ = AppState::MASTER_WAIT_RESULT; retries_ = 0; deadlineMs_ = millis() +
        params_.testTimeMs * params_.repeats + LINK_REPLY_TIMEOUT_MS + (1000UL * PWM_SETTLE_CYCLES / actual_.actualHz) + 20;
    } else if (state_ == AppState::MASTER_WAIT_RESULT && type == MessageType::RESULT) {
      ProtocolPacket ack = makePacket(MessageType::ACK); ack.sequence = r.packet.sequence; radio_.sendTo(peer_, ack); pwm_.stop();
      if (!r.packet.passed) finish(false, static_cast<FailReason>(r.packet.reason)); else stagePassed();
    } else if (state_ == AppState::SLAVE_WAIT_START && type == MessageType::PREPARE) {
      params_.testTimeMs = r.packet.testTimeMs; params_.repeats = r.packet.repeats;
      params_.accuracyPct = r.packet.accuracyX100 / 100.0f; requestedHz_ = r.packet.requestedHz;
      ProtocolPacket ready = makePacket(MessageType::READY); ready.sequence = r.packet.sequence; radio_.sendTo(peer_, ready);
    } else if (state_ == AppState::SLAVE_WAIT_START && type == MessageType::START_STAGE) {
      actual_.actualHz = r.packet.actualHz; actual_.actualDutyPct = r.packet.actualDutyX100 / 100.0f;
      if (!startLocalMeasurement(actual_.actualHz, actual_.actualDutyPct)) { finish(false, FailReason::UNSUPPORTED); continue; }
      state_ = AppState::SLAVE_MEASURE; deadlineMs_ = millis() + params_.testTimeMs * params_.repeats + LINK_REPLY_TIMEOUT_MS;
    } else if (state_ == AppState::SLAVE_WAIT_ACK && type == MessageType::ACK && r.packet.sequence == pendingPacket_.sequence) {
      if (pendingPacket_.passed) { ++stageIndex_; state_ = AppState::SLAVE_WAIT_START; display_.show("SLAVE READY", "WAIT PREPARE"); }
      else finish(false, static_cast<FailReason>(pendingPacket_.reason));
    }
  }
}

void App::updateMaster() {
  const uint32_t now = millis();
  if (state_ == AppState::MASTER_DISCOVER) {
    if (now >= deadlineMs_) { finish(false, FailReason::LINK_LOST); return; }
    if (now - lastSendMs_ >= LINK_RETRY_INTERVAL_MS) { radio_.sendBroadcast(pendingPacket_); lastSendMs_ = now; }
    return;
  }
  if (now < deadlineMs_) return;
  if (retries_ >= LINK_PACKET_RETRIES) { finish(false, FailReason::LINK_LOST); return; }
  radio_.sendTo(peer_, pendingPacket_); ++retries_;
  deadlineMs_ = now + (state_ == AppState::MASTER_WAIT_RESULT ?
    params_.testTimeMs * params_.repeats + LINK_REPLY_TIMEOUT_MS : LINK_REPLY_TIMEOUT_MS);
}

void App::updateSlave() {
  if (state_ == AppState::SLAVE_MEASURE) {
    const MeasureState ms = measurement_.update();
    if (ms != MeasureState::PASS && ms != MeasureState::FAIL) return;
    printStageStats(measurement_.stats(), requestedHz_);
    pendingPacket_ = makePacket(MessageType::RESULT); pendingPacket_.passed = ms == MeasureState::PASS;
    pendingPacket_.reason = static_cast<uint8_t>(measurement_.reason()); pendingPacket_.periods = measurement_.stats().periods;
    pendingPacket_.minPeriodTicks = measurement_.stats().minPeriod; pendingPacket_.maxPeriodTicks = measurement_.stats().maxPeriod;
    pendingPacket_.sequence = ++sequence_; radio_.sendTo(peer_, pendingPacket_);
    state_ = AppState::SLAVE_WAIT_ACK; retries_ = 0; deadlineMs_ = millis() + LINK_REPLY_TIMEOUT_MS;
  } else if (state_ == AppState::SLAVE_WAIT_ACK && millis() >= deadlineMs_) {
    if (retries_++ >= LINK_PACKET_RETRIES) finish(false, FailReason::LINK_LOST);
    else { radio_.sendTo(peer_, pendingPacket_); deadlineMs_ = millis() + LINK_REPLY_TIMEOUT_MS; }
  }
}

void App::sendAbort() { if (havePeer_) sendCurrent(MessageType::ABORT); }

void App::abortTest() { sendAbort(); measurement_.abort(); finish(false, FailReason::ABORTED); }

void App::finish(bool pass, FailReason reason) {
  pwm_.stop(); receiver_.stop();
  if (state_ != AppState::IDLE && state_ != AppState::MENU) radio_.end();
  state_ = AppState::FINISHED; pendingReason_ = reason;
  char one[24];
  if (pass) { snprintf(one, sizeof(one), "PASS %luHz-%lu", params_.startHz, params_.endHz); display_.show(one, "START=REPEAT"); }
  else { snprintf(one, sizeof(one), "FAIL AT %lu", requestedHz_); display_.show(one, failName(reason)); }
  Serial.printf("TEST %s: %s\n", pass ? "PASS" : "FAIL", failName(reason));
}

void App::printConfiguration() {
  const char *board = TARGET_IS_C3 ? "ESP32-C3" : "ESP32-S3";
  uint8_t mac[6] = {}; esp_read_mac(mac, ESP_MAC_WIFI_STA);
  Serial.printf("\nOptical Channel Tester | %s | mode=%s\n", board, roleName(static_cast<Role>(settings_.role)));
  Serial.printf("MAC=%02X:%02X:%02X:%02X:%02X:%02X\n", mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
  Serial.printf("GPIO PWM=%u RX=%u START=%u MODE=%u SDA=%u SCL=%u\n", GPIO_PWM, GPIO_RX,
    GPIO_BUTTON_START, GPIO_BUTTON_MODE, GPIO_SDA, GPIO_SCL);
  Serial.printf("Test %lu..%lu step %lu Hz, accuracy %.2f%%, %lums x%u, duty %u%%\n",
    params_.startHz, params_.endHz, params_.stepHz, params_.accuracyPct, params_.testTimeMs, params_.repeats, params_.dutyPct);
  stageCount_ = frequencyPointCount(params_.startHz, params_.endHz, params_.stepHz);
  Serial.printf("Frequencies (%lu): ", stageCount_);
  for (uint32_t i = 0; i < stageCount_; ++i) Serial.printf("%lu%s", frequencyAt(params_.startHz, params_.endHz, params_.stepHz, i), i + 1 == stageCount_ ? "\n" : ",");
  Serial.printf("ALL nominal: %llu us | RX=%s\n", actualNominalTotalUs(), receiver_.highRateBackend() ? "RMT DMA" : "RMT ping-pong");
}

uint64_t App::actualNominalTotalUs() {
  uint64_t total = 0;
  const uint32_t count = frequencyPointCount(params_.startHz, params_.endHz, params_.stepHz);
  for (uint32_t i = 0; i < count; ++i) {
    ActualPwm preview = {};
    const uint32_t requested = frequencyAt(params_.startHz, params_.endHz, params_.stepHz, i);
    const uint32_t actualHz = pwm_.preview(requested, params_.dutyPct, preview) ? preview.actualHz : requested;
    total += (1000000ULL * PWM_SETTLE_CYCLES + actualHz - 1) / actualHz;
    total += static_cast<uint64_t>(params_.testTimeMs) * 1000ULL * params_.repeats;
  }
  return total;
}

void App::printStageStats(const StageStats &s, uint32_t hz) {
  if (!s.periods) return;
  Serial.printf("STATS %luHz periods=%lu period ticks min/avg/max=%lu/%llu/%lu active=%lu/%llu/%lu\n",
    hz, s.periods, s.minPeriod, s.periodSum / s.periods, s.maxPeriod,
    s.minActive, s.activeSum / s.periods, s.maxActive);
  if (s.reason != FailReason::NONE) Serial.printf("FIRST BAD repeat=%u period=%lu f=%.3f duty=%.3f reason=%s\n",
    s.firstBadRepeat, s.firstBadPeriod, s.badFrequency, s.badDuty, failName(s.reason));
}