#include <Wire.h>

// I2C
#define I2C_SDA       12
#define I2C_SCL       13
#define POT_I2C_ADDR  0x08
#define BNO055_ADDR   0x28
#define DISPLAY_ADDR  0x42

// Display I2C register map (must match display_proto.ino). Payloads are
// little-endian (LSB first), matching how the display parses them.
#define DISP_REG_BATTERY      0x01  // uint16 millivolts
#define DISP_REG_FLASHLIGHT   0x02  // uint8  0 = off, else on
#define DISP_REG_AMMO         0x03  // uint16 rounds remaining
#define DISP_REG_HEADING      0x04  // uint16 degrees 0..359
#define DISP_REG_BATTERY_PCT  0x05  // uint8  battery charge 0..100 %
#define DISP_REG_LOAD_TIMER   0x06  // 2 bytes [state][percent]:
                                    //   state  0 = ready/idle, 1 = on (timeout
                                    //          counting down), 2 = cooldown
                                    //   percent 0..100 = time remaining in phase
#define DISP_REG_CAL          0x07  // 2 bytes [phase][value]: magazine
                                    //   calibration UI (phases below)

// Load-timer phase values sent in DISP_REG_LOAD_TIMER.
#define LOAD_TIMER_READY     0
#define LOAD_TIMER_ACTIVE    1
#define LOAD_TIMER_COOLDOWN  2

// Calibration UI phases sent in DISP_REG_CAL (must match display.ino).
#define CAL_PHASE_OFF        0   // normal HUD
#define CAL_PHASE_PROMPT     1   // "fully load the magazine, then insert it"
#define CAL_PHASE_RECORDING  2   // value = rounds fired so far
#define CAL_PHASE_SAVED      3   // value = discovered capacity
#define CAL_PHASE_ERROR      4   // failed/aborted, nothing saved

// Button events read back from the display slave (1 byte, clear-on-read).
// Button 3 on the display drives magazine calibration: long press enters (or
// cancels) calibration mode, short click saves the finished calibration.
#define DISP_BTN_EV_SHORT  0x01
#define DISP_BTN_EV_LONG   0x02

// Commands written to the magazine slave (single byte). While calibrating,
// reads from the magazine return [MAG_CAL_MARKER][plateaus recorded] instead
// of the 2-byte ammo count.
#define MAG_CMD_CAL_START  0xA1
#define MAG_CMD_CAL_SAVE   0xA2
#define MAG_CMD_CAL_ABORT  0xA3
#define MAG_CAL_MARKER     0xCA

// BNO055 registers
#define BNO055_CHIP_ID_REG    0x00
#define BNO055_OPR_MODE_REG   0x3D
#define BNO055_PWR_MODE_REG   0x3E
#define BNO055_SYS_TRIGGER_REG 0x3F
#define BNO055_EULER_H_LSB    0x1A
#define BNO055_CHIP_ID_VALUE  0xA0

// BNO055 operating modes
#define BNO055_MODE_CONFIG    0x00
#define BNO055_MODE_NDOF      0x0C

// Pins
#define BTN1_PIN      14   // BUT1 — left flashlight / load button (active low)
#define BTN2_PIN      15   // BUT2 — right flashlight / load button (active low)
#define LOAD_PIN      10   // RP2350 GPIO 10
#define BATTERY_ADC   26   // GPIO26 / ADC0

// Battery voltage divider values
// R_TOP = resistor from battery + to ADC pin
// R_BOTTOM = resistor from ADC pin to GND
#define R_TOP         100000.0
#define R_BOTTOM      100000.0

// Battery charge mapping (single-cell LiPo). These are display-only limits used
// to turn the measured voltage into a 0..100% bar -- tweak them to match the
// usable range of your pack. Voltage <= EMPTY reads 0%, >= FULL reads 100%.
#define BATT_FULL_MV   4000   // fully charged (100%); a full LiPo reads ~4.0 V
                              // under load on this hardware, not the 4.2 V
                              // open-circuit figure
#define BATT_EMPTY_MV  3300   // considered empty (0%)

// Pot-to-ammo translation now lives on the magazine firmware: it reads the pot,
// maps it to a round count via its own threshold table, infers whether a round
// is chambered (the charging stroke moves the follower without a shot being
// fired) and reports rounds AVAILABLE TO FIRE directly over I2C. This master
// just forwards that count to the display.
#define AMMO_DISCONNECTED 0xFFFF   // sentinel sent when the magazine read fails

