/*
 * sleep.cpp — IMU Mouse low-power manager  (LSM6DS3TR-C + nRF52840)
 * See sleep.h for full documentation.
 */

#include "sleep.h"
#include "imu.h"       // config.h REG_* macros
#include <bluefruit.h> // sd_app_evt_wait()
#include "Wire.h"

// Registers not in config.h 
#define SLP_CTRL2_G      0x11
#define SLP_CTRL6_C      0x15
#define SLP_CTRL7_G      0x16
#define SLP_WAKE_UP_DUR  0x5C
#define SLP_WAKE_UP_SRC  0x1B

// LSM6DS3 I2C address — SA0 tied LOW on XIAO Sense → 0x6A (NOT 0x6B)
#define LSM_ADDR     0x6A

#define XL_ODR_26HZ  0x20
#define XL_ODR_416HZ 0x60
#define G_ODR_416HZ  0x60

// Module state 
volatile SleepStage sleepStage  = SLEEP_AWAKE;
volatile bool       imuWakeFlag = false;

static unsigned long idleEnteredMs  = 0;  // when motion pipeline first went idle
static unsigned long wakeSettleMs   = 0;  // when sleepManagerWakeIMU() was called
static unsigned long lpEnteredMs    = 0;  // when IMU LP was entered (for recal decision)

static bool          wasIdle        = false;
static uint8_t       savedCtrl1XL   = XL_ODR_416HZ;
static uint8_t       savedCtrl2G    = G_ODR_416HZ;
static volatile bool pendingWakeSettle = false;  // always set on wake — 120ms blackout
static volatile bool pendingWakeRecal  = false;  // set only when recal is also needed

// Only force recalibration after waking from deep sleep, or after the gyro
// has been off long enough for thermal drift to matter (~5 minutes).
static constexpr unsigned long RECAL_AFTER_LP_MS = 5UL * 60UL * 1000UL;

// I2C helpers — Wire1 at 0x6A (SA0 LOW on XIAO nRF52840 Sense) 
static uint8_t lsmRead(uint8_t reg) {
  Wire1.beginTransmission(LSM_ADDR);
  Wire1.write(reg);
  Wire1.endTransmission(false);
  Wire1.requestFrom(LSM_ADDR, (uint8_t)1);
  return Wire1.available() ? Wire1.read() : 0xFF;
}

static void lsmWrite(uint8_t reg, uint8_t val) {
  Wire1.beginTransmission(LSM_ADDR);
  Wire1.write(reg);
  Wire1.write(val);
  Wire1.endTransmission();
}

// ISR 
static void imuInt1ISR() {
  imuWakeFlag = true;
}

// Arm wakeup interrupt 
static void armWakeupInterrupt() {
  lsmWrite(SLP_WAKE_UP_DUR, (uint8_t)((SLEEP_WAKEUP_DUR & 0x03) << 4));

  // Preserve bit7 (tap single/double enable) written by tap.cpp
  uint8_t wuth = lsmRead(REG_WAKE_UP_THS);
  wuth = (wuth & 0xC0) | (SLEEP_WAKEUP_THS & 0x3F);
  lsmWrite(REG_WAKE_UP_THS, wuth);

  // INTERRUPTS_ENABLE=1, SLOPE_FDS=0 (HP filter is gated in LP — must be 0)
  uint8_t tcfg = lsmRead(REG_TAP_CFG);
  tcfg |=  (1 << 7);
  tcfg &= ~(1 << 4);
  lsmWrite(REG_TAP_CFG, tcfg);

  // OR in INT1_WU (bit5), preserve tap routing bits
  uint8_t md1 = lsmRead(REG_MD1_CFG);
  md1 |= (1 << 5);
  lsmWrite(REG_MD1_CFG, md1);

  // Clear any stale latch BEFORE re-arming the edge interrupt.
  // If INT1 is already high from a previous event, RISING will never fire
  // because no low→high edge will occur. Reading WAKE_UP_SRC de-asserts INT1.
  (void)lsmRead(SLP_WAKE_UP_SRC);
  // Small delay to let the INT1 line settle low before we arm the edge detect
  delay(2);

  // Re-attach with RISING — guaranteed clean edge now that latch is cleared
  detachInterrupt(digitalPinToInterrupt(IMU_INT1_PIN));
  attachInterrupt(digitalPinToInterrupt(IMU_INT1_PIN), imuInt1ISR, RISING);
  imuWakeFlag = false;  // discard any flag set before the re-arm

  Serial.print("[SLEEP] armWakeup — TAP_CFG=0x"); Serial.print(lsmRead(REG_TAP_CFG), HEX);
  Serial.print(" MD1_CFG=0x"); Serial.print(lsmRead(REG_MD1_CFG), HEX);
  Serial.print(" WAKE_UP_THS=0x"); Serial.println(lsmRead(REG_WAKE_UP_THS), HEX);
}

