/*
 * IMU BLE Mouse — Seeed XIAO nRF52840 Sense  (v2 — Full Featured)
 * ================================================================
 * Board BSP : Adafruit nRF52  (NOT Seeed mbed BSP)
 *   Board manager URL: https://adafruit.github.io/arduino-board-index/package_adafruit_index.json
 *   Select board: "Seeed XIAO nRF52840 Sense" (listed under Adafruit nRF52)
 *
 * Required Libraries:
 *   - Seeed Arduino LSM6DS3
 *   - Adafruit nRF52 BSP
 *
 * New in v2:
 *   1. BLE Configuration Service   — UUID 0x1234 with writable characteristics
 *   2. EEPROM persistence          — config saved to flash via InternalFileSystem
 *   3. BLE calibration trigger     — write 0x01 to CAL characteristic
 *   4. Motion scaling curve select — LINEAR / SQUARE / SQRT
 *   5. Factory Reset command       — write 0xFF to CMD characteristic
 *   6. Auto-recalibrate on idle    — after AUTO_RECAL_MINUTES minutes of stillness
 *   7. Axis flip flags             — flip X and/or Y via BLE config
 *
 * ── BLE Config Service layout (UUID 0x1234) ────────────────────
 *   Characteristic    | UUID   | Len | Description
 *   ──────────────────|────────|─────|──────────────────────────
 *   Sensitivity       | 0x1235 |  4  | float, cursor speed
 *   Dead Zone         | 0x1236 |  4  | float, noise floor rad/s
 *   Accel Strength    | 0x1237 |  4  | float, pointer accel
 *   Curve Select      | 0x1238 |  1  | 0=LINEAR 1=SQUARE 2=SQRT
 *   Axis Flip         | 0x1239 |  1  | bit0=flipX bit1=flipY
 *   Command           | 0x123A |  1  | 0x01=Calibrate 0xFF=FactoryReset
 */

#include <bluefruit.h>
#include <Adafruit_LittleFS.h>
#include <InternalFileSystem.h>
#include "LSM6DS3.h"
#include "Wire.h"

// ─── Debug ────────────────────────────────────────────────────────────────────
// #define DEBUG

// ─── BLE Standard Services ────────────────────────────────────────────────────
BLEDis         bledis;
BLEHidAdafruit blehid;
BLEBas         blebas;

// ─── BLE Config Service & Characteristics ────────────────────────────────────
BLEService     cfgService(0x1234);
BLECharacteristic cfgSensitivity (0x1235);
BLECharacteristic cfgDeadZone    (0x1236);
BLECharacteristic cfgAccelStr    (0x1237);
BLECharacteristic cfgCurve       (0x1238);
BLECharacteristic cfgAxisFlip    (0x1239);
BLECharacteristic cfgCommand     (0x123A);

// ─── IMU ──────────────────────────────────────────────────────────────────────
LSM6DS3 imu(I2C_MODE, 0x6A);

// ─── Pin Definitions ──────────────────────────────────────────────────────────
#define PIN_VBAT_ENABLE (14)
#define PIN_VBAT_READ   (32)
#define PIN_CHG         (17)

// ─── EEPROM / Persistence ─────────────────────────────────────────────────────
#define CONFIG_FILENAME   "/imu_mouse_cfg.bin"
#define CONFIG_MAGIC      0xDEAD1234UL

using namespace Adafruit_LittleFS_Namespace;
File cfgFile(InternalFS);

// ─── Motion Scaling Curves ────────────────────────────────────────────────────
enum CurveType : uint8_t {
  CURVE_LINEAR = 0,
  CURVE_SQUARE = 1,
  CURVE_SQRT   = 2
};

// ─── Config Struct (persisted) ────────────────────────────────────────────────
struct Config {
  uint32_t  magic;
  float     sensitivity;
  float     deadZone;
  float     accelStrength;
  CurveType curve;
  uint8_t   axisFlip;   // bit0=flipX, bit1=flipY
};

