#include <Wire.h>
#include <EEPROM.h>

#define DEBUG     1

#define I2C_ADDR 0x08
#define POT_PIN  26

#define WINDOW_SIZE 10

// Upper bound on how many rounds any magazine design can hold. Sizes the
// calibration buffers and the flash record; the real capacity of THIS magazine
// is discovered during calibration (rounds fired = capacity).
#define MAG_MAX_CAPACITY 32

// --- Calibration I2C protocol (master -> magazine writes) -------------------
// The master (controller) writes a single command byte to enter/exit
// calibration. While calibrating, reads from this device return
// [CAL_MARKER][plateaus recorded] instead of the ammo count, so the master can
// show live progress on the display.
#define MAG_CMD_CAL_START 0xA1  // begin recording: clears any previous session
#define MAG_CMD_CAL_SAVE  0xA2  // build table from recorded plateaus, write flash
#define MAG_CMD_CAL_ABORT 0xA3  // leave calibration, discard recordings
#define CAL_MARKER        0xCA  // first reply byte while in calibration mode
                                // (impossible as an ammo high byte: 0xCA00 = 51712 rounds)

// --- Plateau detection tuning ------------------------------------------------
// During calibration the smoothed pot reading is watched for "plateaus": the
// value must sit still (within CAL_STABLE_TOL counts) for CAL_STABLE_MS, and a
// new plateau must be at least CAL_MIN_STEP_DELTA counts ABOVE the previous one
// (readings increase as rounds leave the magazine). Plateau sequence for a
// magazine of capacity N: the first plateau after CAL_START is the fully-loaded
// reading (N rounds), then one plateau per follower step down to 0 rounds, and
// finally the OVERTRAVEL plateau -- after the last chambered round is fired the
// bolt is held open by the follower itself, which rises measurably past the
// "0 rounds" position. That final reading is the "truly out of ammo" state.
#define CAL_STABLE_TOL      20    // counts of wiggle still considered "stable"
#define CAL_STABLE_MS       500   // how long the reading must hold to count
#define CAL_MIN_STEP_DELTA  50    // min counts between consecutive plateaus

// --- Flash (EEPROM-emulation) calibration record -----------------------------
// The Arduino-Pico core emulates EEPROM in the last flash sector; commit()
// performs the actual flash write. Layout is a single struct at offset 0 with
// a magic + checksum so a blank/corrupt sector falls back to the compiled-in
// defaults below.
#define EEPROM_SIZE   256
#define CAL_MAGIC     0x4D414743UL   // "MAGC"
#define CAL_VERSION   2              // v2 added the overtravel reading; v1
                                     // records are rejected (recalibrate)

// Sentinel for "no overtravel reading available" (compiled-in defaults only;
// every saved calibration includes one). With no overtravel reading the
// true-empty state can't be detected and the count degrades gracefully to
// mag-only + chamber inference.
#define OVERTRAVEL_NONE 0xFFFF

struct CalRecord {
  uint32_t magic;
  uint8_t  version;
  uint8_t  capacity;                          // rounds this magazine holds
  uint16_t readings[MAG_MAX_CAPACITY + 1];    // avg pot reading per round count
  uint16_t overtravel;                        // follower past "0": mag empty,
                                              // bolt held open, chamber empty
  uint16_t checksum;                          // sum of all preceding bytes
};

// Compiled-in default table, used until a calibration has been saved to flash
// (or if the flash record is corrupt). Same data as the original hard-coded
// table: average raw pot reading (12-bit, 0..4095) for each round count on the
// 10-round prototype magazine. Readings DECREASE as rounds are added.
#define DEFAULT_CAPACITY 10
static const uint16_t defaultRoundAverage[DEFAULT_CAPACITY + 1] = {
  4095,  // 0 rounds (empty)   -- measured
  3850,  // 1                  -- measured
  3670,  // 2                  -- measured
  3520,  // 3                  -- measured
  3390,  // 4                  -- estimate
  3200,  // 5                  -- estimate
  3030,  // 6                  -- estimate
  2860,  // 7                  -- estimate
  2690,  // 8                  -- estimate
  2540,  // 9                  -- estimate
  2380,  // 10                 -- estimate
};

