#include <Arduino.h>

#include "driver/rtc_io.h"
#include "rom/ets_sys.h"
#include "rom/rtc.h"
#include "esp_attr.h"
#include "esp_sleep.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "soc/rtc.h"
#include "soc/rtc_cntl_reg.h"
#include "soc/rtc_io_reg.h"
#include "soc/timer_group_reg.h"
#include "soc/uart_reg.h"
#include <stdio.h>
#include <string.h>

RTC_IRAM_ATTR void deepsleep_for_us(uint64_t duration_us);

#define S_TO_NS 1000000ULL
// Boot counter value, stored in RTC_SLOW_MEM
static size_t RTC_DATA_ATTR boots;
static size_t RTC_DATA_ATTR max_boots;

// Function which runs after exit from deep sleep
static void RTC_IRAM_ATTR wake_stub();
void setup(){
  
}
void loop(void) {
  if (rtc_get_reset_reason(0) == DEEPSLEEP_RESET) {
    printf("I'm back! Wake up from (extended) deep sleep. Should be 60s elapsed.\n");
    printf("Boot count=%d\n", boots);
    while (1) {
      ;
    }
  } else {
    printf("Not a deep sleep wake up\n");
  }

  printf("Going to deep sleep for 5x12=60s second\n");
  boots = 0;
  max_boots = 12;
  // hold high during deep sleep
  esp_sleep_pd_config(ESP_PD_DOMAIN_RTC_PERIPH, ESP_PD_OPTION_ON);

  // Set the wake stub function
  esp_set_deep_sleep_wake_stub(&wake_stub);

  esp_sleep_enable_timer_wakeup(5 * S_TO_NS);
  esp_deep_sleep_start();
}

static const char RTC_RODATA_ATTR wake_fmt_str[] = "count=%d\n";
static const char RTC_RODATA_ATTR sleep_fmt_str[] = "sleeping\n";

static void RTC_IRAM_ATTR wake_stub() {
  // Increment the Boot counter
  boots++;
  // and print the Boot counter value:
  ets_printf(wake_fmt_str, boots);

  if (boots >= max_boots) {
    // On revision 0 of ESP32, this function must be called:
    esp_default_wake_deep_sleep();
    // Return from the wake stub function to continue
    // booting the firmware.
    return;
  }

  // Print status
  ets_printf(sleep_fmt_str);
  // Wait for UART to end transmitting.
  // feed the watchdog
  REG_WRITE(TIMG_WDTFEED_REG(0), 1);
  while (REG_GET_FIELD(UART_STATUS_REG(0), UART_ST_UTX_OUT)) {

  }
  deepsleep_for_us(5 * S_TO_NS);

  // Set the pointer of the wake stub function.
  REG_WRITE(RTC_ENTRY_ADDR_REG, (uint32_t)&wake_stub);
  // Go to sleep.
  CLEAR_PERI_REG_MASK(RTC_CNTL_STATE0_REG, RTC_CNTL_SLEEP_EN);
  SET_PERI_REG_MASK(RTC_CNTL_STATE0_REG, RTC_CNTL_SLEEP_EN);
  // A few CPU cycles may be necessary for the sleep to start...
  while (true) {
    ;
  }
  // never reaches here.
}

// Comment out this line if you're using the internal RTC RC (150KHz) oscillator.
//#define USE_EXTERNAL_RTC_CRYSTAL
#ifdef USE_EXTERNAL_RTC_CRYSTAL
#define DEEP_SLEEP_TIME_OVERHEAD_US (650 + 100 * 240 / CONFIG_ESP32_DEFAULT_CPU_FREQ_MHZ)
#else
#define DEEP_SLEEP_TIME_OVERHEAD_US (250 + 100 * 240 / CONFIG_ESP32_DEFAULT_CPU_FREQ_MHZ)
#endif // USE_EXTERNAL_RTC_CRYSTAL