Config cfg;

// ─── Default Parameters ───────────────────────────────────────────────────────
const Config CFG_DEFAULTS = {
  CONFIG_MAGIC,
  600.0f,   // sensitivity
  0.060f,   // dead zone
  0.08f,    // accel strength
  CURVE_LINEAR,
  0x00      // no flips
};

// ─── Fixed Parameters ─────────────────────────────────────────────────────────
const float ALPHA            = 0.96f;
const int   LOOP_RATE_MS     = 10;
const int   BIAS_SAMPLES     = 200;
const int   IDLE_FRAMES      = 150;

// Auto-recalibrate: recalibrate after this many minutes of continuous idle
const unsigned long AUTO_RECAL_MINUTES = 5;
const unsigned long AUTO_RECAL_MS      = AUTO_RECAL_MINUTES * 60UL * 1000UL;

const unsigned long BATT_REPORT_MS  = 10000;
const unsigned long HEARTBEAT_MS    = 2000;
const int           HEARTBEAT_DUR   = 30;

const float BATT_FULL     = 4.20f;
const float BATT_EMPTY    = 3.00f;
const float BATT_CRITICAL = 3.10f;

// ─── State ────────────────────────────────────────────────────────────────────
float angleX = 0.0f, angleY = 0.0f;
float accumX = 0.0f, accumY = 0.0f;
float biasGX = 0.0f, biasGY = 0.0f, biasGZ = 0.0f;
int   idleFrames    = 0;
bool  pendingCal    = false;   // set by BLE write callback
bool  pendingReset  = false;   // set by BLE write callback

unsigned long lastTime        = 0;
unsigned long lastBattTime    = 0;
unsigned long lastHeartbeat   = 0;
unsigned long idleStartMs     = 0;  // when continuous idle began (0 = not idle)

// ─── EEPROM Helpers ───────────────────────────────────────────────────────────
void loadConfig() {
  InternalFS.begin();
  cfgFile.open(CONFIG_FILENAME, FILE_O_READ);
  if (cfgFile) {
    cfgFile.read(&cfg, sizeof(cfg));
    cfgFile.close();
    if (cfg.magic != CONFIG_MAGIC) {
      Serial.println("[CFG] Bad magic — using defaults");
      cfg = CFG_DEFAULTS;
    } else {
      Serial.println("[CFG] Loaded from flash");
    }
  } else {
    Serial.println("[CFG] No file — using defaults");
    cfg = CFG_DEFAULTS;
  }
}

void saveConfig() {
  InternalFS.remove(CONFIG_FILENAME);
  cfgFile.open(CONFIG_FILENAME, FILE_O_WRITE);
  if (cfgFile) {
    cfgFile.write((uint8_t*)&cfg, sizeof(cfg));
    cfgFile.close();
    Serial.println("[CFG] Saved to flash");
  } else {
    Serial.println("[CFG] ERROR: could not open file for write");
  }
}

void factoryReset() {
  Serial.println("[CFG] Factory reset!");
  cfg = CFG_DEFAULTS;
  saveConfig();
  // Push defaults back to BLE characteristics
  cfgSensitivity.write((uint8_t*)&cfg.sensitivity,    4);
  cfgDeadZone.write   ((uint8_t*)&cfg.deadZone,       4);
  cfgAccelStr.write   ((uint8_t*)&cfg.accelStrength,  4);
  cfgCurve.write      ((uint8_t*)&cfg.curve,          1);
  cfgAxisFlip.write   ((uint8_t*)&cfg.axisFlip,       1);
}

