#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

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

// Post-wake stillness gate: don't calibrate while the device is moving.
// Each axis of the gyro must stay below RECAL_STILL_DPS for RECAL_STILL_FRAMES
// consecutive loop ticks before calibration fires. If the device keeps moving
// past RECAL_WAIT_MAX_MS, recal is skipped and the pre-sleep bias is kept.
static constexpr float         RECAL_STILL_DPS    = 10.0f;  // deg/s per axis
static constexpr uint8_t       RECAL_STILL_FRAMES = 10;     // ~100 ms at 10 ms/frame
static constexpr unsigned long RECAL_WAIT_MAX_MS  = 3000;   // give up after 3 s

static uint8_t       recalStillFrames = 0;
static unsigned long recalLastWarnMs  = 0;

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

static volatile bool btnWakeFlag = false;
static void btnWakeISR() {
  btnWakeFlag = 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)
  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);

  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.
  (void)lsmRead(SLP_WAKE_UP_SRC);
  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;

#ifdef DEBUG
  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);
#endif
}

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

  // Arm button wake interrupt
  #if BTN_LEFT_PIN != BTN_PIN_NONE
    btnWakeFlag = false;
    attachInterrupt(digitalPinToInterrupt(BTN_LEFT_PIN), btnWakeISR, FALLING);
  #endif

  // Turn off all LEDs for sleep
  digitalWrite(LED_RED, HIGH); digitalWrite(LED_GREEN, HIGH); digitalWrite(LED_BLUE, HIGH);

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

#ifdef DEBUG
  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);
#endif
}

// 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, HIGH); digitalWrite(LED_GREEN, HIGH); digitalWrite(LED_BLUE, HIGH);

  while (!imuWakeFlag && !btnWakeFlag) {
    (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();

  // Detach button wake interrupt — normal polling takes over
  #if BTN_LEFT_PIN != BTN_PIN_NONE
    detachInterrupt(digitalPinToInterrupt(BTN_LEFT_PIN));
  #endif

  // 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.
  extern float accumX, accumY;
  extern float prevAx, prevAy, prevAz;
  extern unsigned long lastTime;
  accumX = accumY = 0.0f;
  prevAx = 0.0f; prevAy = 0.0f; prevAz = 0.0f;
  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 (IMU motion or button press)
  bool woke = false;
  if (imuWakeFlag) {
    imuWakeFlag = false;
    Serial.print("[SLEEP] INT1 fired - stage="); Serial.println((int)sleepStage);
    woke = true;
  }
  if (btnWakeFlag) {
    btnWakeFlag = false;
    Serial.print("[SLEEP] Button fired - stage="); Serial.println((int)sleepStage);
    woke = true;
  }
  if (woke) {
    if (sleepStage == SLEEP_DEEP || sleepStage == SLEEP_IMU_LP) {
      sleepManagerWakeIMU();
    } else {
      (void)lsmRead(SLP_WAKE_UP_SRC);
    }
  }

  // Post-wake recalibration — gated on device being still
  if (pendingWakeRecal) {
    if (nowMs - wakeSettleMs < 120) return true;  // initial hardware settle

    // Sample gyro magnitude; each axis must be below threshold
    float gx = fabsf(imu.readFloatGyroX());
    float gy = fabsf(imu.readFloatGyroY());
    float gz = fabsf(imu.readFloatGyroZ());
    bool still = (gx < RECAL_STILL_DPS && gy < RECAL_STILL_DPS && gz < RECAL_STILL_DPS);

    if (!still) {
      recalStillFrames = 0;
      // Rate-limited log so we don't flood serial while waiting
      if (nowMs - recalLastWarnMs >= 500) {
        recalLastWarnMs = nowMs;
        Serial.print("[SLEEP] Waiting for still: gx="); Serial.print(gx, 1);
        Serial.print(" gy="); Serial.print(gy, 1);
        Serial.print(" gz="); Serial.println(gz, 1);
      }
      if (nowMs - wakeSettleMs >= RECAL_WAIT_MAX_MS) {
        // Device never settled — keep pre-sleep bias rather than corrupt it
        pendingWakeRecal = false;
        recalStillFrames = 0;
        recalLastWarnMs  = 0;
        Serial.println("[SLEEP] Recal skipped — device still moving after timeout");
      }
      return true;
    }

    // Device is still — accumulate consecutive still frames
    if (++recalStillFrames < RECAL_STILL_FRAMES) return true;

    pendingWakeRecal = false;
    recalStillFrames = 0;
    recalLastWarnMs  = 0;
    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;
}