#include <Button.h>
#include <Adafruit_NeoPixel.h>
#include <Wire.h>

// without USB CDC, Serial falls back to USART1 on PA9/PA10, hijacking the mode and right buttons
#ifndef USBD_USE_CDC
#error "Set Tools > USB support to CDC (generic 'Serial' supersede U(S)ART)"
#endif

/*
Records the IMU to internal flash to work out what a crash looks like, download with tools/crash_dump.py

MODE        start/stop recording, starting erases the previous recording, hold 5s to power off
LEFT/RIGHT  add marker 1/2 to the recording, e.g. "crash now" and "bump now"

DASH        dim blue: empty, green: recording stored (bar shows length), amber: erasing,
            red: recording (bar shows storage used), blue: downloading, purple blinking: error

ONBOARD LED fallback for when the dash has no power without the battery
            off: idle, fast flicker: erasing, slow blink: recording, short flicker: marker added,
            on: downloading, rapid blink: error

serial commands: r record, s stop, d dump, i info, l live sensor reading
*/

constexpr int LEDS_TAIL_PIN = PA6;
constexpr int LEDS_TAIL_LENGTH = 10;

constexpr int LEDS_LEFT_PIN = PA5;
constexpr int LEDS_RIGHT_PIN = PA7;
constexpr int LEDS_HEAD_LENGTH = 39;

constexpr int LEDS_DASH_PIN = PB0;
constexpr int LEDS_DASH_LENGTH = 2;
constexpr int LEDS_DASH_BRIGHTNESS = 100;

constexpr int PIN_ALIVE = PA2;
constexpr int PIN_BUCK = PA1;

constexpr int PIN_PEDAL = PA0;
constexpr int PIN_PEDAL_FORWARD = PA3;
constexpr int PIN_PEDAL_REVERSE = PA4;

constexpr int PIN_BUTTON_MODE = PA9;
constexpr int PIN_BUTTON_LEFT = PA8;
constexpr int PIN_BUTTON_RIGHT = PA10;

constexpr int PIN_IMU_SCL = PB6;
constexpr int PIN_IMU_SDA = PB7;

constexpr int LED_ON = LOW;
constexpr int LED_OFF = HIGH;

constexpr int BUCK_ON = HIGH;
constexpr int BUCK_OFF = LOW;

constexpr unsigned long SHUTDOWN_HOLD = 5000; // ms
constexpr unsigned long IDLE_TIMEOUT = 180000; // ms, not while recording
constexpr unsigned long STATUS_INTERVAL = 50; // ms
constexpr unsigned long MARKER_FLASH = 150; // ms

// MPU9250/9255
constexpr uint8_t IMU_ADDRESSES[] = {0x68, 0x69}; // AD0 low, AD0 high
constexpr uint8_t IMU_REG_SMPLRT_DIV = 0x19;
constexpr uint8_t IMU_REG_CONFIG = 0x1A;
constexpr uint8_t IMU_REG_GYRO_CONFIG = 0x1B;
constexpr uint8_t IMU_REG_ACCEL_CONFIG = 0x1C;
constexpr uint8_t IMU_REG_ACCEL_CONFIG2 = 0x1D;
constexpr uint8_t IMU_REG_FIFO_EN = 0x23;
constexpr uint8_t IMU_REG_INT_PIN_CFG = 0x37;
constexpr uint8_t IMU_REG_INT_ENABLE = 0x38;
constexpr uint8_t IMU_REG_INT_STATUS = 0x3A;
constexpr uint8_t IMU_REG_USER_CTRL = 0x6A;
constexpr uint8_t IMU_REG_PWR_MGMT_1 = 0x6B;
constexpr uint8_t IMU_REG_FIFO_COUNT_H = 0x72;
constexpr uint8_t IMU_REG_FIFO_R_W = 0x74;
constexpr uint8_t IMU_REG_WHO_AM_I = 0x75;

constexpr uint8_t IMU_FIFO_EN_ACCEL_GYRO = 0x78;
constexpr uint8_t IMU_USER_CTRL_FIFO_EN = 0x40;
constexpr uint8_t IMU_USER_CTRL_FIFO_RST = 0x04;
constexpr uint8_t IMU_INT_FIFO_OFLOW = 0x10;
constexpr uint8_t IMU_INT_PIN_CFG_BYPASS = 0x02; // exposes the magnetometer on the main bus