// ─── BLE Write Callbacks ──────────────────────────────────────────────────────
void onSensitivityWrite(uint16_t conn_hdl, BLECharacteristic* chr,
                        uint8_t* data, uint16_t len) {
  if (len == 4) { memcpy(&cfg.sensitivity, data, 4); saveConfig(); }
}
void onDeadZoneWrite(uint16_t conn_hdl, BLECharacteristic* chr,
                     uint8_t* data, uint16_t len) {
  if (len == 4) { memcpy(&cfg.deadZone, data, 4); saveConfig(); }
}
void onAccelStrWrite(uint16_t conn_hdl, BLECharacteristic* chr,
                     uint8_t* data, uint16_t len) {
  if (len == 4) { memcpy(&cfg.accelStrength, data, 4); saveConfig(); }
}
void onCurveWrite(uint16_t conn_hdl, BLECharacteristic* chr,
                  uint8_t* data, uint16_t len) {
  if (len == 1 && data[0] <= 2) {
    cfg.curve = (CurveType)data[0];
    saveConfig();
    Serial.print("[CFG] Curve -> "); Serial.println(cfg.curve);
  }
}
void onAxisFlipWrite(uint16_t conn_hdl, BLECharacteristic* chr,
                     uint8_t* data, uint16_t len) {
  if (len == 1) {
    cfg.axisFlip = data[0];
    saveConfig();
    Serial.print("[CFG] AxisFlip -> 0x"); Serial.println(cfg.axisFlip, HEX);
  }
}
void onCommandWrite(uint16_t conn_hdl, BLECharacteristic* chr,
                    uint8_t* data, uint16_t len) {
  if (len < 1) return;
  if (data[0] == 0x01) {
    pendingCal = true;
    Serial.println("[CMD] Calibration requested via BLE");
  } else if (data[0] == 0xFF) {
    pendingReset = true;
    Serial.println("[CMD] Factory reset requested via BLE");
  }
}

// ─── BLE Config Service Setup ─────────────────────────────────────────────────
void setupConfigService() {
  cfgService.begin();

  // Each characteristic: READ | WRITE, no response needed for writes
  auto props = CHR_PROPS_READ | CHR_PROPS_WRITE;

  cfgSensitivity.setProperties(props);
  cfgSensitivity.setPermission(SECMODE_OPEN, SECMODE_OPEN);
  cfgSensitivity.setFixedLen(4);
  cfgSensitivity.setWriteCallback(onSensitivityWrite);
  cfgSensitivity.begin();
  cfgSensitivity.write((uint8_t*)&cfg.sensitivity, 4);

  cfgDeadZone.setProperties(props);
  cfgDeadZone.setPermission(SECMODE_OPEN, SECMODE_OPEN);
  cfgDeadZone.setFixedLen(4);
  cfgDeadZone.setWriteCallback(onDeadZoneWrite);
  cfgDeadZone.begin();
  cfgDeadZone.write((uint8_t*)&cfg.deadZone, 4);

  cfgAccelStr.setProperties(props);
  cfgAccelStr.setPermission(SECMODE_OPEN, SECMODE_OPEN);
  cfgAccelStr.setFixedLen(4);
  cfgAccelStr.setWriteCallback(onAccelStrWrite);
  cfgAccelStr.begin();
  cfgAccelStr.write((uint8_t*)&cfg.accelStrength, 4);

  cfgCurve.setProperties(props);
  cfgCurve.setPermission(SECMODE_OPEN, SECMODE_OPEN);
  cfgCurve.setFixedLen(1);
  cfgCurve.setWriteCallback(onCurveWrite);
  cfgCurve.begin();
  cfgCurve.write((uint8_t*)&cfg.curve, 1);

  cfgAxisFlip.setProperties(props);
  cfgAxisFlip.setPermission(SECMODE_OPEN, SECMODE_OPEN);
  cfgAxisFlip.setFixedLen(1);
  cfgAxisFlip.setWriteCallback(onAxisFlipWrite);
  cfgAxisFlip.begin();
  cfgAxisFlip.write((uint8_t*)&cfg.axisFlip, 1);

  cfgCommand.setProperties(CHR_PROPS_WRITE);
  cfgCommand.setPermission(SECMODE_OPEN, SECMODE_OPEN);
  cfgCommand.setFixedLen(1);
  cfgCommand.setWriteCallback(onCommandWrite);
  cfgCommand.begin();
}