// Disarm — restore registers for normal HP operation 
static void disarmWakeupInterrupt() {
  // Restore SLOPE_FDS=1 for tap detection HP filter path
  uint8_t tcfg = lsmRead(REG_TAP_CFG);
  tcfg |= (1 << 4);
  lsmWrite(REG_TAP_CFG, tcfg);

  // Clear INT1_WU to stop wakeup engine driving INT1 at full rate
  uint8_t md1 = lsmRead(REG_MD1_CFG);
  md1 &= ~(1 << 5);
  lsmWrite(REG_MD1_CFG, md1);
}

// Enter IMU LP 
static void enterImuLP() {
  if (sleepStage >= SLEEP_IMU_LP) return;

  savedCtrl1XL = lsmRead(REG_CTRL1_XL);
  savedCtrl2G  = lsmRead(SLP_CTRL2_G);

  lsmWrite(SLP_CTRL2_G,  savedCtrl2G & 0x0F);                      // gyro off
  lsmWrite(REG_CTRL1_XL, (savedCtrl1XL & 0x0F) | XL_ODR_26HZ);    // accel 26 Hz
  lsmWrite(SLP_CTRL6_C,  lsmRead(SLP_CTRL6_C) | (1 << 4));         // XL_HM_MODE=1
  lsmWrite(SLP_CTRL7_G,  lsmRead(SLP_CTRL7_G) | (1 << 7));         // G_HM_MODE=1

  armWakeupInterrupt();

  lpEnteredMs = millis();
  sleepStage = SLEEP_IMU_LP;
  Serial.print("[SLEEP] IMU LP entered — idle for ");
  Serial.print((millis() - idleEnteredMs) / 1000); Serial.println("s");

  // Full register readback — confirms chip actually accepted all writes
  Serial.print("[SLEEP] CTRL1_XL=0x"); Serial.print(lsmRead(REG_CTRL1_XL), HEX);
  Serial.print(" CTRL2_G=0x");  Serial.print(lsmRead(SLP_CTRL2_G), HEX);
  Serial.print(" CTRL6_C=0x");  Serial.print(lsmRead(SLP_CTRL6_C), HEX);
  Serial.print(" CTRL7_G=0x");  Serial.println(lsmRead(SLP_CTRL7_G), HEX);
  Serial.print("[SLEEP] TAP_CFG=0x");    Serial.print(lsmRead(REG_TAP_CFG), HEX);
  Serial.print(" MD1_CFG=0x");          Serial.print(lsmRead(REG_MD1_CFG), HEX);
  Serial.print(" WAKE_UP_THS=0x");      Serial.print(lsmRead(REG_WAKE_UP_THS), HEX);
  Serial.print(" WAKE_UP_DUR=0x");      Serial.println(lsmRead(SLP_WAKE_UP_DUR), HEX);
  // Expected when working:
  //  CTRL1_XL=0x2_ (ODR=26Hz, _ = FS bits)   CTRL2_G=0x0_ (ODR=off)
  //  CTRL6_C bit4=1 (XL_HM_MODE)              CTRL7_G bit7=1 (G_HM_MODE)
  //  TAP_CFG=0x80  MD1_CFG=0x20  WAKE_UP_THS=0x01
}