RTC_IRAM_ATTR void deepsleep_for_us(uint64_t duration_us) {
  // Feed watchdog
  REG_WRITE(TIMG_WDTFEED_REG(0), 1);
  // Get RTC calibration
  uint32_t period = REG_READ(RTC_SLOW_CLK_CAL_REG);
  // Calculate sleep duration in microseconds
  int64_t sleep_duration = (int64_t)duration_us - (int64_t)DEEP_SLEEP_TIME_OVERHEAD_US;
  if (sleep_duration < 0) {
    sleep_duration = 0;
  }
  // Convert microseconds to RTC clock cycles
  int64_t rtc_count_delta = (sleep_duration << RTC_CLK_CAL_FRACT) / period;
  // Feed watchdog
  REG_WRITE(TIMG_WDTFEED_REG(0), 1);
  // Get current RTC time
  SET_PERI_REG_MASK(RTC_CNTL_TIME_UPDATE_REG, RTC_CNTL_TIME_UPDATE);
  while (GET_PERI_REG_MASK(RTC_CNTL_TIME_UPDATE_REG, RTC_CNTL_TIME_VALID) == 0) {
    ets_delay_us(1);
  }
  SET_PERI_REG_MASK(RTC_CNTL_INT_CLR_REG, RTC_CNTL_TIME_VALID_INT_CLR);
  uint64_t now = READ_PERI_REG(RTC_CNTL_TIME0_REG);
  now |= ((uint64_t)READ_PERI_REG(RTC_CNTL_TIME1_REG)) << 32;
  // Set wakeup time
  uint64_t future = now + rtc_count_delta;
  WRITE_PERI_REG(RTC_CNTL_SLP_TIMER0_REG, future & UINT32_MAX);
  WRITE_PERI_REG(RTC_CNTL_SLP_TIMER1_REG, future >> 32);
  // Start RTC deepsleep timer
  REG_SET_FIELD(RTC_CNTL_WAKEUP_STATE_REG, RTC_CNTL_WAKEUP_ENA, RTC_TIMER_TRIG_EN); // Wake up on timer
  WRITE_PERI_REG(RTC_CNTL_SLP_REJECT_CONF_REG, 0);                                  // Clear sleep rejection cause
  // Go to sleep
  CLEAR_PERI_REG_MASK(RTC_CNTL_STATE0_REG, RTC_CNTL_SLEEP_EN);
  SET_PERI_REG_MASK(RTC_CNTL_STATE0_REG, RTC_CNTL_SLEEP_EN);
}

#include <Arduino.h>

#include "driver/rtc_io.h"
#include "rom/ets_sys.h"
#include "rom/rtc.h"
#include "esp_attr.h"
#include "esp_sleep.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "soc/rtc.h"
#include "soc/rtc_cntl_reg.h"
#include "soc/rtc_io_reg.h"
#include "soc/timer_group_reg.h"
#include "soc/uart_reg.h"
#include <stdio.h>
#include <string.h>

RTC_IRAM_ATTR void deepsleep_for_us(uint64_t duration_us);

#define S_TO_NS 1000000ULL
// Boot counter value, stored in RTC_SLOW_MEM
static size_t RTC_DATA_ATTR boots;
static size_t RTC_DATA_ATTR max_boots;

// Function which runs after exit from deep sleep
static void RTC_IRAM_ATTR wake_stub();
void setup(){
  
}
void loop(void) {
  if (rtc_get_reset_reason(0) == DEEPSLEEP_RESET) {
    printf("I'm back! Wake up from (extended) deep sleep. Should be 60s elapsed.\n");
    printf("Boot count=%d\n", boots);
    while (1) {
      ;
    }
  } else {
    printf("Not a deep sleep wake up\n");
  }

  printf("Going to deep sleep for 5x12=60s second\n");
  boots = 0;
  max_boots = 12;
  // hold high during deep sleep
  esp_sleep_pd_config(ESP_PD_DOMAIN_RTC_PERIPH, ESP_PD_OPTION_ON);

  // Set the wake stub function
  esp_set_deep_sleep_wake_stub(&wake_stub);

  esp_sleep_enable_timer_wakeup(5 * S_TO_NS);
  esp_deep_sleep_start();
}

static const char RTC_RODATA_ATTR wake_fmt_str[] = "count=%d\n";
static const char RTC_RODATA_ATTR sleep_fmt_str[] = "sleeping\n";