// ─── Battery ──────────────────────────────────────────────────────────────────
float readBatteryVoltage() {
  pinMode(PIN_VBAT_ENABLE, OUTPUT);
  digitalWrite(PIN_VBAT_ENABLE, LOW);
  delay(1);

  pinMode(PIN_VBAT_READ, INPUT);
  analogReference(AR_INTERNAL_3_0);
  analogReadResolution(12);
  for (int i = 0; i < 5; i++) { analogRead(PIN_VBAT_READ); delay(1); }

  int32_t raw = 0;
  for (int i = 0; i < 16; i++) raw += analogRead(PIN_VBAT_READ);
  raw /= 16;

  digitalWrite(PIN_VBAT_ENABLE, HIGH);
  analogReference(AR_DEFAULT);
  analogReadResolution(10);

  float v = (raw / 4096.0f) * 3.0f * 2.0f;

  Serial.print("[BATT DBG] raw="); Serial.print(raw);
  Serial.print(" ("); Serial.print(v, 3); Serial.print("V)");
  Serial.print("  CHG pin="); Serial.println(digitalRead(PIN_CHG));
  return v;
}

int batteryPercent(float v) {
  return (int) constrain((v - BATT_EMPTY) / (BATT_FULL - BATT_EMPTY) * 100.0f, 0, 100);
}

void updateBattery() {
  float v   = readBatteryVoltage();
  int   pct = batteryPercent(v);
  bool  charging = (digitalRead(PIN_CHG) == LOW);

  blebas.write(pct);

  Serial.print("[BATT] ");
  Serial.print(v, 2); Serial.print("V  ");
  Serial.print(pct);  Serial.print("%");
  if (charging)             Serial.print("  [CHARGING]");
  else if (pct >= 99)       Serial.print("  [FULL]");
  else if (v < BATT_CRITICAL) Serial.print("  [CRITICAL - CHARGE NOW]");
  else                      Serial.print("  [ON BATTERY]");
  Serial.println();

  if (!charging && v < BATT_CRITICAL) {
    pinMode(LED_RED, OUTPUT);
    for (int i = 0; i < 6; i++) {
      digitalWrite(LED_RED, LOW);  delay(80);
      digitalWrite(LED_RED, HIGH); delay(80);
    }
  }
}

// ─── Gyro Calibration ─────────────────────────────────────────────────────────
void calibrateGyroBias() {
  Serial.println("[CAL] Hold still — calibrating gyro bias...");
  pinMode(LED_BLUE, OUTPUT);

  double sumX = 0, sumY = 0, sumZ = 0;
  for (int i = 0; i < BIAS_SAMPLES; i++) {
    sumX += imu.readFloatGyroX();
    sumY += imu.readFloatGyroY();
    sumZ += imu.readFloatGyroZ();
    digitalWrite(LED_BLUE, (i % 20 < 10));
    delay(5);
  }

  biasGX = (float)(sumX / BIAS_SAMPLES);
  biasGY = (float)(sumY / BIAS_SAMPLES);
  biasGZ = (float)(sumZ / BIAS_SAMPLES);

  // Reset angle state to avoid a jump after recal
  angleX = 0.0f;
  angleY = 0.0f;
  accumX = 0.0f;
  accumY = 0.0f;

  digitalWrite(LED_BLUE, HIGH);
  Serial.print("[CAL] Done. Bias — gx:"); Serial.print(biasGX, 4);
  Serial.print(" gy:"); Serial.print(biasGY, 4);
  Serial.print(" gz:"); Serial.println(biasGZ, 4);
}