constexpr int SAMPLE_RATE = 500; // Hz
constexpr uint8_t IMU_DLPF_184HZ = 0x01; // 1kHz internal rate, used for gyro and accel
constexpr uint8_t IMU_ACCEL_16G = 0x18;
constexpr uint8_t IMU_GYRO_2000DPS = 0x18;
constexpr uint16_t ACCEL_LSB_PER_G = 2048;
constexpr uint16_t GYRO_LSB_PER_DPS_X10 = 164;

constexpr int FIFO_SIZE = 512;
constexpr int FIFO_SAMPLE_SIZE = 12; // accel xyz + gyro xyz, big-endian
constexpr int FIFO_READ_BATCH = 2; // samples per I2C read, stays within the 32 byte Wire buffer
constexpr unsigned long FIFO_INTERVAL = 10; // ms, the fifo holds ~84ms at 500Hz

// AK8963 magnetometer inside the MPU9250/9255
constexpr uint8_t MAG_ADDRESS = 0x0C;
constexpr uint8_t MAG_REG_WIA = 0x00;
constexpr uint8_t MAG_REG_ST1 = 0x02;
constexpr uint8_t MAG_REG_HXL = 0x03;
constexpr uint8_t MAG_REG_CNTL1 = 0x0A;
constexpr uint8_t MAG_WIA = 0x48;
constexpr uint8_t MAG_ST1_DRDY = 0x01;
constexpr uint8_t MAG_ST2_HOFL = 0x08;
constexpr uint8_t MAG_CNTL1_POWER_DOWN = 0x00;
constexpr uint8_t MAG_CNTL1_16BIT_100HZ = 0x16;
constexpr unsigned long MAG_INTERVAL = 10; // ms

// sectors 0-3 (64KB) hold this sketch, sectors 4-7 (448KB) hold the recording
constexpr uint32_t STORAGE_START = 0x08010000;
constexpr uint32_t STORAGE_END = 0x08080000;
constexpr uint32_t STORAGE_SIZE = STORAGE_END - STORAGE_START;
constexpr uint32_t STORAGE_SECTOR_FIRST = FLASH_SECTOR_4;
constexpr int STORAGE_SECTOR_COUNT = 4;

constexpr float RECORD_BYTES_PER_SECOND = SAMPLE_RATE * (1 + FIFO_SAMPLE_SIZE) + (1000 / MAG_INTERVAL) * (1 + 6);

// recording format, a stream of tagged records, keep in sync with tools/crash_dump.py
enum RecordTag : uint8_t {
  TAG_HEADER = 0x01, // 12 bytes: "QLR", version, WHO_AM_I, mag ok, rate, accel LSB/g, gyro LSB/dps * 10 (uint16 LE)
  TAG_IMU = 0x02, // 12 bytes: raw fifo sample, accel xyz + gyro xyz, int16 BE
  TAG_MAG = 0x03, // 6 bytes: raw magnetometer xyz, int16 LE, 0.15uT/LSB without sensitivity adjustment
  TAG_INPUTS = 0x04, // 1 byte: raw pin levels, see INPUT_BIT_*, low is active
  TAG_MARKER = 0x05, // 1 byte: marker number
  TAG_OVERFLOW = 0x06, // 0 bytes: fifo overflowed, an unknown number of samples was lost
  TAG_END = 0xFF, // erased flash
};

constexpr uint8_t RECORD_VERSION = 1;

constexpr uint8_t INPUT_BIT_PEDAL = 0;
constexpr uint8_t INPUT_BIT_FORWARD = 1;
constexpr uint8_t INPUT_BIT_REVERSE = 2;
constexpr uint8_t INPUT_BIT_MODE = 3;
constexpr uint8_t INPUT_BIT_LEFT = 4;
constexpr uint8_t INPUT_BIT_RIGHT = 5;

// end of the program in flash, from the linker script
extern "C" uint32_t _sidata, _sdata, _edata;