// Enter deep sleep 
static void enterDeepSleep() {
  if (sleepStage >= SLEEP_DEEP) return;
  if (sleepStage < SLEEP_IMU_LP) enterImuLP();

  sleepStage = SLEEP_DEEP;
  Serial.println("[SLEEP] Deep sleep — WFE on INT1");
  Serial.flush();

  digitalWrite(LED_RED, LOW); delay(80); digitalWrite(LED_RED, HIGH);

  while (!imuWakeFlag) {
    (void)lsmRead(SLP_WAKE_UP_SRC);
    sd_app_evt_wait();
  }
}

// Public: wake 
void sleepManagerWakeIMU() {
  (void)lsmRead(SLP_WAKE_UP_SRC);

  lsmWrite(SLP_CTRL6_C, lsmRead(SLP_CTRL6_C) & ~(1 << 4));
  lsmWrite(SLP_CTRL7_G, lsmRead(SLP_CTRL7_G) & ~(1 << 7));
  lsmWrite(REG_CTRL1_XL, savedCtrl1XL);
  lsmWrite(SLP_CTRL2_G,  savedCtrl2G);

  disarmWakeupInterrupt();

  // Only recalibrate if gyro was off long enough for thermal drift to accumulate,
  // or if waking from full deep sleep. Short LP naps reuse the existing bias.
  unsigned long lpDuration = millis() - lpEnteredMs;
  bool needsRecal = (sleepStage == SLEEP_DEEP) || (lpDuration >= RECAL_AFTER_LP_MS);

  wakeSettleMs      = millis();
  pendingWakeSettle = true;          // always block output for 120ms
  pendingWakeRecal  = needsRecal;    // additionally recalibrate if needed

  wasIdle       = false;
  idleEnteredMs = 0;
  lpEnteredMs   = 0;

  // Reset motion filter state to prevent a cursor jump on the first frame.
  // After sleep: angleX/Y are stale, gravX/Y/Z drifted, accumX/Y is dirty,
  // and lastTime is old so dt would be huge on the first loop iteration.
  // Zeroing these here means the first frame integrates 0 motion cleanly.
  extern float angleX, angleY;
  extern float accumX, accumY;
  extern float gravX, gravY, gravZ;
  extern float prevAx, prevAy, prevAz;
  extern unsigned long lastTime;
  angleX = angleY = 0.0f;
  accumX = accumY = 0.0f;
  // Reseed gravity from current accel so projection is correct immediately.
  // Can't call imu.readFloat* here (gyro not fully settled), but accel is
  // already running — read it directly via Wire1.
  // Simpler: just reset to neutral [0,0,1] and let the LP filter converge
  // over the first ~20 frames (200 ms) of real use.
  gravX = 0.0f; gravY = 0.0f; gravZ = 1.0f;
  prevAx = 0.0f; prevAy = 0.0f; prevAz = 0.0f;
  // Set lastTime to now so the first dt = 0 rather than (now - sleepEntryTime)
  lastTime = millis();

  sleepStage = SLEEP_AWAKE;
  if (needsRecal)
    Serial.println("[SLEEP] Awake — gyro settling, recal needed");
  else
    Serial.println("[SLEEP] Awake — short LP, reusing existing bias");
}

// Public: init 
void sleepManagerInit() {
  pinMode(IMU_INT1_PIN, INPUT);
  attachInterrupt(digitalPinToInterrupt(IMU_INT1_PIN), imuInt1ISR, RISING);

  // Sanity check: WHO_AM_I should return 0x6A for LSM6DS3
  uint8_t whoami = imuReadReg(0x0F);
  Serial.print("[SLEEP] WHO_AM_I=0x"); Serial.print(whoami, HEX);
  if (whoami == 0x6A) Serial.println(" (OK)");
  else                Serial.println(" (WRONG — I2C not working, sleep disabled)");

  if (whoami != 0x6A) return;  // don't arm anything if we can't talk to the IMU

  Serial.print("[SLEEP] Init — INT1 pin="); Serial.print(IMU_INT1_PIN);
  Serial.print(", WU_THS="); Serial.print(SLEEP_WAKEUP_THS);
  Serial.print(" (~"); Serial.print(SLEEP_WAKEUP_THS * 7.8f, 0); Serial.print(" mg)");
  Serial.print(", IMU_LP after "); Serial.print(SLEEP_IMU_IDLE_MS / 1000); Serial.print("s");
  Serial.print(", deep after ");   Serial.print(SLEEP_DEEP_IDLE_MS / 1000); Serial.println("s");
}