// ─── Motion Scaling ───────────────────────────────────────────────────────────
float applyAcceleration(float delta) {
  // Pointer acceleration on top of curve
  return delta * (1.0f + fabsf(delta) * cfg.accelStrength);
}

float applyCurve(float v) {
  switch (cfg.curve) {
    case CURVE_SQUARE:
      return (v >= 0.0f ? 1.0f : -1.0f) * v * v;
    case CURVE_SQRT:
      return (v >= 0.0f ? 1.0f : -1.0f) * sqrtf(fabsf(v));
    case CURVE_LINEAR:
    default:
      return v;
  }
}

// ─── BLE Advertising ──────────────────────────────────────────────────────────
void startAdvertising() {
  Bluefruit.Advertising.addFlags(BLE_GAP_ADV_FLAGS_LE_ONLY_GENERAL_DISC_MODE);
  Bluefruit.Advertising.addTxPower();
  Bluefruit.Advertising.addAppearance(BLE_APPEARANCE_HID_MOUSE);
  Bluefruit.Advertising.addService(blehid);
  Bluefruit.Advertising.addService(blebas);
  Bluefruit.Advertising.addName();
  Bluefruit.Advertising.restartOnDisconnect(true);
  Bluefruit.Advertising.setInterval(32, 244);
  Bluefruit.Advertising.setFastTimeout(30);
  Bluefruit.Advertising.start(0);
}

// ─────────────────────────────────────────────────────────────────────────────
void setup() {
  Serial.begin(115200);
  while (!Serial) delay(10);

  pinMode(PIN_CHG,  INPUT_PULLUP);
  pinMode(LED_RED,  OUTPUT); digitalWrite(LED_RED,  HIGH);
  pinMode(LED_BLUE, OUTPUT); digitalWrite(LED_BLUE, HIGH);

  // ── Load persisted config ─────────────────────────────────────────────────
  loadConfig();

  // ── IMU ───────────────────────────────────────────────────────────────────
  if (imu.begin() != 0) {
    Serial.println("[ERROR] IMU init failed.");
    while (1) { digitalWrite(LED_RED, !digitalRead(LED_RED)); delay(100); }
  }
  Serial.println("[OK] IMU initialised");

  updateBattery();
  calibrateGyroBias();

  // ── BLE ───────────────────────────────────────────────────────────────────
  Bluefruit.begin();
  Bluefruit.setTxPower(4);
  Bluefruit.setName("IMU Mouse");
  Bluefruit.Periph.setConnInterval(6, 12);

  bledis.setManufacturer("Seeed Studio");
  bledis.setModel("XIAO nRF52840 Sense");
  bledis.begin();

  blehid.begin();
  blebas.begin();
  blebas.write(100);

  // Config service must begin AFTER Bluefruit.begin()
  setupConfigService();

  startAdvertising();
  Serial.println("[OK] BLE advertising — pair 'IMU Mouse' on your host");
  Serial.println("     Config service UUID 0x1234 available for tuning");

  lastTime      = millis();
  lastBattTime  = millis();
  lastHeartbeat = millis();
  idleStartMs   = 0;
}