static void RTC_IRAM_ATTR wake_stub() {
  // Increment the Boot counter
  boots++;
  // and print the Boot counter value:
  ets_printf(wake_fmt_str, boots);

  if (boots >= max_boots) {
    // On revision 0 of ESP32, this function must be called:
    esp_default_wake_deep_sleep();
    // Return from the wake stub function to continue
    // booting the firmware.
    return;
  }

  // Print status
  ets_printf(sleep_fmt_str);
  // Wait for UART to end transmitting.
  // feed the watchdog
  REG_WRITE(TIMG_WDTFEED_REG(0), 1);
  while (REG_GET_FIELD(UART_STATUS_REG(0), UART_ST_UTX_OUT)) {

  }
  deepsleep_for_us(5 * S_TO_NS);

  // Set the pointer of the wake stub function.
  REG_WRITE(RTC_ENTRY_ADDR_REG, (uint32_t)&wake_stub);
  // Go to sleep.
  CLEAR_PERI_REG_MASK(RTC_CNTL_STATE0_REG, RTC_CNTL_SLEEP_EN);
  SET_PERI_REG_MASK(RTC_CNTL_STATE0_REG, RTC_CNTL_SLEEP_EN);
  // A few CPU cycles may be necessary for the sleep to start...
  while (true) {
    ;
  }
  // never reaches here.
}

// Comment out this line if you're using the internal RTC RC (150KHz) oscillator.
//#define USE_EXTERNAL_RTC_CRYSTAL
#ifdef USE_EXTERNAL_RTC_CRYSTAL
#define DEEP_SLEEP_TIME_OVERHEAD_US (650 + 100 * 240 / CONFIG_ESP32_DEFAULT_CPU_FREQ_MHZ)
#else
#define DEEP_SLEEP_TIME_OVERHEAD_US (250 + 100 * 240 / CONFIG_ESP32_DEFAULT_CPU_FREQ_MHZ)
#endif // USE_EXTERNAL_RTC_CRYSTAL

RTC_IRAM_ATTR void deepsleep_for_us(uint64_t duration_us) {
  // Feed watchdog
  REG_WRITE(TIMG_WDTFEED_REG(0), 1);
  // Get RTC calibration
  uint32_t period = REG_READ(RTC_SLOW_CLK_CAL_REG);
  // Calculate sleep duration in microseconds
  int64_t sleep_duration = (int64_t)duration_us - (int64_t)DEEP_SLEEP_TIME_OVERHEAD_US;
  if (sleep_duration < 0) {
    sleep_duration = 0;
  }
  // Convert microseconds to RTC clock cycles
  int64_t rtc_count_delta = (sleep_duration << RTC_CLK_CAL_FRACT) / period;
  // Feed watchdog
  REG_WRITE(TIMG_WDTFEED_REG(0), 1);
  // Get current RTC time
  SET_PERI_REG_MASK(RTC_CNTL_TIME_UPDATE_REG, RTC_CNTL_TIME_UPDATE);
  while (GET_PERI_REG_MASK(RTC_CNTL_TIME_UPDATE_REG, RTC_CNTL_TIME_VALID) == 0) {
    ets_delay_us(1);
  }
  SET_PERI_REG_MASK(RTC_CNTL_INT_CLR_REG, RTC_CNTL_TIME_VALID_INT_CLR);
  uint64_t now = READ_PERI_REG(RTC_CNTL_TIME0_REG);
  now |= ((uint64_t)READ_PERI_REG(RTC_CNTL_TIME1_REG)) << 32;
  // Set wakeup time
  uint64_t future = now + rtc_count_delta;
  WRITE_PERI_REG(RTC_CNTL_SLP_TIMER0_REG, future & UINT32_MAX);
  WRITE_PERI_REG(RTC_CNTL_SLP_TIMER1_REG, future >> 32);
  // Start RTC deepsleep timer
  REG_SET_FIELD(RTC_CNTL_WAKEUP_STATE_REG, RTC_CNTL_WAKEUP_ENA, RTC_TIMER_TRIG_EN); // Wake up on timer
  WRITE_PERI_REG(RTC_CNTL_SLP_REJECT_CONF_REG, 0);                                  // Clear sleep rejection cause
  // Go to sleep
  CLEAR_PERI_REG_MASK(RTC_CNTL_STATE0_REG, RTC_CNTL_SLEEP_EN);
  SET_PERI_REG_MASK(RTC_CNTL_STATE0_REG, RTC_CNTL_SLEEP_EN);
}