// I2C bus tolerance. The magazine pot (POT_I2C_ADDR) lives on a connector that
// is unplugged/replugged constantly during normal use, and it shares SDA/SCL
// with the display and compass. A device that vanishes mid-transfer can leave
// the shared bus wedged (SDA held low), which would otherwise hang the master
// AND kill every other peripheral. Bounding each transfer and enabling auto
// recovery (end() -> 9-clock SCL unstick -> begin()) keeps the master
// responsive and self-healing no matter when the magazine comes or goes.
#define I2C_TIMEOUT_MS  25

// Timing
const unsigned long debounceDelay = 50;
const unsigned long doubleTapWindow = 200;
const unsigned long strobeInterval = 50;
const unsigned long batteryLogInterval = 2000;
const unsigned long potPollInterval = 50;
const unsigned long compassPollInterval = 100;
const unsigned long loadTimeLimit = 60000;
const unsigned long loadCooldown = 30000;

// Button state
bool lastRawButtonState = false;
bool debouncedButtonState = false;
unsigned long lastDebounceTime = 0;

// Load state
bool loadState = false;
bool strobeMode = false;
bool strobeOutputState = false;

// Load timeout/cooldown: the load auto-shuts off after loadTimeLimit, then is
// locked out for loadCooldown before it can be turned on again.
unsigned long loadOnStartTime = 0;
bool loadInCooldown = false;
unsigned long loadCooldownStartTime = 0;

// Tap tracking
bool waitingForSecondTap = false;
unsigned long firstTapTime = 0;

// Strobe timing
unsigned long lastStrobeTime = 0;

// Battery logging
unsigned long lastBatteryLogTime = 0;

// Pot polling. The magazine is a hot-pluggable slave, so we track its presence
// and only redraw the ammo / log on a connected<->disconnected transition (or a
// genuine value change) rather than spamming the bus and console every poll.
unsigned long lastPotPollTime = 0;
int lastAmmoSent = -1;          // last ammo value pushed to the display
int magazinePresentState = -1;  // -1 = unknown, 0 = disconnected, 1 = connected

// Compass (heading is reported to the display only; it no longer drives the load)
unsigned long lastCompassPollTime = 0;
float compassHeading = 0.0;
bool bnoFound = false;

// Last flashlight/load status pushed to the display (so we only send on change).
// Initialized to a sentinel so the first call always transmits.
int lastFlashlightSent = -1;

// Load-timer reporting to the display. While a phase is active we resend the
// progress periodically so the bar animates; otherwise we only send on change.
const unsigned long loadTimerSendInterval = 200;
unsigned long lastLoadTimerSendTime = 0;
int lastLoadTimerStateSent = -1;

// --- Magazine calibration state machine -------------------------------------
// Entered by a long press of the display's action button. The user is told to
// fully load the magazine and insert it; when the magazine (re)appears on the
// bus it's commanded into calibration mode and records a pot plateau for every
// follower position: the full-mag baseline, one step per round stripped, and
// finally the overtravel plateau (after the last chambered round is fired the
// follower rises past "empty" to hold the bolt open). The user must therefore
// fire the gun COMPLETELY dry before saving. A short click saves: the magazine
// reverses its recording into a threshold table (capacity = plateaus - 2,
// since the baseline and overtravel plateaus aren't rounds) and writes it to
// its flash. A long press at any point cancels.
enum CalState {
  CAL_IDLE,       // not calibrating; normal ammo polling
  CAL_PROMPT,     // waiting for the user to load + insert the magazine
  CAL_RECORDING,  // magazine is recording plateaus; user fires rounds
  CAL_DONE_WAIT   // showing the saved/error screen before returning to HUD
};
CalState calState = CAL_IDLE;
unsigned long calStateSince = 0;
const unsigned long buttonPollInterval = 100;   // display button-event poll
const unsigned long calPollInterval    = 100;   // magazine probe/progress poll
const unsigned long calDoneShowMs      = 2500;  // how long saved/error stays up
const unsigned long calMagLostTimeout  = 2000;  // mag silent this long = fail
unsigned long lastButtonPollTime = 0;
unsigned long lastCalPollTime = 0;
unsigned long calLastMagOkTime = 0;
bool calMagSeenAbsent = false;   // must see the mag ABSENT before arming, so a
                                 // still-inserted mag isn't calibrated as-is