// ─────────────────────────────────────────────────────────────────────────────
void loop() {
  unsigned long now = millis();

  // ── Deferred commands (from BLE callbacks, safe to run on main thread) ────
  if (pendingCal) {
    pendingCal = false;
    calibrateGyroBias();
  }
  if (pendingReset) {
    pendingReset = false;
    factoryReset();
  }

  // ── Heartbeat LED ─────────────────────────────────────────────────────────
  if (now - lastHeartbeat >= HEARTBEAT_MS) {
    lastHeartbeat = now;
    int led = Bluefruit.connected() ? LED_BLUE : LED_RED;
    digitalWrite(led, LOW);
    delay(HEARTBEAT_DUR);
    digitalWrite(led, HIGH);
  }

  // ── Battery ───────────────────────────────────────────────────────────────
  if (now - lastBattTime >= BATT_REPORT_MS) {
    lastBattTime = now;
    updateBattery();
  }

  // ── IMU rate limit ────────────────────────────────────────────────────────
  if (now - lastTime < (unsigned long)LOOP_RATE_MS) return;

  float dt = (now - lastTime) / 1000.0f;
  lastTime = now;
  if (dt <= 0.0f || dt > 0.5f) return;

  // ── Read IMU ──────────────────────────────────────────────────────────────
  float gx = (imu.readFloatGyroX() - biasGX) * (PI / 180.0f);
  float gy = (imu.readFloatGyroY() - biasGY) * (PI / 180.0f);
  float gz = (imu.readFloatGyroZ() - biasGZ) * (PI / 180.0f);

  float ax = imu.readFloatAccelX();
  float ay = imu.readFloatAccelY();
  float az = imu.readFloatAccelZ();

  // ── Complementary filter ──────────────────────────────────────────────────
  angleX = ALPHA * (angleX + gx * dt) + (1.0f - ALPHA) * atan2f(ax, sqrtf(ay*ay + az*az));
  angleY = ALPHA * (angleY + gy * dt) + (1.0f - ALPHA) * atan2f(ay, sqrtf(ax*ax + az*az));

  // ── Dead zone ─────────────────────────────────────────────────────────────
  float filteredGy = (fabsf(gy) > cfg.deadZone) ? gy : 0.0f;
  float filteredGz = (fabsf(gz) > cfg.deadZone) ? gz : 0.0f;

  // ── Idle detection + auto-recalibrate ─────────────────────────────────────
  bool moving = (filteredGy != 0.0f || filteredGz != 0.0f);
  if (moving) {
    idleFrames  = 0;
    idleStartMs = 0;
  } else {
    idleFrames++;
    if (idleStartMs == 0) idleStartMs = now;   // mark start of idle streak
  }

  bool idle = (idleFrames >= IDLE_FRAMES);

  // Auto-recalibrate after AUTO_RECAL_MS of continuous stillness
  if (idle && idleStartMs != 0 && (now - idleStartMs >= AUTO_RECAL_MS)) {
    Serial.println("[AUTO-CAL] Long idle detected — recalibrating...");
    idleStartMs = 0;   // reset so we don't retrigger immediately
    calibrateGyroBias();
    return;
  }

  if (idle) {
    accumX = 0.0f;
    accumY = 0.0f;
#ifdef DEBUG
    Serial.println("[IDLE]");
#endif
    return;
  }

  // ── Delta + curve + acceleration + sub-pixel accumulation ─────────────────
  float rawX = -filteredGz * cfg.sensitivity * dt;
  float rawY =  filteredGy * cfg.sensitivity * dt;

  rawX = applyCurve(rawX);
  rawY = applyCurve(rawY);

  rawX = applyAcceleration(rawX);
  rawY = applyAcceleration(rawY);

  // ── Axis flip ─────────────────────────────────────────────────────────────
  if (cfg.axisFlip & 0x01) rawX = -rawX;   // flip X
  if (cfg.axisFlip & 0x02) rawY = -rawY;   // flip Y

  accumX += rawX;
  accumY += rawY;

  int8_t moveX = (int8_t) constrain((int)accumX, -127, 127);
  int8_t moveY = (int8_t) constrain((int)accumY, -127, 127);

  accumX -= moveX;
  accumY -= moveY;

  // ── BLE HID ───────────────────────────────────────────────────────────────
  if (Bluefruit.connected() && (moveX != 0 || moveY != 0)) {
    blehid.mouseMove(moveX, moveY);
  }

#ifdef DEBUG
  Serial.print("gy:"); Serial.print(gy, 3);
  Serial.print(" gz:"); Serial.print(gz, 3);
  Serial.print(" | mx:"); Serial.print(moveX);
  Serial.print(" my:"); Serial.println(moveY);
#endif
}