// Active table used by potToAmmo(). Index == number of rounds in the magazine.
// The decision thresholds are NOT stored directly: each round count owns the
// band of readings closest to its average, so the boundary between count n and
// n+1 is the midpoint (g_roundAverage[n] + g_roundAverage[n+1]) / 2, computed
// at runtime in potToAmmo().
uint8_t  g_capacity = DEFAULT_CAPACITY;
uint16_t g_roundAverage[MAG_MAX_CAPACITY + 1];
uint16_t g_overtravelReading = OVERTRAVEL_NONE;  // set by calibration (v2 flash)

// Operating mode, shared with the I2C request ISR.
#define MODE_NORMAL 0
#define MODE_CAL    1
volatile uint8_t g_mode = MODE_NORMAL;
volatile uint8_t g_pendingCmd = 0;      // last command byte from the master
volatile uint8_t g_calStepsForIsr = 0;  // plateau count, mirrored for the ISR

uint16_t samples[WINDOW_SIZE];
uint8_t  sampleIdx = 0;
uint32_t sampleSum = 0;
bool     bufferFull = false;

// Rounds AVAILABLE TO FIRE (rounds in the magazine plus the chambered round,
// when we know about it), as reported to the master. Not the raw follower
// count -- see the chamber-inference block below.
volatile uint16_t ammoCount = 0;

// --- Chamber inference ---------------------------------------------------
// The follower only measures rounds in the MAGAZINE, but the number the user
// wants is rounds available to fire = mag + (1 if a round is chambered). The
// charging stroke is where they diverge: the follower drops one step but the
// round only moved into the chamber, so the displayed count must not change.
//
// Chamber state is tracked HERE (not on the eport) on purpose: this MCU loses
// power the instant the magazine is pulled, and anything done to the chamber
// while the mag is out (ejecting a live round by hand, etc.) is unobservable
// anyway -- so a power cycle doubling as "chamber state unknown" is exactly
// the right reset. Rules:
//   - Boot (= insertion; the bolt is always closed then, since the hold-open
//     is the follower itself): chamber UNKNOWN, report the mag count alone.
//     Pessimistic by design: a retained chambered round from a tactical
//     reload under-reports by exactly 1 until the first shot resolves it.
//   - Follower drops while chamber is not LOADED: that round was stripped
//     into the chamber -> chamber = LOADED, reported total unchanged.
//   - Follower drops while chamber is LOADED: the previously chambered round
//     left the gun (fired, or ejected by a second rack) -> total drops.
//   - Overtravel reading: mag empty + bolt held open -> chamber is certainly
//     EMPTY, report 0. This is the absolute resync point every time the gun
//     runs dry, so the error can never accumulate.
//
// Count transitions are only committed after the mapped value has held steady
// for a debounce window, because the strip happens during the most violent
// part of the bolt cycle when the follower is bouncing. Downward steps and
// overtravel commit fast; upward steps (physically impossible while inserted,
// so almost certainly noise) need to persist much longer.
#define CH_UNKNOWN 0
#define CH_EMPTY   1
#define CH_LOADED  2

#define COUNT_STABLE_DOWN_MS  150   // strips / overtravel
#define COUNT_STABLE_UP_MS   1000   // upward "corrections" (noise recovery)

uint8_t       g_chamber = CH_UNKNOWN;
bool          magStateValid = false;   // baseline committed since boot/cal-exit?
uint16_t      committedCount = 0;      // accepted rounds-in-mag
bool          committedOvertravel = false;
uint16_t      candCount = 0;           // candidate awaiting stability
bool          candOvertravel = false;
unsigned long candSince = 0;
bool          candValid = false;

// Calibration session state (loop-only).
uint16_t      calReadings[MAG_MAX_CAPACITY + 2];  // plateaus, in firing order
uint8_t       calStepCount = 0;
uint16_t      calCandidate = 0;         // reading we're watching for stability
unsigned long calCandidateSince = 0;
bool          calCandidateStable = false;

// Translate a raw pot reading into a round count using the active table. A
// reading is assigned to whichever round count it is closest to: walk from
// empty (0) toward full and stop at the first count whose lower boundary (the
// midpoint to the next count) the reading has not yet crossed. Overtravel
// readings (see isOvertravel) also map to 0 here; callers that care check
// isOvertravel() first.
uint16_t potToAmmo(uint16_t raw) {
  for (uint16_t n = 0; n < g_capacity; n++) {
    uint16_t threshold = (g_roundAverage[n] + g_roundAverage[n + 1]) / 2;
    if (raw >= threshold) {
      return n;
    }
  }
  return g_capacity;
}