int  calPlateauCount = 0;        // plateaus recorded (rounds fired + 1)
int  calLastRoundsSent = -1;     // last rounds-fired value pushed to display

void setup() {
  Serial.begin(115200);
  delay(2000);

  Wire.setSDA(I2C_SDA);
  Wire.setSCL(I2C_SCL);
  Wire.begin();
  // Make every bus transfer non-blocking and self-healing so a magazine that
  // disconnects (even mid-transfer) can't hang the loop or wedge the shared bus
  // for the display/compass. On timeout the core runs the 9-clock SCL recovery
  // and re-inits the master automatically.
  Wire.setTimeout(I2C_TIMEOUT_MS, true);

  pinMode(LOAD_PIN, OUTPUT);
  digitalWrite(LOAD_PIN, LOW);

  // Ambidextrous flashlight / load buttons on CN5 (active low, internal pull-ups).
  pinMode(BTN1_PIN, INPUT_PULLUP);
  pinMode(BTN2_PIN, INPUT_PULLUP);

  analogReadResolution(12);

  // BNO055 can be slow to boot — give it extra time
  Serial.println("Waiting for BNO055 to boot...");
  delay(1000);

  // Scan I2C bus to show what's present
  Serial.println("I2C scan:");
  bool displaySeen = false;
  for (uint8_t addr = 1; addr < 127; addr++) {
    Wire.beginTransmission(addr);
    if (Wire.endTransmission() == 0) {
      Serial.print("  Found device at 0x");
      if (addr < 16) Serial.print("0");
      Serial.println(addr, HEX);
      if (addr == DISPLAY_ADDR) displaySeen = true;
    }
  }

  // Explicit verdict on the display so we don't have to read the scan by eye.
  if (displaySeen) {
    Serial.println("Display: ACKed at 0x42. Bus path to display is OK.");
  } else {
    Serial.println("Display: NOT FOUND at 0x42 -- check wiring (SDA->GP12, SCL->GP13),");
    Serial.println("         common ground, SDA/SCL not swapped, and that the slave is running.");
  }

  // Try to init BNO055 directly via registers
  unsigned long bnoStart = millis();
  while (millis() - bnoStart < 3000) {
    if (bnoReadByte(BNO055_CHIP_ID_REG) == BNO055_CHIP_ID_VALUE) {
      bnoFound = true;
      break;
    }
    delay(100);
  }

  if (bnoFound) {
    bnoWriteByte(BNO055_OPR_MODE_REG, BNO055_MODE_CONFIG);
    delay(25);
    bnoWriteByte(BNO055_PWR_MODE_REG, 0x00); // Normal power mode
    bnoWriteByte(BNO055_SYS_TRIGGER_REG, 0x00);
    delay(10);
    bnoWriteByte(BNO055_OPR_MODE_REG, BNO055_MODE_NDOF);
    delay(20);
    Serial.println("BNO055 found and configured.");
  } else {
    Serial.println("BNO055 not found. Continuing without compass.");
  }

  Serial.println("Ready.");
  Serial.println("Single press toggles load. Double-tap starts strobe. Press again to exit strobe.");
}

void loop() {
  handleButton();
  handleStrobe();
  updateLoadTimer();
  logBatteryVoltage();
  pollDisplayButtons();
  runCalibration();
  pollPotValue();       // no-op while calibration owns the magazine
  pollCompass();
  updateLoadOutput();
  updateDisplayFlashlight();
  updateDisplayLoadTimer();
}

// --- Display I2C senders ---------------------------------------------------
// The display expects [reg][LSB][MSB] for 16-bit values and [reg][value] for
// 8-bit values, so write the low byte first to match its little-endian parse.

// Tracks the last endTransmission() result for the display so we only log when
// it changes (these senders fire often and would otherwise spam the console).
uint8_t g_lastDispTxStatus = 0xFF;

