/*
 * IMU BLE Mouse - Seeed XIAO nRF52840 Sense  (v3.4)
 * ================================================================
 * MINIMUM (just working mouse, no BLE config):
 *   leave only FEATURE_BATTERY_MONITOR + FEATURE_BOOT_LOOP_DETECT
 *
 * RECOMMENDED first test:
 *   enable FEATURE_CONFIG_SERVICE, keep TAP + STREAM + TELEMETRY off
 *
 * Dependencies:
 *   FEATURE_TELEMETRY  requires FEATURE_CONFIG_SERVICE
 *   FEATURE_IMU_STREAM requires FEATURE_CONFIG_SERVICE
 * ================================================================
 */

#include "config.h"
#include "imu.h"
#include "ble_config.h"
#include "battery.h"
#include "tap.h"
#include "buttons.h"
#include <bluefruit.h>
#include <Adafruit_LittleFS.h>
#include <InternalFileSystem.h>
#include "Wire.h"
#include "sleep.h"

// Boot-loop detection 
#ifdef FEATURE_BOOT_LOOP_DETECT
  static uint32_t __attribute__((section(".noinit"))) bootCount;
  static uint32_t __attribute__((section(".noinit"))) bootMagic;
#endif

// BLE Standard Services 
BLEDis         bledis;
BLEHidAdafruit blehid;
#ifdef FEATURE_BATTERY_MONITOR
  BLEBas blebas;
#endif

// Persistence 
using namespace Adafruit_LittleFS_Namespace;
File cfgFile(InternalFS);

// Config definitions 
Config cfg;
const Config CFG_DEFAULTS = {
  CONFIG_MAGIC, 600.0f, 0.060f, 0.08f, CURVE_LINEAR, 0x00, CHARGE_SLOW,
  /*tapThreshold=*/12, /*tapAction=*/TAP_ACTION_LEFT, /*tapKey=*/0, /*tapMod=*/0,
  /*jerkThreshold=*/2000.0f, /*tapFreezeEnabled=*/1, /*featureFlags=*/FLAG_ALL_DEFAULT
};

// Telemetry definition 
#ifdef FEATURE_TELEMETRY
TelemetryPacket telem = {};
#endif

// Tuning constants 
const float ALPHA              = 0.96f;
const int   LOOP_RATE_MS       = 10;
const float SMOOTH_LOW_RPS     = 0.15f;  // below this → heavy EMA smoothing (~8°/s)
const float SMOOTH_HIGH_RPS    = 0.50f;  // above this → no smoothing (~29°/s)
const int   BIAS_SAMPLES       = 200;
const int   IDLE_FRAMES        = 150;
const unsigned long BATT_REPORT_MS       = 20000;
const unsigned long TELEMETRY_MS         = 1000;
const unsigned long HEARTBEAT_MS         = 10000;
const int           HEARTBEAT_DUR        = 30;
const unsigned long BOOT_SAFE_MS         = 5000;
#ifdef FEATURE_IMU_STREAM
  const unsigned long IMU_STREAM_RATE_MS = 20;
#endif
const float BATT_FULL                    = 4.20f;
const float BATT_EMPTY                   = 3.00f;
const float BATT_CRITICAL                = 3.10f;
#ifdef FEATURE_TAP_DETECTION
  const unsigned long CLICK_HOLD_MS = 60;
#endif
#ifdef FEATURE_TEMP_COMPENSATION
  const float TEMP_COMP_COEFF_DPS_C = 0.004f;
#endif
#ifdef FEATURE_AUTO_RECAL
  const unsigned long AUTO_RECAL_MS = 5UL * 60UL * 1000UL;
#endif

// Global state definitions 
float angleX = 0, angleY = 0;
float accumX = 0, accumY = 0;
// Low-pass filtered gravity estimate in device frame (for roll-independent axis projection)
float gravX = 0, gravY = 0, gravZ = 1.0f;
float biasGX = 0, biasGY = 0, biasGZ = 0;
float calTempC    = 25.0f;
float cachedTempC = 25.0f;