Button buttonMode(PIN_BUTTON_MODE);
Button buttonLeft(PIN_BUTTON_LEFT);
Button buttonRight(PIN_BUTTON_RIGHT);

Adafruit_NeoPixel ledsTail(LEDS_TAIL_LENGTH, LEDS_TAIL_PIN, NEO_GRBW + NEO_KHZ800);
Adafruit_NeoPixel ledsHeadLeft(LEDS_HEAD_LENGTH, LEDS_LEFT_PIN, NEO_GRBW + NEO_KHZ800);
Adafruit_NeoPixel ledsHeadRight(LEDS_HEAD_LENGTH, LEDS_RIGHT_PIN, NEO_GRBW + NEO_KHZ800);
Adafruit_NeoPixel ledsDash(LEDS_DASH_LENGTH, LEDS_DASH_PIN, NEO_GRBW + NEO_KHZ800);

const uint32_t COLOR_OFF = ledsDash.Color(0, 0, 0, 0);
const uint32_t COLOR_DIM = ledsDash.Color(0, 0, 0, 10);
const uint32_t COLOR_WHITE = ledsDash.Color(0, 0, 0, 255);
const uint32_t COLOR_EMPTY = ledsDash.Color(0, 0, 40, 0);
const uint32_t COLOR_STORED = ledsDash.Color(0, 255, 0, 0);
const uint32_t COLOR_ERASING = ledsDash.Color(255, 60, 0, 0);
const uint32_t COLOR_RECORDING = ledsDash.Color(255, 0, 0, 0);
const uint32_t COLOR_DUMPING = ledsDash.Color(0, 0, 255, 0);
const uint32_t COLOR_ERROR = ledsDash.Color(255, 0, 255, 0);

enum State {
  STATE_IDLE,
  STATE_ERASING,
  STATE_RECORDING,
  STATE_DUMPING,
  STATE_ERROR,
  STATE_OFF,
};

State state = STATE_IDLE;
const char* errorReason = "";

uint8_t imuAddress = 0; // 0 when not found
uint8_t imuWhoAmI = 0;
bool magFound = false;

uint32_t storageAddress = STORAGE_START; // next word to program
uint32_t storageWord = 0xFFFFFFFF;
int storageWordBytes = 0;
bool storageFull = false;
uint32_t recordedBytes = 0;
int sectorsErased = 0;

unsigned long recordStart = 0;
unsigned long lastFifoRead = 0;
unsigned long lastMagRead = 0;
unsigned long lastStatus = 0;
unsigned long lastActivity = 0;
unsigned long markerFlashUntil = 0;
uint8_t lastInputs = 0;
uint32_t overflowCount = 0;

unsigned long modePressStart = 0;
bool modePressed = false;

// I2C

void writeRegister(uint8_t address, uint8_t reg, uint8_t value) {
  Wire.beginTransmission(address);
  Wire.write(reg);
  Wire.write(value);
  Wire.endTransmission();
}

bool readRegisters(uint8_t address, uint8_t reg, uint8_t* buffer, uint8_t length) {
  Wire.beginTransmission(address);
  Wire.write(reg);

  if (Wire.endTransmission(false) != 0 || Wire.requestFrom(address, length) != length) {
    return false;
  }

  for (int i = 0; i < length; i++) {
    buffer[i] = Wire.read();
  }

  return true;
}

uint8_t readRegister(uint8_t address, uint8_t reg) {
  uint8_t value = 0;
  readRegisters(address, reg, &value, 1);

  return value;
}

// IMU

