/*
 * IMU BLE Mouse — Seeed XIAO nRF52840 Sense  (v3.4)
 * ================================================================
 * Feature flags — comment out any line to disable that feature.
 * ATT table size is computed automatically from enabled features.
 * Start with minimal flags to isolate the SoftDevice RAM issue,
 * then re-enable one at a time.
 *
 * 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
 *
 * ── Feature flag index ───────────────────────────────────────────
 *   FEATURE_CONFIG_SERVICE    Custom GATT service (ConfigBlob + Command)
 *   FEATURE_TELEMETRY         +24-byte notify characteristic, 1 Hz
 *   FEATURE_IMU_STREAM        +14-byte notify characteristic, ~100 Hz
 *   FEATURE_TAP_DETECTION     LSM6DS3 hardware tap engine → L/R clicks
 *   FEATURE_TEMP_COMPENSATION Gyro drift correction by temperature delta
 *   FEATURE_AUTO_RECAL        Recalibrate after AUTO_RECAL_MS idle
 *   FEATURE_BATTERY_MONITOR   ADC battery read + BLE Battery Service
 *   FEATURE_BOOT_LOOP_DETECT  .noinit crash counter → safe mode
 *
 * 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 <bluefruit.h>
#include <Adafruit_LittleFS.h>
#include <InternalFileSystem.h>
#include "Wire.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
};

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

// ─── Tuning constants ─────────────────────────────────────────────────────────
const float ALPHA              = 0.96f;
const int   LOOP_RATE_MS       = 10;
const int   BIAS_SAMPLES       = 200;
const int   IDLE_FRAMES        = 150;
const unsigned long BATT_REPORT_MS       = 10000;
const unsigned long TELEMETRY_MS         = 1000;
const unsigned long HEARTBEAT_MS         = 2000;
const int           HEARTBEAT_DUR        = 30;
const unsigned long BOOT_SAFE_MS         = 5000;
#ifdef FEATURE_IMU_STREAM
  const unsigned long IMU_STREAM_RATE_MS = 50;
#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;
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;
#endif

bool pendingCal   = false;
bool pendingReset = false;

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_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); }
    }
  #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(12, 24);  // 15-30ms — less aggressive, prevents stream disconnect

  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); }
  }
  Serial.println("[OK] IMU ready");

  #ifdef FEATURE_TAP_DETECTION
    setupTapDetection();
  #endif

  cachedTempC = readIMUTemp();

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

  calibrateGyroBias();

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

  blehid.begin();

  #ifdef FEATURE_BATTERY_MONITOR
    blebas.begin(); blebas.write(100);
  #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
  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

  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);
  }

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

  #ifdef FEATURE_TAP_DETECTION
    processTaps(now);
  #endif

  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;

  cachedTempC = readIMUTemp();

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

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

  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));

  float fGy = (fabsf(gy) > cfg.deadZone) ? gy : 0.0f;
  float fGz = (fabsf(gz) > cfg.deadZone) ? gz : 0.0f;

  bool moving = (fGy != 0.0f || fGz != 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 (idle && idleStartMs != 0 && (now - idleStartMs >= AUTO_RECAL_MS)) {
      Serial.println("[AUTO-CAL] Long idle — recalibrating...");
      idleStartMs = 0; calibrateGyroBias(); return;
    }
  #endif

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

  if (idle) {
    accumX = accumY = 0.0f;
    flags |= 0x01;
  } else {
    float rawX = applyAcceleration(applyCurve(-fGz * cfg.sensitivity * dt));
    float rawY = applyAcceleration(applyCurve(-fGy * cfg.sensitivity * dt));
    if (cfg.axisFlip & 0x01) rawX = -rawX;
    if (cfg.axisFlip & 0x02) rawY = -rawY;
    accumX += rawX; accumY += rawY;
    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;
      ImuPacket pkt;
      pkt.gyroY_mDPS = (int16_t)constrain(gy*(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;
      cfgImuStream.notify((uint8_t*)&pkt, sizeof(pkt));
    }
  #endif

  #ifdef DEBUG
    Serial.print("T="); Serial.print(cachedTempC,1);
    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
}