// True when the reading is past the "0 rounds" position: the magazine is empty
// AND the follower is holding the bolt open, which on this design can only
// mean the chamber is empty too (the bolt hold-open IS the follower, and a
// rearward bolt has extracted whatever was chambered). This is the only
// ground-truth "totally out of ammo" signal we have. Boundary is the midpoint
// between the calibrated 0-rounds reading and the overtravel reading.
bool isOvertravel(uint16_t raw) {
  if (g_overtravelReading == OVERTRAVEL_NONE) return false;
  uint16_t threshold =
      (uint16_t)(((uint32_t)g_roundAverage[0] + g_overtravelReading) / 2);
  return raw >= threshold;
}

// Forget everything we thought we knew about the chamber and the follower.
// Called on calibration exit; boot gets the same effect from the initializers.
void resetChamberTracking() {
  g_chamber = CH_UNKNOWN;
  magStateValid = false;
  candValid = false;
}

// Publish rounds-available to the master-facing register.
void publishAmmo() {
  ammoCount = committedOvertravel
                  ? 0
                  : committedCount + (g_chamber == CH_LOADED ? 1 : 0);
}

// Normal-mode counterpart of updateCalibration(): map the smoothed pot value
// to (rounds-in-mag, overtravel), debounce it, and run the chamber rules on
// each committed transition. See the chamber-inference block up top.
void updateAmmo(uint16_t value) {
  bool     ot    = isOvertravel(value);
  uint16_t count = ot ? 0 : potToAmmo(value);

  // (Re)start the stability timer whenever the mapped state moves.
  if (!candValid || count != candCount || ot != candOvertravel) {
    candCount = count;
    candOvertravel = ot;
    candSince = millis();
    candValid = true;
  }

  if (!magStateValid) {
    // First reading after boot/cal-exit: report it immediately (chamber is
    // UNKNOWN, so mag count alone -- or 0 on overtravel), and latch it as the
    // baseline once stable. No strip inference on the baseline itself.
    ammoCount = ot ? 0 : count;
    if (millis() - candSince >= COUNT_STABLE_DOWN_MS) {
      magStateValid = true;
      committedCount = candCount;
      committedOvertravel = candOvertravel;
      if (committedOvertravel) g_chamber = CH_EMPTY;
      publishAmmo();
    }
    return;
  }

  if (candCount == committedCount && candOvertravel == committedOvertravel) {
    return;  // steady state
  }

  // Rounds can't be ADDED to an inserted magazine, so an upward count is
  // noise until proven otherwise. Everything else (strips, overtravel, bolt
  // eased forward over an empty mag) commits on the short window.
  bool upward = candCount > committedCount;
  if (millis() - candSince < (upward ? COUNT_STABLE_UP_MS
                                     : COUNT_STABLE_DOWN_MS)) {
    return;
  }

  uint16_t oldCount = committedCount;
  bool wasOvertravel = committedOvertravel;
  committedCount = candCount;
  committedOvertravel = candOvertravel;

  if (committedOvertravel || wasOvertravel) {
    // In (or leaving) overtravel the chamber is empty: entering it means the
    // last round was fired (or a rack ejected it); leaving it toward "0"
    // means the bolt was eased forward over the still-empty magazine.
    g_chamber = CH_EMPTY;
  } else if (committedCount < oldCount && g_chamber != CH_LOADED) {
    // First strip with the chamber not known-loaded: the round went INTO the
    // chamber, not out of the gun. Reported total stays flat because the +1
    // now compensates the mag's -1. (If several steps were missed at once,
    // only one of them can be the chambering; the rest fell out of the total
    // via committedCount, which is correct -- they were fired.)
    g_chamber = CH_LOADED;
  }
  // Strips while LOADED need no flag change: each cycle fires/ejects the
  // chambered round and chambers the next, so total = mag + 1 keeps tracking.

  publishAmmo();

#if DEBUG
  Serial.print("Mag state: count ");
  Serial.print(committedCount);
  Serial.print(committedOvertravel ? " (overtravel)" : "");
  Serial.print(", chamber ");
  Serial.print(g_chamber == CH_LOADED ? "LOADED"
               : g_chamber == CH_EMPTY ? "EMPTY" : "UNKNOWN");
  Serial.print(" -> available ");
  Serial.println(ammoCount);
#endif
}