void connectImu() {
  for (const uint8_t address : IMU_ADDRESSES) {
    uint8_t whoAmI = 0;

    if (readRegisters(address, IMU_REG_WHO_AM_I, &whoAmI, 1)) {
      imuAddress = address;
      imuWhoAmI = whoAmI;

      break;
    }
  }

  if (imuAddress == 0) {
    return;
  }

  writeRegister(imuAddress, IMU_REG_PWR_MGMT_1, 0x80); // reset
  delay(100);
  writeRegister(imuAddress, IMU_REG_PWR_MGMT_1, 0x01); // gyro PLL clock
  delay(10);

  writeRegister(imuAddress, IMU_REG_CONFIG, IMU_DLPF_184HZ);
  writeRegister(imuAddress, IMU_REG_SMPLRT_DIV, 1000 / SAMPLE_RATE - 1);
  writeRegister(imuAddress, IMU_REG_GYRO_CONFIG, IMU_GYRO_2000DPS);
  writeRegister(imuAddress, IMU_REG_ACCEL_CONFIG, IMU_ACCEL_16G);
  writeRegister(imuAddress, IMU_REG_ACCEL_CONFIG2, IMU_DLPF_184HZ);
  writeRegister(imuAddress, IMU_REG_INT_ENABLE, IMU_INT_FIFO_OFLOW);
  writeRegister(imuAddress, IMU_REG_INT_PIN_CFG, IMU_INT_PIN_CFG_BYPASS);

  magFound = readRegister(MAG_ADDRESS, MAG_REG_WIA) == MAG_WIA;

  if (magFound) {
    writeRegister(MAG_ADDRESS, MAG_REG_CNTL1, MAG_CNTL1_POWER_DOWN);
    delay(10);
    writeRegister(MAG_ADDRESS, MAG_REG_CNTL1, MAG_CNTL1_16BIT_100HZ);
  }
}

void startFifo() {
  writeRegister(imuAddress, IMU_REG_FIFO_EN, 0);
  writeRegister(imuAddress, IMU_REG_USER_CTRL, IMU_USER_CTRL_FIFO_RST);
  readRegister(imuAddress, IMU_REG_INT_STATUS); // clears a stale overflow flag
  writeRegister(imuAddress, IMU_REG_USER_CTRL, IMU_USER_CTRL_FIFO_EN);
  writeRegister(imuAddress, IMU_REG_FIFO_EN, IMU_FIFO_EN_ACCEL_GYRO);
}

void stopFifo() {
  writeRegister(imuAddress, IMU_REG_FIFO_EN, 0);
  writeRegister(imuAddress, IMU_REG_USER_CTRL, 0);
}

// storage

uint32_t programEnd() {
  return (uint32_t)&_sidata + ((uint32_t)&_edata - (uint32_t)&_sdata);
}

bool hasHeader() {
  const uint8_t* storage = (const uint8_t*)STORAGE_START;

  return storage[0] == TAG_HEADER && storage[1] == 'Q' && storage[2] == 'L' && storage[3] == 'R';
}

// recordings end at the first erased tag, the last programmed word marks where to stop downloading
uint32_t findRecordedBytes() {
  if (!hasHeader()) {
    return 0;
  }

  for (uint32_t address = STORAGE_END - 4; address >= STORAGE_START; address -= 4) {
    if (*(const uint32_t*)address != 0xFFFFFFFF) {
      return address + 4 - STORAGE_START;
    }
  }

  return 0;
}

void flushWord() {
  if (storageWordBytes == 0) {
    return;
  }

  // unwritten bytes stay 0xFF, which reads as the end tag
  if (HAL_FLASH_Program(FLASH_TYPEPROGRAM_WORD, storageAddress, storageWord) != HAL_OK) {
    storageFull = true;
    errorReason = "FLASH WRITE FAILED";
  }

  storageAddress += 4;
  storageWord = 0xFFFFFFFF;
  storageWordBytes = 0;
}

void storeByte(uint8_t value) {
  storageWord &= ~(0xFFu << (storageWordBytes * 8));
  storageWord |= (uint32_t)value << (storageWordBytes * 8);
  storageWordBytes++;

  if (storageWordBytes == 4) {
    flushWord();
  }
}

// only whole records are stored, so a full recording still parses cleanly
void storeRecord(uint8_t tag, const uint8_t* payload, int length) {
  if (storageFull || storageAddress + storageWordBytes + 1 + length > STORAGE_END) {
    storageFull = true;
    return;
  }

  storeByte(tag);

  for (int i = 0; i < length; i++) {
    storeByte(payload[i]);
  }
}