// Returns true if the display ACKed. Logs (on change) so we can tell whether
// the slave at DISPLAY_ADDR is even acknowledging on the bus.
bool displayEnd() {
  uint8_t status = Wire.endTransmission();
  if (status != g_lastDispTxStatus) {
    g_lastDispTxStatus = status;
    if (status == 0) {
      Serial.println("Display I2C: ACK OK from 0x42.");
    } else {
      Serial.print("Display I2C: endTransmission error ");
      Serial.print(status);
      // 2 = address NACK (no slave answering), 3 = data NACK, 5 = timeout.
      if (status == 2) Serial.print(" (no ACK on address 0x42 - slave not on this bus?)");
      Serial.println();
    }
  }
  return status == 0;
}

void displaySendU16(uint8_t reg, uint16_t value) {
  Wire.beginTransmission(DISPLAY_ADDR);
  Wire.write(reg);
  Wire.write((uint8_t)(value & 0xFF));
  Wire.write((uint8_t)(value >> 8));
  displayEnd();
}

void displaySendU8(uint8_t reg, uint8_t value) {
  Wire.beginTransmission(DISPLAY_ADDR);
  Wire.write(reg);
  Wire.write(value);
  displayEnd();
}

void displaySendU8x2(uint8_t reg, uint8_t a, uint8_t b) {
  Wire.beginTransmission(DISPLAY_ADDR);
  Wire.write(reg);
  Wire.write(a);
  Wire.write(b);
  displayEnd();
}

// Maps a battery voltage (mV) to a 0..100% charge using the configurable
// BATT_EMPTY_MV / BATT_FULL_MV limits. Linear approximation -- good enough for
// a charge bar; a real fuel gauge would use a discharge curve.
uint8_t batteryPercent(uint16_t mv) {
  if (mv <= BATT_EMPTY_MV) return 0;
  if (mv >= BATT_FULL_MV) return 100;
  return (uint8_t)((uint32_t)(mv - BATT_EMPTY_MV) * 100UL /
                   (BATT_FULL_MV - BATT_EMPTY_MV));
}

// Pushes the flashlight/load status to the display, but only when it changes
// so we don't flood the bus. "On" covers manual load and strobe mode.
void updateDisplayFlashlight() {
  bool on = loadState || strobeMode;
  if ((int)on != lastFlashlightSent) {
    lastFlashlightSent = (int)on;
    displaySendU8(DISP_REG_FLASHLIGHT, on ? 1 : 0);
  }
}

// Reports the load timeout/cooldown to the display as a phase + "percent
// remaining" so it can draw a progress bar. Resends periodically while a phase
// is running so the bar animates; only on change when idle.
void updateDisplayLoadTimer() {
  uint8_t state;
  uint8_t pct;

  if (loadInCooldown) {
    state = LOAD_TIMER_COOLDOWN;
    unsigned long elapsed = millis() - loadCooldownStartTime;
    pct = (elapsed >= loadCooldown)
            ? 0
            : (uint8_t)(100UL - (elapsed * 100UL / loadCooldown));
  } else if (loadState) {
    state = LOAD_TIMER_ACTIVE;
    unsigned long elapsed = millis() - loadOnStartTime;
    pct = (elapsed >= loadTimeLimit)
            ? 0
            : (uint8_t)(100UL - (elapsed * 100UL / loadTimeLimit));
  } else {
    state = LOAD_TIMER_READY;
    pct = 0;
  }

  bool stateChanged = ((int)state != lastLoadTimerStateSent);
  bool periodicDue = (state != LOAD_TIMER_READY) &&
                     (millis() - lastLoadTimerSendTime >= loadTimerSendInterval);

  if (stateChanged || periodicDue) {
    lastLoadTimerStateSent = state;
    lastLoadTimerSendTime = millis();
    displaySendU8x2(DISP_REG_LOAD_TIMER, state, pct);
  }
}

// --- Magazine calibration ----------------------------------------------------

// Writes a single command byte to the magazine. Returns true if it ACKed.
bool magSendCmd(uint8_t cmd) {
  Wire.beginTransmission(POT_I2C_ADDR);
  Wire.write(cmd);
  return Wire.endTransmission() == 0;
}

// Address-only probe: is the magazine on the bus right now?
bool magProbe() {
  Wire.beginTransmission(POT_I2C_ADDR);
  return Wire.endTransmission() == 0;
}

// Leave calibration: clear the display's cal screen and reset the normal ammo
// polling trackers so the (possibly re-calibrated) count is re-sent and
// presence is re-announced.
void exitCalibration() {
  displaySendU8x2(DISP_REG_CAL, CAL_PHASE_OFF, 0);
  calState = CAL_IDLE;
  lastAmmoSent = -1;
  magazinePresentState = -1;
}