// Sum of the struct's bytes, excluding the trailing checksum field itself.
uint16_t calChecksum(const CalRecord &rec) {
  const uint8_t *p = (const uint8_t *)&rec;
  uint16_t sum = 0;
  for (size_t i = 0; i < offsetof(CalRecord, checksum); i++) sum += p[i];
  return sum;
}

// Load the calibration record from flash into the active table; fall back to
// the compiled-in defaults if the record is missing or invalid.
void loadCalibration() {
  CalRecord rec;
  EEPROM.get(0, rec);

  bool valid = rec.magic == CAL_MAGIC &&
               rec.version == CAL_VERSION &&
               rec.capacity >= 1 && rec.capacity <= MAG_MAX_CAPACITY &&
               rec.checksum == calChecksum(rec);

  if (valid) {
    g_capacity = rec.capacity;
    for (uint8_t n = 0; n <= g_capacity; n++) g_roundAverage[n] = rec.readings[n];
    g_overtravelReading = rec.overtravel;
#if DEBUG
    Serial.print("Loaded calibration from flash: capacity ");
    Serial.print(g_capacity);
    Serial.print(", overtravel ");
    Serial.println(g_overtravelReading);
#endif
  } else {
    g_capacity = DEFAULT_CAPACITY;
    for (uint8_t n = 0; n <= g_capacity; n++) g_roundAverage[n] = defaultRoundAverage[n];
    g_overtravelReading = OVERTRAVEL_NONE;  // defaults predate overtravel data
#if DEBUG
    Serial.println("No valid calibration in flash; using compiled-in defaults.");
#endif
  }
}

// Persist the active table to flash. commit() briefly stalls the core for the
// flash write; the master isn't polling us at that moment (it just sent SAVE
// and is showing the confirmation screen), and its bus timeout would recover
// it anyway.
void saveCalibration() {
  CalRecord rec = {};
  rec.magic = CAL_MAGIC;
  rec.version = CAL_VERSION;
  rec.capacity = g_capacity;
  for (uint8_t n = 0; n <= g_capacity; n++) rec.readings[n] = g_roundAverage[n];
  rec.overtravel = g_overtravelReading;
  rec.checksum = calChecksum(rec);

  EEPROM.put(0, rec);
  bool ok = EEPROM.commit();
#if DEBUG
  Serial.print("Calibration flash write: ");
  Serial.println(ok ? "OK" : "FAILED");
#endif
}

// Reset the plateau tracker for a new calibration session.
void startCalibration() {
  calStepCount = 0;
  g_calStepsForIsr = 0;
  calCandidate = 0;
  calCandidateSince = millis();
  calCandidateStable = false;
  g_mode = MODE_CAL;
#if DEBUG
  Serial.println("Calibration START: waiting for fully-loaded plateau...");
#endif
}

// Turn the recorded plateaus into the active table and save it. Recorded in
// firing order: calReadings[0] is the fully-loaded reading (capacity rounds),
// calReadings[stepCount-2] is empty-mag/bolt-closed (0 rounds, last round
// chambered), and calReadings[stepCount-1] is the overtravel reading (last
// round fired, bolt held open). So capacity = stepCount - 2 and the round
// table is the recording reversed, minus the final overtravel entry. Needs at
// least 3 plateaus (full + one strip + overtravel, i.e. a 1-round magazine
// fully fired out); otherwise the old table is kept. The user must fire the
// gun completely dry -- including the last chambered round -- before saving,
// so the overtravel reading is always captured.
void finishCalibration() {
  if (calStepCount < 3) {
#if DEBUG
    Serial.println("Calibration SAVE ignored: fewer than 3 plateaus recorded.");
#endif
    return;
  }

  uint8_t capacity = calStepCount - 2;
  for (uint8_t n = 0; n <= capacity; n++) {
    g_roundAverage[n] = calReadings[capacity - n];
  }
  g_overtravelReading = calReadings[calStepCount - 1];
  g_capacity = capacity;

#if DEBUG
  Serial.print("Calibration SAVE: capacity ");
  Serial.print(capacity);
  Serial.print(", overtravel ");
  Serial.print(g_overtravelReading);
  Serial.print(", table (0..full): ");
  for (uint8_t n = 0; n <= capacity; n++) {
    Serial.print(g_roundAverage[n]);
    Serial.print(n < capacity ? ", " : "\n");
  }
#endif

  saveCalibration();
}