void storeHeader() {
  const uint8_t header[] = {
    'Q', 'L', 'R',
    RECORD_VERSION,
    imuWhoAmI,
    magFound,
    SAMPLE_RATE & 0xFF, SAMPLE_RATE >> 8,
    ACCEL_LSB_PER_G & 0xFF, ACCEL_LSB_PER_G >> 8,
    GYRO_LSB_PER_DPS_X10 & 0xFF, GYRO_LSB_PER_DPS_X10 >> 8,
  };

  storeRecord(TAG_HEADER, header, sizeof(header));
}

bool eraseStorage() {
  HAL_FLASH_Unlock();
  __HAL_FLASH_CLEAR_FLAG(FLASH_FLAG_EOP | FLASH_FLAG_OPERR | FLASH_FLAG_WRPERR | FLASH_FLAG_PGAERR | FLASH_FLAG_PGPERR | FLASH_FLAG_PGSERR);

  // one sector at a time so the dash can show progress, each 128KB sector takes about a second
  for (sectorsErased = 0; sectorsErased < STORAGE_SECTOR_COUNT; sectorsErased++) {
    renderStatus(millis());

    FLASH_EraseInitTypeDef erase = {};
    erase.TypeErase = FLASH_TYPEERASE_SECTORS;
    erase.Sector = STORAGE_SECTOR_FIRST + sectorsErased;
    erase.NbSectors = 1;
    erase.VoltageRange = FLASH_VOLTAGE_RANGE_3;

    uint32_t sectorError = 0;

    if (HAL_FLASHEx_Erase(&erase, &sectorError) != HAL_OK) {
      HAL_FLASH_Lock();
      return false;
    }
  }

  return true;
}

// recording

uint8_t readInputs() {
  return digitalRead(PIN_PEDAL) << INPUT_BIT_PEDAL
    | digitalRead(PIN_PEDAL_FORWARD) << INPUT_BIT_FORWARD
    | digitalRead(PIN_PEDAL_REVERSE) << INPUT_BIT_REVERSE
    | digitalRead(PIN_BUTTON_MODE) << INPUT_BIT_MODE
    | digitalRead(PIN_BUTTON_LEFT) << INPUT_BIT_LEFT
    | digitalRead(PIN_BUTTON_RIGHT) << INPUT_BIT_RIGHT;
}

void recordInputs() {
  const uint8_t inputs = readInputs();

  if (inputs != lastInputs) {
    lastInputs = inputs;
    storeRecord(TAG_INPUTS, &inputs, 1);
  }
}

void recordMarker(uint8_t marker) {
  storeRecord(TAG_MARKER, &marker, 1);
  markerFlashUntil = millis() + MARKER_FLASH;

  Serial.print("MARKER ");
  Serial.println(marker);
}

void readFifo() {
  const uint8_t status = readRegister(imuAddress, IMU_REG_INT_STATUS);

  uint8_t countBytes[2] = {0, 0};
  readRegisters(imuAddress, IMU_REG_FIFO_COUNT_H, countBytes, 2);
  const int count = (countBytes[0] & 0x1F) << 8 | countBytes[1];

  // once full, new samples overwrite the oldest and the stream loses alignment
  if (status & IMU_INT_FIFO_OFLOW || count > FIFO_SIZE - FIFO_SAMPLE_SIZE) {
    storeRecord(TAG_OVERFLOW, nullptr, 0);
    startFifo();
    overflowCount++;

    return;
  }

  uint8_t buffer[FIFO_SAMPLE_SIZE * FIFO_READ_BATCH];

  for (int samples = count / FIFO_SAMPLE_SIZE; samples > 0; samples -= FIFO_READ_BATCH) {
    const int batch = min(samples, FIFO_READ_BATCH);

    if (!readRegisters(imuAddress, IMU_REG_FIFO_R_W, buffer, batch * FIFO_SAMPLE_SIZE)) {
      return;
    }

    for (int i = 0; i < batch; i++) {
      storeRecord(TAG_IMU, buffer + i * FIFO_SAMPLE_SIZE, FIFO_SAMPLE_SIZE);
    }
  }
}

void readMag() {
  if (!(readRegister(MAG_ADDRESS, MAG_REG_ST1) & MAG_ST1_DRDY)) {
    return;
  }

  // reading through ST2 is required to release the next measurement
  uint8_t data[7];

  if (readRegisters(MAG_ADDRESS, MAG_REG_HXL, data, sizeof(data)) && !(data[6] & MAG_ST2_HOFL)) {
    storeRecord(TAG_MAG, data, 6);
  }
}