#ifdef FEATURE_TAP_DETECTION
  bool    clickButtonDown = false;
  uint8_t clickButton     = 0;
  unsigned long clickDownMs= 0;
  uint32_t statLeftClicks  = 0;
  uint32_t statRightClicks = 0;
#endif

#ifdef FEATURE_IMU_STREAM
  bool imuStreamEnabled = false;
  uint32_t streamNotifyFails = 0;
  uint32_t streamNotifyOk   = 0;
  unsigned long lastStreamDiag = 0;
  // Back-off state: after STREAM_BACKOFF_THRESH consecutive fails, skip notifies
  // for STREAM_BACKOFF_MS to let the SoftDevice HVN TX semaphore drain.
  // Without this, every notify() blocks for BLE_GENERIC_TIMEOUT (100ms).
  uint8_t       streamConsecFails = 0;
  unsigned long streamBackoffUntil = 0;
  const uint8_t        STREAM_BACKOFF_THRESH = 2;    // fails before backing off
  const unsigned long  STREAM_BACKOFF_MS     = 500;  // cooldown window
#endif

uint32_t loopStalls = 0;          // loop iterations where dt > 20ms (behind schedule)
bool pendingCal   = false;
bool pendingReset = false;
#ifdef FEATURE_OTA
  bool pendingOTA = false;
#endif

// Jerk-based shock detection - freeze cursor during tap impacts, doesn't work well yet!
unsigned long shockFreezeUntil = 0;
float prevAx = 0, prevAy = 0, prevAz = 0;     // previous frame's accel for Δa
const unsigned long SHOCK_FREEZE_MS = 80;       // hold freeze after last spike

ChargeStatus lastChargeStatus = CHGSTAT_DISCHARGING;

int           idleFrames  = 0;
unsigned long idleStartMs = 0;
unsigned long lastTime      = 0;
unsigned long lastBattTime  = 0;
unsigned long lastHeartbeat = 0;
unsigned long lastTelemetry = 0;
unsigned long bootStartMs   = 0;
#ifdef FEATURE_IMU_STREAM
  unsigned long lastImuStream = 0;
#endif

#ifdef FEATURE_TELEMETRY
  uint16_t statRecalCount = 0;
  float    statBiasRms    = 0.0f;
#endif

bool safeMode         = false;
bool bootCountCleared = false;

// Advertising 
static 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);
  #ifdef FEATURE_BATTERY_MONITOR
    Bluefruit.Advertising.addService(blebas);
  #endif
  Bluefruit.Advertising.addName();
  Bluefruit.Advertising.restartOnDisconnect(true);
  Bluefruit.Advertising.setInterval(32, 244);
  Bluefruit.Advertising.setFastTimeout(30);
  Bluefruit.Advertising.start(0);
}