// Watch the smoothed pot value for a new stable plateau and record it. Runs
// every loop pass while calibrating.
void updateCalibration(uint16_t value) {
  // Buffer holds positions capacity..0 plus the overtravel plateau.
  if (calStepCount >= MAG_MAX_CAPACITY + 2) return;  // buffer full, stop recording

  // Reading moved: restart the stability timer on the new value.
  int delta = (int)value - (int)calCandidate;
  if (delta < 0) delta = -delta;
  if (delta > CAL_STABLE_TOL) {
    calCandidate = value;
    calCandidateSince = millis();
    calCandidateStable = false;
    return;
  }

  if (calCandidateStable || millis() - calCandidateSince < CAL_STABLE_MS) return;
  calCandidateStable = true;

  // First plateau (fully loaded) is always accepted; later ones must be a real
  // step UP from the previous plateau (rounds leaving raise the reading).
  if (calStepCount > 0 &&
      (int)value < (int)calReadings[calStepCount - 1] + CAL_MIN_STEP_DELTA) {
    return;
  }

  calReadings[calStepCount++] = value;
  g_calStepsForIsr = calStepCount;
#if DEBUG
  Serial.print("Calibration plateau ");
  Serial.print(calStepCount);
  Serial.print(": ");
  Serial.println(value);
#endif
}

// Master wrote to us: a single command byte. Just record it -- the real work
// (including flash writes) happens in loop(), not in the ISR.
void receiveHandler(int numBytes) {
  if (numBytes < 1) return;
  uint8_t cmd = Wire.read();
  while (Wire.available()) Wire.read();
  g_pendingCmd = cmd;
}

void requestHandler() {
  if (g_mode == MODE_CAL) {
    uint8_t buf[2] = { CAL_MARKER, g_calStepsForIsr };
    Wire.write(buf, 2);
  } else {
    uint16_t a = ammoCount;
    uint8_t buf[2] = {
      (uint8_t)(a >> 8),
      (uint8_t)(a & 0xFF)
    };
    Wire.write(buf, 2);
  }
}

void setup() {
  analogReadResolution(12);
  pinMode(POT_PIN, INPUT);

#if DEBUG
  Serial.begin(115200);
#endif

  EEPROM.begin(EEPROM_SIZE);
  loadCalibration();

  Wire.setSDA(12);
  Wire.setSCL(13);
  Wire.begin(I2C_ADDR);
  Wire.onReceive(receiveHandler);
  Wire.onRequest(requestHandler);
}

void loop() {
  uint16_t raw = analogRead(POT_PIN);

  sampleSum -= samples[sampleIdx];
  samples[sampleIdx] = raw;
  sampleSum += raw;
  sampleIdx = (sampleIdx + 1) % WINDOW_SIZE;

  if (!bufferFull && sampleIdx == 0)
    bufferFull = true;

  uint8_t count = bufferFull ? WINDOW_SIZE : sampleIdx;
  uint16_t potValue = (count > 0) ? sampleSum / count : raw;

  // Act on any command the master wrote since the last pass.
  noInterrupts();
  uint8_t cmd = g_pendingCmd;
  g_pendingCmd = 0;
  interrupts();

  switch (cmd) {
    case MAG_CMD_CAL_START:
      startCalibration();
      break;
    case MAG_CMD_CAL_SAVE:
      if (g_mode == MODE_CAL) {
        finishCalibration();
        g_mode = MODE_NORMAL;
        // Re-baseline: after a full calibration run the gun sits at
        // overtravel, which the fresh tracking will read as a certain 0.
        resetChamberTracking();
      }
      break;
    case MAG_CMD_CAL_ABORT:
      if (g_mode == MODE_CAL) {
#if DEBUG
        Serial.println("Calibration ABORTED; keeping previous table.");
#endif
        g_mode = MODE_NORMAL;
        resetChamberTracking();
      }
      break;
    default:
      break;
  }

  if (g_mode == MODE_CAL) {
    updateCalibration(potValue);
  } else {
    updateAmmo(potValue);
  }

#if DEBUG
  Serial.print(potValue);
  Serial.print(" -> ");
  if (g_mode == MODE_CAL) {
    Serial.print("CAL plateaus: ");
    Serial.println(calStepCount);
  } else {
    Serial.println(ammoCount);
  }
#endif

  delay(10);
}