// Poll the display slave for queued button-3 events (1 byte, clear-on-read).
void pollDisplayButtons() {
  if (millis() - lastButtonPollTime < buttonPollInterval) return;
  lastButtonPollTime = millis();

  uint8_t n = Wire.requestFrom(DISPLAY_ADDR, 1);
  if (n != 1) {
    while (Wire.available()) Wire.read();
    return;
  }
  uint8_t ev = Wire.read();
  if (ev & DISP_BTN_EV_LONG)  handleCalLongPress();
  if (ev & DISP_BTN_EV_SHORT) handleCalShortPress();
}

// Long press: enter calibration from idle, cancel it anywhere else.
void handleCalLongPress() {
  if (calState == CAL_IDLE) {
    Serial.println("Calibration: entering magazine calibration mode.");
    calState = CAL_PROMPT;
    calStateSince = millis();
    calMagSeenAbsent = false;
    calPlateauCount = 0;
    calLastRoundsSent = -1;
    displaySendU8x2(DISP_REG_CAL, CAL_PHASE_PROMPT, 0);
  } else if (calState == CAL_PROMPT || calState == CAL_RECORDING) {
    Serial.println("Calibration: cancelled by long press.");
    magSendCmd(MAG_CMD_CAL_ABORT);   // best effort; mag may not be inserted
    exitCalibration();
  }
}

// Short click while recording: save. Needs at least 3 plateaus (the full-mag
// baseline, at least one strip, and the final overtravel plateau from firing
// the gun dry); otherwise there's nothing to calibrate. Must match the
// magazine's own >= 3 check in its finishCalibration().
void handleCalShortPress() {
  if (calState != CAL_RECORDING) return;

  if (calPlateauCount >= 3) {
    uint8_t capacity = (uint8_t)(calPlateauCount - 2);
    Serial.print("Calibration: saving, capacity ");
    Serial.println(capacity);
    magSendCmd(MAG_CMD_CAL_SAVE);    // mag builds the table + writes flash
    displaySendU8x2(DISP_REG_CAL, CAL_PHASE_SAVED, capacity);
  } else {
    Serial.println("Calibration: save pressed with nothing recorded -- failing.");
    magSendCmd(MAG_CMD_CAL_ABORT);
    displaySendU8x2(DISP_REG_CAL, CAL_PHASE_ERROR, 0);
  }
  calState = CAL_DONE_WAIT;
  calStateSince = millis();
}

// Drives the calibration state machine (magazine probing/progress polling and
// the timed return to the HUD). Button transitions live in the handlers above.
void runCalibration() {
  if (calState == CAL_IDLE) return;

  if (calState == CAL_DONE_WAIT) {
    if (millis() - calStateSince >= calDoneShowMs) exitCalibration();
    return;
  }

  if (millis() - lastCalPollTime < calPollInterval) return;
  lastCalPollTime = millis();

  if (calState == CAL_PROMPT) {
    // Require an absent->present transition so the user actually pulls the
    // mag to load it; a mag left inserted from before doesn't arm recording.
    bool present = magProbe();
    if (!present) {
      calMagSeenAbsent = true;
    } else if (calMagSeenAbsent && magSendCmd(MAG_CMD_CAL_START)) {
      Serial.println("Calibration: magazine inserted, recording started.");
      calState = CAL_RECORDING;
      calStateSince = millis();
      calLastMagOkTime = millis();
      calPlateauCount = 0;
      calLastRoundsSent = -1;
      displaySendU8x2(DISP_REG_CAL, CAL_PHASE_RECORDING, 0);
    }
    return;
  }

  // CAL_RECORDING: poll the magazine for [marker][plateaus] progress.
  uint8_t bytesReceived = Wire.requestFrom(POT_I2C_ADDR, 2);
  if (bytesReceived == 2) {
    uint8_t marker = Wire.read();
    uint8_t steps = Wire.read();
    if (marker == MAG_CAL_MARKER) {
      calLastMagOkTime = millis();
      calPlateauCount = steps;
      // Plateau #1 is the full-mag baseline and plateau #2 is the charging
      // stroke (a round chambered, nothing fired yet), so rounds fired =
      // plateaus - 2. The final overtravel plateau lands the moment the last
      // chambered round fires, so this stays exact through the end of the run.
      int roundsFired = (steps > 2) ? steps - 2 : 0;
      if (roundsFired != calLastRoundsSent) {
        calLastRoundsSent = roundsFired;
        displaySendU8x2(DISP_REG_CAL, CAL_PHASE_RECORDING, (uint8_t)roundsFired);
        Serial.print("Calibration: rounds fired ");
        Serial.println(roundsFired);
      }
      return;
    }
  }
  while (Wire.available()) Wire.read();

  // Bad read or wrong marker: tolerate brief glitches, fail if the magazine
  // has been gone too long (pulled mid-calibration, wedged, etc.).
  if (millis() - calLastMagOkTime > calMagLostTimeout) {
    Serial.println("Calibration: magazine lost -- aborting.");
    magSendCmd(MAG_CMD_CAL_ABORT);
    displaySendU8x2(DISP_REG_CAL, CAL_PHASE_ERROR, 0);
    calState = CAL_DONE_WAIT;
    calStateSince = millis();
  }
}