// Setup 
void setup() {
  Serial.begin(115200);
  unsigned long serialWait = millis();
  while (!Serial && (millis() - serialWait < 2000)) { delay(10); }

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

  // Boot-loop detection 
  #ifdef FEATURE_BOOT_LOOP_DETECT
    if (bootMagic != 0xCAFEBABE) { bootMagic = 0xCAFEBABE; bootCount = 0; }
    bootCount++;
    Serial.print("[BOOT] count="); Serial.println(bootCount);
    if (bootCount >= 3) {
      bootCount = 0; safeMode = true;
      Serial.println("[BOOT] Boot loop - safe mode (no config service)");
      InternalFS.begin(); InternalFS.remove(CONFIG_FILENAME);
      for (int i=0; i<3; i++) { digitalWrite(LED_RED,LOW); delay(150); digitalWrite(LED_RED,HIGH); delay(150); }  // fault: red
    }
  #endif

  loadConfig();
  applyChargeMode(cfg.chargeMode);

  // configAttrTableSize MUST be called before Bluefruit.begin().
  // Value must be >= what HID+DIS+BAS+custom services actually need.
  // Too small crashes just as hard as too large. Floor is 1536.
  Serial.print("[BLE] ATT table: "); Serial.print(ATT_TABLE_SIZE); Serial.println(" bytes");
  Bluefruit.configAttrTableSize(ATT_TABLE_SIZE);
  Bluefruit.configPrphBandwidth(BANDWIDTH_MAX);  // Raise MTU so 24-byte telemetry fits in one notify

  Bluefruit.begin(1, 0);
  Bluefruit.setTxPower(4);
  Bluefruit.setName(safeMode ? "IMU Mouse (safe)" : "IMU Mouse");
  Bluefruit.Periph.setConnInterval(16, 32);  // 20-40ms - wider interval reduces SoftDevice TX stalls

  Wire1.begin();  // LSM6DS3 is on internal I2C bus (Wire1), must init before imu.begin()
  if (imu.begin() != 0) {
    Serial.println("[ERROR] IMU init failed");
    while(1) { digitalWrite(LED_RED, !digitalRead(LED_RED)); delay(100); }  // fault: red rapid blink
  }
  Serial.println("[OK] IMU ready");

  #ifdef FEATURE_TAP_DETECTION
    if (cfg.featureFlags & FLAG_TAP_ENABLED) setupTapDetection();
  #endif

  #ifdef FEATURE_PHYSICAL_BUTTONS
    setupPhysicalButtons();
  #endif

  cachedTempC = readIMUTemp();

  #ifdef FEATURE_BATTERY_MONITOR
    initBatteryADC();
    updateBattery();
  #endif

  calibrateGyroBias();
  // Seed previous-accel for jerk detection so first frame doesn't spike
  prevAx = imu.readFloatAccelX(); prevAy = imu.readFloatAccelY(); prevAz = imu.readFloatAccelZ();

  sleepManagerInit();

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

  blehid.begin();

  #ifdef FEATURE_BATTERY_MONITOR
    blebas.begin(); blebas.write(batteryPercent(readBatteryVoltage()));
  #endif

  #ifdef FEATURE_CONFIG_SERVICE
    if (!safeMode) {
      setupConfigService();
      Serial.println("[OK] Config service started");
    } else {
      Serial.println("[SAFE] Config service skipped");
    }
  #endif

  startAdvertising();
  Serial.print("[OK] Advertising - features:");
  #ifdef FEATURE_CONFIG_SERVICE
    Serial.print(" CFG");
  #endif
  #ifdef FEATURE_TELEMETRY
    Serial.print(" TELEM");
  #endif
  #ifdef FEATURE_IMU_STREAM
    Serial.print(" STREAM");
  #endif
  #ifdef FEATURE_TAP_DETECTION
    Serial.print(" TAP");
  #endif
  #ifdef FEATURE_TEMP_COMPENSATION
    Serial.print(" TEMPCOMP");
  #endif
  #ifdef FEATURE_AUTO_RECAL
    Serial.print(" AUTORECAL");
  #endif
  #ifdef FEATURE_BATTERY_MONITOR
    Serial.print(" BATT");
  #endif
  #ifdef FEATURE_BOOT_LOOP_DETECT
    Serial.print(" BOOTDET");
  #endif
  #ifdef FEATURE_PHYSICAL_BUTTONS
    Serial.print(" PHYSBTN");
  #endif
  Serial.print(" SLEEP");
  Serial.println();

  bootStartMs = millis();
  lastTime = lastBattTime = lastHeartbeat = lastTelemetry = millis();
}

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

  // Clear boot counter after BOOT_SAFE_MS of stable running
  #ifdef FEATURE_BOOT_LOOP_DETECT
    if (!bootCountCleared && (now - bootStartMs >= BOOT_SAFE_MS)) {
      bootCount = 0; bootCountCleared = true;
      Serial.println("[BOOT] Stable - counter cleared");
    }
  #endif

  // Serial commands: 'c' = calibrate, 'r' = factory reset
  while (Serial.available()) {
    char cmd = Serial.read();
    if (cmd == 'c') { Serial.println("[SERIAL] Calibrate"); pendingCal   = true; }
    if (cmd == 'r') { Serial.println("[SERIAL] Reset");     pendingReset = true; }
    #ifdef FEATURE_OTA
      if (cmd == 'o') { Serial.println("[SERIAL] OTA DFU"); pendingOTA = true; }
    #endif
  }

  if (pendingCal)   { pendingCal = false; calibrateGyroBias(); prevAx = imu.readFloatAccelX(); prevAy = imu.readFloatAccelY(); prevAz = imu.readFloatAccelZ(); }
  if (pendingReset) { pendingReset = false; factoryReset(); }
  #ifdef FEATURE_OTA
    if (pendingOTA) {
      pendingOTA = false;
      Serial.println("[OTA] Disconnecting BLE and entering bootloader DFU mode...");
      Serial.flush();
      // Gracefully close the BLE connection first so the host can detect the
      // disconnect and be ready to see DfuTarg advertise after the reboot.
      if (Bluefruit.connected()) {
        Bluefruit.disconnect(0);
        delay(300);
      }
      delay(200);
      enterOTADfu();  // Adafruit nRF52 core: sets GPREGRET correctly and resets into bootloader OTA mode
    }
  #endif

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

  #ifdef FEATURE_BATTERY_MONITOR
    if (now - lastBattTime >= BATT_REPORT_MS) { lastBattTime = now; updateBattery(); }
  #endif

  #ifdef FEATURE_TAP_DETECTION
    if (cfg.featureFlags & FLAG_TAP_ENABLED) processTaps(now);
  #endif

  #ifdef FEATURE_PHYSICAL_BUTTONS
    processPhysicalButtons();
  #endif
  {
    bool idle_for_sleep = (sleepStage == SLEEP_IMU_LP) ? true
                          : (idleFrames >= IDLE_FRAMES);
    if (sleepManagerUpdate(now, idle_for_sleep, Bluefruit.connected())) return;
  }

  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;
  // Threshold 50ms: intentional heartbeat blink (30ms) won't false-trigger;
  // real SoftDevice stalls (100ms+) and unexpected delays still get flagged.
  if (dt > 0.050f) { loopStalls++; Serial.print("[STALL] dt="); Serial.print(dt*1000.f,1); Serial.print("ms stalls="); Serial.println(loopStalls); }

  cachedTempC = readIMUTemp();

  #ifdef FEATURE_TELEMETRY
    if (!safeMode && (now - lastTelemetry >= TELEMETRY_MS)) {
      lastTelemetry = now; pushTelemetry(now);
    }
  #endif

  // Gyro reads with optional temperature compensation.
  float correction = 0.0f;
  #ifdef FEATURE_TEMP_COMPENSATION
    if (cfg.featureFlags & FLAG_TEMP_COMP_ENABLED)
      correction = TEMP_COMP_COEFF_DPS_C * (cachedTempC - calTempC);
  #endif
  float gx = (imu.readFloatGyroX() - biasGX - correction) * (PI/180.0f);
  float gy = (imu.readFloatGyroY() - biasGY - correction) * (PI/180.0f);
  float gz = (imu.readFloatGyroZ() - biasGZ - correction) * (PI/180.0f);

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

  // Jerk-based shock detection - freeze cursor during tap impacts, doesn't work well yet
  float jx = (ax - prevAx) / dt, jy = (ay - prevAy) / dt, jz = (az - prevAz) / dt;
  float jerkSq = jx*jx + jy*jy + jz*jz;
  prevAx = ax; prevAy = ay; prevAz = az;
  bool shocked = cfg.tapFreezeEnabled && ((jerkSq > cfg.jerkThreshold) || (now < shockFreezeUntil));
  if (cfg.tapFreezeEnabled && jerkSq > cfg.jerkThreshold) shockFreezeUntil = now + SHOCK_FREEZE_MS;

  // Complementary filter
  if (shocked) {
    angleX += gx * dt;
    angleY += gz * dt;
  } else {
    angleX = ALPHA*(angleX + gx*dt) + (1.0f - ALPHA)*atan2f(ax, sqrtf(ay*ay + az*az));
    angleY = ALPHA*(angleY + gz*dt) + (1.0f - ALPHA)*atan2f(ay, sqrtf(ax*ax + az*az));
  }

  // Gravity-based axis decomposition
  const float GRAV_LP = 0.05f;
  if (!shocked) {
    gravX += GRAV_LP * (ax - gravX);
    gravY += GRAV_LP * (ay - gravY);
    gravZ += GRAV_LP * (az - gravZ);
  }

  float gN = sqrtf(gravX*gravX + gravY*gravY + gravZ*gravZ);
  if (gN < 0.3f) gN = 1.0f;
  float gnx = gravX/gN, gny = gravY/gN, gnz = gravZ/gN;

  float ry = -gnz, rz = gny;
  float rN = sqrtf(ry*ry + rz*rz);
  if (rN < 0.01f) { ry = -1.0f; rz = 0.0f; rN = 1.0f; }
  ry /= rN; rz /= rN;

  // Yaw  (cursor X) = angular velocity component around gravity (vertical)
  // Pitch (cursor Y) = angular velocity component around screen-right
  float yawRate   = gx*gnx + gy*gny + gz*gnz;
  float pitchRate = -(gy*ry + gz*rz);

  // Projected rates amplify residual gyro bias (especially GY drift on pitch axis).
  float fYaw   = (fabsf(yawRate)   > cfg.deadZone)       ? yawRate   : 0.0f;
  float fPitch = (fabsf(pitchRate) > cfg.deadZone * 3.0f) ? pitchRate : 0.0f;

  #ifdef DEBUG
  { static unsigned long lastDiag = 0;
    if (now - lastDiag >= 500) { lastDiag = now;
      Serial.print("[PROJ] grav="); Serial.print(gnx,2); Serial.print(","); Serial.print(gny,2); Serial.print(","); Serial.print(gnz,2);
      Serial.print(" R="); Serial.print(ry,2); Serial.print(","); Serial.print(rz,2);
      Serial.print(" gyro="); Serial.print(gx,2); Serial.print(","); Serial.print(gy,2); Serial.print(","); Serial.print(gz,2);
      Serial.print(" yaw="); Serial.print(yawRate,3); Serial.print(" pitch="); Serial.println(pitchRate,3);
    }
  }
  #endif

  bool moving = (fPitch != 0.0f || fYaw != 0.0f);
  if  (moving) { idleFrames = 0; idleStartMs = 0; }
  else         { idleFrames++; if (idleStartMs == 0) idleStartMs = now; }
  bool idle = (idleFrames >= IDLE_FRAMES);

  #ifdef FEATURE_AUTO_RECAL
    if ((cfg.featureFlags & FLAG_AUTO_RECAL_ENABLED) && idle && idleStartMs != 0 && (now - idleStartMs >= AUTO_RECAL_MS)) {
      Serial.println("[AUTO-CAL] Long idle - recalibrating...");
      idleStartMs = 0; calibrateGyroBias(); prevAx = imu.readFloatAccelX(); prevAy = imu.readFloatAccelY(); prevAz = imu.readFloatAccelZ(); return;
    }
  #endif

  int8_t  moveX = 0, moveY = 0;
  uint8_t flags = 0;

  static float smoothX = 0.0f, smoothY = 0.0f;

  if (shocked) {
    // Shock freeze - discard accumulated sub-pixel motion and suppress output
    smoothX = smoothY = 0.0f;
    accumX = accumY = 0.0f;
    flags |= 0x08;  // bit3 = shock freeze active
  } else if (idle) {
    smoothX = smoothY = 0.0f;
    accumX = accumY = 0.0f;
    flags |= 0x01;
  } else {
    float rawX = applyAcceleration(applyCurve(-fYaw   * cfg.sensitivity * dt));
    float rawY = applyAcceleration(applyCurve(-fPitch * cfg.sensitivity * dt));
    if (cfg.axisFlip & 0x01) rawX = -rawX;
    if (cfg.axisFlip & 0x02) rawY = -rawY;
    // Tiered velocity smoothing: heavy EMA when nearly still, none when fast.
    // Thresholds are in rad/s (angular rate), independent of sensitivity setting.
    float speed = sqrtf(fYaw*fYaw + fPitch*fPitch);
    float alpha = (speed < SMOOTH_LOW_RPS)  ? 0.25f :
                  (speed < SMOOTH_HIGH_RPS) ? 0.65f : 1.00f;
    smoothX = smoothX * (1.0f - alpha) + rawX * alpha;
    smoothY = smoothY * (1.0f - alpha) + rawY * alpha;
    accumX += smoothX; accumY += smoothY;
    moveX = (int8_t)constrain((int)accumX, -127, 127);
    moveY = (int8_t)constrain((int)accumY, -127, 127);
    accumX -= moveX; accumY -= moveY;
    if (Bluefruit.connected() && (moveX != 0 || moveY != 0)) blehid.mouseMove(moveX, moveY);
  }

  #ifdef FEATURE_IMU_STREAM
    if (!safeMode && imuStreamEnabled && Bluefruit.connected()
        && (now - lastImuStream >= IMU_STREAM_RATE_MS)) {
      lastImuStream = now;

      if (now < streamBackoffUntil) {
        // Backing off - host TX buffer congested, skip to avoid 100ms block
      } else {
        ImuPacket pkt;
        pkt.gyroX_mDPS = (int16_t)constrain(gx*(180.f/PI)*1000.f, -32000, 32000);
        pkt.gyroZ_mDPS = (int16_t)constrain(gz*(180.f/PI)*1000.f, -32000, 32000);
        pkt.accelX_mg  = (int16_t)constrain(ax*1000.f, -32000, 32000);
        pkt.accelY_mg  = (int16_t)constrain(ay*1000.f, -32000, 32000);
        pkt.accelZ_mg  = (int16_t)constrain(az*1000.f, -32000, 32000);
        pkt.moveX      = moveX;
        pkt.moveY      = moveY;
        pkt.flags      = flags;
        pkt._pad       = 0;
        if (cfgImuStream.notify((uint8_t*)&pkt, sizeof(pkt))) {
          streamNotifyOk++;
          streamConsecFails = 0;
        } else {
          streamNotifyFails++;
          streamConsecFails++;
          if (streamConsecFails >= STREAM_BACKOFF_THRESH) {
            streamBackoffUntil = now + STREAM_BACKOFF_MS;
            streamConsecFails  = 0;
            Serial.print("[STREAM] TX congested - backing off ");
            Serial.print(STREAM_BACKOFF_MS); Serial.println("ms");
          }
        }
      }

      // Periodic stream health report every 10 seconds
      if (now - lastStreamDiag >= 10000) {
        lastStreamDiag = now;
        Serial.print("[STREAM] ok="); Serial.print(streamNotifyOk);
        Serial.print(" fail="); Serial.print(streamNotifyFails);
        Serial.print(" rate="); Serial.print((streamNotifyOk * 1000UL) / 10000); Serial.println("pkt/s");
        streamNotifyOk = 0; streamNotifyFails = 0;
      }
    }
  #endif
}