void printRecordedLength(uint32_t bytes) {
  Serial.print(bytes);
  Serial.print(" BYTES, ~");
  Serial.print(bytes / RECORD_BYTES_PER_SECOND, 1);
  Serial.println(" S");
}

void startRecording() {
  if (state != STATE_IDLE) {
    return;
  }

  Serial.println("ERASING");
  state = STATE_ERASING;

  if (!eraseStorage()) {
    state = STATE_ERROR;
    errorReason = "FLASH ERASE FAILED";
    Serial.println(errorReason);

    return;
  }

  storageAddress = STORAGE_START;
  storageWord = 0xFFFFFFFF;
  storageWordBytes = 0;
  storageFull = false;
  recordedBytes = 0;
  overflowCount = 0;

  storeHeader();

  lastInputs = readInputs();
  storeRecord(TAG_INPUTS, &lastInputs, 1);

  startFifo();

  recordStart = millis();
  lastFifoRead = recordStart;
  lastMagRead = recordStart;
  state = STATE_RECORDING;

  Serial.println("RECORDING");
}

void stopRecording() {
  if (state != STATE_RECORDING) {
    return;
  }

  stopFifo();
  flushWord();
  HAL_FLASH_Lock();

  recordedBytes = storageAddress - STORAGE_START;
  state = STATE_IDLE;
  lastActivity = millis();

  Serial.print(storageFull ? "STORAGE FULL, STOPPED, " : "STOPPED, ");
  printRecordedLength(recordedBytes);

  if (overflowCount > 0) {
    Serial.print("FIFO OVERFLOWS ");
    Serial.println(overflowCount);
  }
}

// serial

void printInfo() {
  Serial.print("IMU ");

  if (imuAddress == 0) {
    Serial.println("NOT FOUND ON PB6/PB7");
  } else {
    Serial.print("0x");
    Serial.print(imuAddress, HEX);
    Serial.print(", WHO_AM_I 0x");
    Serial.print(imuWhoAmI, HEX);
    Serial.println(magFound ? ", MAG OK" : ", MAG NOT FOUND");
  }

  Serial.print("PROGRAM ENDS 0x");
  Serial.print(programEnd(), HEX);
  Serial.print(", STORAGE 0x");
  Serial.print(STORAGE_START, HEX);
  Serial.print(" - 0x");
  Serial.println(STORAGE_END, HEX);

  if (state == STATE_ERROR) {
    Serial.print("ERROR ");
    Serial.println(errorReason);
  }

  Serial.print("RECORDING ");

  if (state == STATE_RECORDING) {
    Serial.print("IN PROGRESS, ");
    printRecordedLength(storageAddress - STORAGE_START);
  } else {
    printRecordedLength(recordedBytes);
  }

  Serial.println("INFO END");
}

// reads straight from the data registers instead of the fifo, handy for checking orientation and calibration
void printLive() {
  constexpr uint8_t IMU_REG_ACCEL_XOUT_H = 0x3B;

  Serial.print("CONFIG 0x");
  Serial.print(readRegister(imuAddress, IMU_REG_CONFIG), HEX);
  Serial.print(", GYRO_CONFIG 0x");
  Serial.print(readRegister(imuAddress, IMU_REG_GYRO_CONFIG), HEX);
  Serial.print(", ACCEL_CONFIG 0x");
  Serial.print(readRegister(imuAddress, IMU_REG_ACCEL_CONFIG), HEX);
  Serial.print(", ACCEL_CONFIG2 0x");
  Serial.println(readRegister(imuAddress, IMU_REG_ACCEL_CONFIG2), HEX);

  uint8_t data[14];

  if (!readRegisters(imuAddress, IMU_REG_ACCEL_XOUT_H, data, sizeof(data))) {
    Serial.println("LIVE READ FAILED");
    return;
  }

  // accel xyz, temperature, gyro xyz
  const char* names[] = {"AX", "AY", "AZ", "T", "GX", "GY", "GZ"};
  const float scales[] = {ACCEL_LSB_PER_G, ACCEL_LSB_PER_G, ACCEL_LSB_PER_G, 1, GYRO_LSB_PER_DPS_X10 / 10.0f, GYRO_LSB_PER_DPS_X10 / 10.0f, GYRO_LSB_PER_DPS_X10 / 10.0f};

  for (int i = 0; i < 7; i++) {
    const int16_t raw = (int16_t)(data[i * 2] << 8 | data[i * 2 + 1]);

    Serial.print(names[i]);
    Serial.print(" ");
    Serial.print(raw);

    if (i == 3) {
      Serial.print(" (");
      Serial.print(raw / 333.87f + 21, 1); // MPU9250 datasheet
      Serial.print("C)");
    } else {
      Serial.print(" (");
      Serial.print(raw / scales[i], 3);
      Serial.print(")");
    }

    Serial.print(i < 6 ? ", " : "\n");
  }
}