void handleButton() {
  // Either panel button (L/R) counts as pressed; both are active-low.
  bool rawButtonPressed =
      (digitalRead(BTN1_PIN) == LOW) || (digitalRead(BTN2_PIN) == LOW);

  if (rawButtonPressed != lastRawButtonState) {
    lastDebounceTime = millis();
  }

  if ((millis() - lastDebounceTime) > debounceDelay) {
    if (rawButtonPressed != debouncedButtonState) {
      debouncedButtonState = rawButtonPressed;

      if (debouncedButtonState == true) {
        handleButtonPress();
      }
    }
  }

  lastRawButtonState = rawButtonPressed;

  // If we got one tap and no second tap arrived, treat it as a single press
  if (waitingForSecondTap && (millis() - firstTapTime > doubleTapWindow)) {
    waitingForSecondTap = false;

    if (!strobeMode) {
      toggleLoad();
    }
  }
}

void handleButtonPress() {
  if (strobeMode) {
    strobeMode = false;
    strobeOutputState = false;

    Serial.println("Strobe OFF");
    return;
  }

  if (!waitingForSecondTap) {
    waitingForSecondTap = true;
    firstTapTime = millis();
  } else {
    // Second tap arrived quickly, activate strobe
    waitingForSecondTap = false;
    strobeMode = true;
    strobeOutputState = false;
    lastStrobeTime = millis();

    Serial.println("Strobe ON");
  }
}

void toggleLoad() {
  if (!loadState) {
    // Trying to turn the load on -- blocked while cooling down.
    if (loadInCooldown) {
      Serial.println("Load is in cooldown; ignoring turn-on.");
      return;
    }
    loadState = true;
    loadOnStartTime = millis();
    Serial.println("Load is now ON");
  } else {
    loadState = false;
    Serial.println("Load is now OFF");
  }
}

// Enforces the load timeout/cooldown each loop: clears an expired cooldown, and
// if the load has been on past loadTimeLimit, shuts it off and starts cooldown.
void updateLoadTimer() {
  if (loadInCooldown && (millis() - loadCooldownStartTime >= loadCooldown)) {
    loadInCooldown = false;
    Serial.println("Load cooldown complete; ready again.");
  }

  if (loadState && (millis() - loadOnStartTime >= loadTimeLimit)) {
    loadState = false;
    loadInCooldown = true;
    loadCooldownStartTime = millis();
    Serial.println("Load timed out; entering cooldown.");
  }
}

void handleStrobe() {
  if (!strobeMode) {
    return;
  }

  if (millis() - lastStrobeTime >= strobeInterval) {
    lastStrobeTime = millis();

    strobeOutputState = !strobeOutputState;
    digitalWrite(LOAD_PIN, strobeOutputState ? HIGH : LOW);
  }
}

void updateLoadOutput() {
  if (strobeMode) return; // strobe handles its own output

  digitalWrite(LOAD_PIN, loadState ? HIGH : LOW);
}