// Public: per-loop update 
// Must be called AFTER idleFrames/idle is updated by the motion pipeline.
// Returns true → caller must skip IMU reads this iteration.
bool sleepManagerUpdate(unsigned long nowMs, bool idle, bool bleConnected) {

  // ISR wakeup 
  if (imuWakeFlag) {
    imuWakeFlag = false;
    Serial.print("[SLEEP] INT1 fired — stage="); Serial.println((int)sleepStage);
    if (sleepStage == SLEEP_DEEP || sleepStage == SLEEP_IMU_LP) {
      sleepManagerWakeIMU();
    } else {
      (void)lsmRead(SLP_WAKE_UP_SRC);  // normal-mode edge, clear latch only
    }
  }

  // Gyro settling after wake 
  if (pendingWakeRecal) {
    if (nowMs - wakeSettleMs >= 120) {
      pendingWakeRecal = false;
      wakeSettleMs     = 0;
      extern void calibrateGyroBias();
      calibrateGyroBias();
      Serial.println("[SLEEP] Post-wake recal done");
    }
    return true;
  }

  // IMU_LP path 
  // main.cpp returns early when we return true, so idleFrames never increments
  // and the motion pipeline's `idle` flag is stale. Use our own idleEnteredMs
  // (captured before LP entry) to drive the deep-sleep countdown independently.
  if (sleepStage == SLEEP_IMU_LP) {
    // Periodic log: confirms loop is running, shows live INT1 pin level and
    // WAKE_UP_SRC register. If INT1_pin never goes high after movement, the
    // wakeup engine is not generating an interrupt — check register values.
    static unsigned long lastLpLog = 0;
    if (nowMs - lastLpLog >= 5000) {
      lastLpLog = nowMs;
      unsigned long lpSecs = idleEnteredMs ? (nowMs - idleEnteredMs) / 1000 : 0;
      Serial.print("[SLEEP] LP tick — idle="); Serial.print(lpSecs);
      Serial.print("s INT1="); Serial.print(digitalRead(IMU_INT1_PIN));
      Serial.print(" WAKE_UP_SRC=0x"); Serial.println(lsmRead(SLP_WAKE_UP_SRC), HEX);
    }
    if (!bleConnected && idleEnteredMs != 0
        && (nowMs - idleEnteredMs >= SLEEP_DEEP_IDLE_MS))
    {
      Serial.println("[SLEEP] Deep sleep threshold reached");
      enterDeepSleep();
    }
    return true;
  }

  // AWAKE path 
  if (!idle) {
    if (wasIdle) {
      Serial.println("[SLEEP] Motion — idle timer reset");
    }
    wasIdle       = false;
    idleEnteredMs = 0;
    return false;
  }

  if (!wasIdle) {
    wasIdle       = true;
    idleEnteredMs = nowMs;
    Serial.println("[SLEEP] Idle started");
  }

  unsigned long idleElapsed = nowMs - idleEnteredMs;

  // Progress report every 5 s while waiting for LP threshold
  #ifdef DEBUG
  { static unsigned long lastReport = 0;
    if (nowMs - lastReport >= 5000) { lastReport = nowMs;
      Serial.print("[SLEEP] idle="); Serial.print(idleElapsed/1000);
      Serial.print("s / LP@"); Serial.print(SLEEP_IMU_IDLE_MS/1000); Serial.println("s");
    }
  }
  #endif

  if (idleElapsed >= SLEEP_IMU_IDLE_MS) {
    enterImuLP();
    return true;
  }

  return false;
}