void dump() {
  // also works while off or in error, so a recording can be downloaded after the idle timeout
  if (state == STATE_RECORDING || state == STATE_ERASING || state == STATE_DUMPING) {
    Serial.println("DUMP UNAVAILABLE");
    return;
  }

  const State previousState = state;
  state = STATE_DUMPING;
  renderStatus(millis());

  Serial.print("DUMP ");
  Serial.println(recordedBytes);

  for (uint32_t offset = 0; offset < recordedBytes; offset += 256) {
    Serial.write((const uint8_t*)(STORAGE_START + offset), min(recordedBytes - offset, (uint32_t)256));
  }

  Serial.println();
  Serial.println("DUMP END");

  state = previousState;
}

void handleSerial() {
  while (Serial.available() > 0) {
    const char command = Serial.read();
    lastActivity = millis();

    if (command == 'r') {
      startRecording();
    } else if (command == 's') {
      stopRecording();
    } else if (command == 'd') {
      dump();
    } else if (command == 'i') {
      printInfo();
    } else if (command == 'l' && imuAddress != 0) {
      printLive();
    }
  }
}

// power

void shutdown() {
  // when STLink or USB prevents full power cut, loop keeps running, only shut down once
  if (state == STATE_OFF) {
    return;
  }

  stopRecording();

  Serial.println("SHUTTING DOWN");
  state = STATE_OFF;
  renderStatus(millis());

  digitalWrite(PIN_BUCK, BUCK_OFF);
  delay(100);

  Serial.println("SELF OFF, GOODBYE");
  digitalWrite(PIN_ALIVE, LOW);
}

// only used when STLink or USB prevents full power cut
void restart() {
  Serial.println("RESTARTING");

  digitalWrite(PIN_BUCK, BUCK_ON);
  state = imuAddress == 0 ? STATE_ERROR : STATE_IDLE;
  lastActivity = millis();
}

// status

void fillBar(uint32_t color, float fraction) {
  const int length = constrain((int)ceil(fraction * LEDS_DASH_LENGTH), 1, LEDS_DASH_LENGTH);
  ledsDash.fill(color, 0, length);
}

void renderStatus(unsigned long now) {
  ledsDash.fill(COLOR_DIM);

  const bool blink = (now / 250) % 2 == 0;

  if (state == STATE_IDLE && recordedBytes == 0) {
    ledsDash.fill(COLOR_EMPTY);
  } else if (state == STATE_IDLE) {
    fillBar(COLOR_STORED, (float)recordedBytes / STORAGE_SIZE);
  } else if (state == STATE_ERASING) {
    fillBar(COLOR_ERASING, (float)(sectorsErased + 1) / STORAGE_SECTOR_COUNT);
  } else if (state == STATE_RECORDING) {
    fillBar(COLOR_RECORDING, (float)(storageAddress - STORAGE_START) / STORAGE_SIZE);

    if (blink) {
      ledsDash.setPixelColor(LEDS_DASH_LENGTH - 1, COLOR_RECORDING);
    }
  } else if (state == STATE_DUMPING) {
    ledsDash.fill(COLOR_DUMPING);
  } else if (state == STATE_ERROR) {
    ledsDash.fill(blink ? COLOR_ERROR : COLOR_OFF);
  } else {
    ledsDash.fill(COLOR_OFF);
  }

  if (state != STATE_OFF && now < markerFlashUntil) {
    ledsDash.fill(COLOR_WHITE);
  }

  ledsDash.show();

  bool onboard = false;

  if (state == STATE_RECORDING && now < markerFlashUntil) {
    onboard = (now / 25) % 2 == 0;
  } else if (state == STATE_ERASING) {
    onboard = (now / 50) % 2 == 0;
  } else if (state == STATE_RECORDING) {
    onboard = (now / 500) % 2 == 0;
  } else if (state == STATE_DUMPING) {
    onboard = true;
  } else if (state == STATE_ERROR) {
    onboard = (now / 100) % 2 == 0;
  }

  digitalWrite(LED_BUILTIN, onboard ? LED_ON : LED_OFF);
}