void pollPotValue() {
  // While calibrating, the magazine is answering with calibration progress
  // (not ammo) and the display is showing the cal screen, so normal ammo
  // polling/forwarding pauses. exitCalibration() resets the trackers below so
  // everything re-syncs when we return.
  if (calState != CAL_IDLE) {
    return;
  }

  if (millis() - lastPotPollTime < potPollInterval) {
    return;
  }

  lastPotPollTime = millis();

  // requestFrom is bounded by I2C_TIMEOUT_MS and self-recovers the bus on a
  // wedged disconnect, so this can never hang the loop. A clean unplug NACKs and
  // returns 0; a half-connected magazine may time out, both land in the else.
  uint8_t bytesReceived = Wire.requestFrom(POT_I2C_ADDR, 2);
  if (bytesReceived == 2) {
    uint8_t highByte = Wire.read();
    uint8_t lowByte = Wire.read();
    // The magazine already translated the pot reading into a round count.
    uint16_t ammo = ((uint16_t)highByte << 8) | lowByte;

    if (magazinePresentState != 1) {
      magazinePresentState = 1;
      Serial.println("Magazine connected.");
    }

    if ((int)ammo != lastAmmoSent) {
      lastAmmoSent = (int)ammo;
      displaySendU16(DISP_REG_AMMO, ammo);
      Serial.print("Magazine ammo: ");
      Serial.println(ammo);
    }
  } else {
    // Drain any stray bytes from a partial/aborted transfer before moving on.
    while (Wire.available()) Wire.read();

    if (magazinePresentState != 0) {
      magazinePresentState = 0;
      Serial.println("Magazine disconnected -> ammo: -- (showing disconnected)");
    }

    if (lastAmmoSent != (int)AMMO_DISCONNECTED) {
      lastAmmoSent = (int)AMMO_DISCONNECTED;
      displaySendU16(DISP_REG_AMMO, AMMO_DISCONNECTED);
    }
  }
}

void pollCompass() {
  if (!bnoFound) return;

  if (millis() - lastCompassPollTime < compassPollInterval) return;
  lastCompassPollTime = millis();

  uint8_t lsb = bnoReadByte(BNO055_EULER_H_LSB);
  uint8_t msb = bnoReadByte(BNO055_EULER_H_LSB + 1);
  int16_t rawHeading = ((int16_t)msb << 8) | lsb;
  compassHeading = rawHeading / 16.0;

  // Normalize to 0..359 before sending to the display.
  int headingDeg = (int)lround(compassHeading) % 360;
  if (headingDeg < 0) headingDeg += 360;
  displaySendU16(DISP_REG_HEADING, (uint16_t)headingDeg);

  //Serial.print("Compass heading: ");
  //Serial.print(compassHeading, 1);
  //Serial.println(" deg");
}

uint8_t bnoReadByte(uint8_t reg) {
  Wire.beginTransmission(BNO055_ADDR);
  Wire.write(reg);
  Wire.endTransmission(false);
  Wire.requestFrom(BNO055_ADDR, (uint8_t)1);
  return Wire.available() ? Wire.read() : 0;
}

void bnoWriteByte(uint8_t reg, uint8_t value) {
  Wire.beginTransmission(BNO055_ADDR);
  Wire.write(reg);
  Wire.write(value);
  Wire.endTransmission();
}

void logBatteryVoltage() {
  if (millis() - lastBatteryLogTime < batteryLogInterval) {
    return;
  }

  lastBatteryLogTime = millis();

  int raw = analogRead(BATTERY_ADC);

  float adcVoltage = raw * (3.3 / 4095.0);
  float dividerRatio = (R_TOP + R_BOTTOM) / R_BOTTOM;
  float batteryVoltage = adcVoltage * dividerRatio;

  // Display expects millivolts as a uint16 (e.g. 3.85 V -> 3850), plus a
  // separate 0..100% charge estimate for the fill bar.
  uint16_t batteryMv = (uint16_t)lround(batteryVoltage * 1000.0);
  uint8_t batteryPct = batteryPercent(batteryMv);
  displaySendU16(DISP_REG_BATTERY, batteryMv);
  displaySendU8(DISP_REG_BATTERY_PCT, batteryPct);

  Serial.print("Battery ADC raw: ");
  Serial.print(raw);

  Serial.print(" | ADC voltage: ");
  Serial.print(adcVoltage, 3);
  Serial.print(" V");

  Serial.print(" | Battery voltage: ");
  Serial.print(batteryVoltage, 3);
  Serial.print(" V");

  Serial.print(" | Charge: ");
  Serial.print(batteryPct);
  Serial.println(" %");
}