void setup() {
  pinMode(LED_BUILTIN, OUTPUT);
  pinMode(PIN_ALIVE, OUTPUT);
  pinMode(PIN_BUCK, OUTPUT);

  pinMode(PIN_PEDAL, INPUT_PULLUP);
  pinMode(PIN_PEDAL_FORWARD, INPUT_PULLUP);
  pinMode(PIN_PEDAL_REVERSE, INPUT_PULLUP);

  digitalWrite(PIN_ALIVE, HIGH);
  digitalWrite(PIN_BUCK, BUCK_ON);

  buttonMode.begin();
  buttonLeft.begin();
  buttonRight.begin();

  Serial.begin(9600);

  // the other strips aren't used, make sure they don't keep showing whatever ran before
  for (Adafruit_NeoPixel* leds : {&ledsTail, &ledsHeadLeft, &ledsHeadRight}) {
    leds->begin();
    leds->clear();
    leds->show();
  }

  ledsDash.setBrightness(LEDS_DASH_BRIGHTNESS);
  ledsDash.begin();

  // setSCL/SDA have to be called before begin()
  Wire.setSCL(PIN_IMU_SCL);
  Wire.setSDA(PIN_IMU_SDA);
  Wire.begin();
  Wire.setClock(400000);

  connectImu();

  recordedBytes = findRecordedBytes();
  lastActivity = millis();

  if (programEnd() > STORAGE_START) {
    state = STATE_ERROR;
    errorReason = "PROGRAM OVERLAPS STORAGE";
  } else if (imuAddress == 0) {
    state = STATE_ERROR;
    errorReason = "IMU NOT FOUND";
  }
}

void loop() {
  const unsigned long now = millis();

  if (state == STATE_OFF && (buttonMode.pressed() || buttonLeft.pressed() || buttonRight.pressed())) {
    restart();
  } else if (state != STATE_OFF) {
    if (buttonMode.pressed()) {
      modePressed = true;
      modePressStart = now;
      lastActivity = now;
    }

    // act on release, so a long press for shutdown doesn't toggle recording first
    if (buttonMode.released() && modePressed) {
      modePressed = false;

      if (state == STATE_RECORDING) {
        stopRecording();
      } else {
        startRecording();
      }
    }

    if (modePressed && now - modePressStart > SHUTDOWN_HOLD) {
      modePressed = false;
      shutdown();
    }

    if (buttonLeft.pressed()) {
      lastActivity = now;

      if (state == STATE_RECORDING) {
        recordMarker(1);
      }
    }

    if (buttonRight.pressed()) {
      lastActivity = now;

      if (state == STATE_RECORDING) {
        recordMarker(2);
      }
    }
  }

  handleSerial();

  if (state == STATE_RECORDING) {
    if (now - lastFifoRead >= FIFO_INTERVAL) {
      lastFifoRead = now;
      readFifo();
    }

    if (magFound && now - lastMagRead >= MAG_INTERVAL) {
      lastMagRead = now;
      readMag();
    }

    recordInputs();

    if (storageFull) {
      stopRecording();
    }
  }

  if (now - lastStatus >= STATUS_INTERVAL) {
    lastStatus = now;
    renderStatus(now);
  }

  if (state != STATE_RECORDING && state != STATE_OFF && now - lastActivity > IDLE_TIMEOUT) {
    shutdown();
  }
}