#include <Arduino.h>
#include <string.h>
#include <avr/wdt.h>

#include <helpers/ArduinoHelpers.h>
#include <target.h>

#include "FixedPoolPacketManager.h"
#include "FlashStore.h"
#include "PublicApp.h"
#include "UITask.h"

uint8_t mcusr_mirror __attribute__((section(".noinit")));
uint8_t wdt_stage_mirror __attribute__((section(".noinit")));

extern char __heap_start;

namespace {
const uint8_t WDT_STAGE_IDLE = 'I';
const uint8_t WDT_STAGE_SETUP = 'S';
const uint8_t WDT_STAGE_APP = 'A';
const uint8_t WDT_STAGE_UI = 'U';
const uint8_t WDT_STAGE_FLASH = 'F';
const uint8_t WDT_STAGE_POWER = 'P';
const uint8_t STACK_CANARY = 0xA5;
const uint8_t STACK_GUARD_BYTES = 64;

uint8_t* stack_canary_start = NULL;
uint8_t* stack_canary_end = NULL;
uint16_t stack_min_free = 0;

static uint8_t* currentStackPointer() {
  uint16_t sp;
  asm volatile(
      "in %A0, __SP_L__" "\n\t"
      "in %B0, __SP_H__"
      : "=r"(sp));
  return (uint8_t*)sp;
}

static uint16_t currentStackFree() {
  uint8_t* heap_end = (uint8_t*)&__heap_start;
  uint8_t* sp = currentStackPointer();
  if (sp <= heap_end + STACK_GUARD_BYTES) {
    return 0;
  }
  return (uint16_t)(sp - heap_end - STACK_GUARD_BYTES);
}

static uint16_t scanStackFree() {
  if (!stack_canary_start || !stack_canary_end || stack_canary_end <= stack_canary_start) {
    return 0;
  }

  uint8_t* heap_end = stack_canary_start;
  if (heap_end >= stack_canary_end) {
    return 0;
  }

  uint8_t* p = stack_canary_end;
  while (p > heap_end && *(p - 1) != STACK_CANARY) {
    --p;
  }
  return (uint16_t)(p - heap_end);
}

static void initStackCanary() {
  uint8_t* start = (uint8_t*)&__heap_start;
  uint8_t* sp = currentStackPointer();
  if (sp <= start + STACK_GUARD_BYTES) {
    stack_canary_start = NULL;
    stack_canary_end = NULL;
    stack_min_free = 0;
    return;
  }

  stack_canary_start = start;
  stack_canary_end = sp - STACK_GUARD_BYTES;
  memset(stack_canary_start, STACK_CANARY, stack_canary_end - stack_canary_start);
  stack_min_free = scanStackFree();
}

static void updateStackMinFree() {
  uint16_t free_now = scanStackFree();
  if (stack_min_free == 0 || free_now < stack_min_free) {
    stack_min_free = free_now;
  }
}
}

void captureResetCause(void) __attribute__((naked, used, section(".init3")));
void captureResetCause(void) {
  mcusr_mirror = MCUSR;
  MCUSR = 0;
  wdt_disable();
}

static void setWatchdogStage(uint8_t stage) {
  wdt_stage_mirror = stage;
}

extern "C" void SRXEWatchdogPoll(void) {
  setWatchdogStage(WDT_STAGE_FLASH);
  wdt_reset();
}

extern "C" void MeshCoreWatchdogStage(uint8_t stage) {
  setWatchdogStage(stage);
}

extern "C" uint16_t MeshCoreStackFreeNow(void) {
  updateStackMinFree();
  return currentStackFree();
}

extern "C" uint16_t MeshCoreStackFreeMin(void) {
  updateStackMinFree();
  return stack_min_free;
}

static ArduinoMillis ms_clock;
static FixedPoolPacketManager<4> packet_mgr;
static FlashStore flash_store;
static PublicRTCClock rtc_clock;
static PublicApp app(radio_driver, ms_clock, packet_mgr, rtc_clock, flash_store);
static UITask ui(board, app);

#ifndef SMARTRESPONSE_BUILD_EPOCH
#define SMARTRESPONSE_BUILD_EPOCH 1715770351UL
#endif

static uint32_t initialClockTime() {
  uint32_t start = SMARTRESPONSE_BUILD_EPOCH;
  uint32_t latest = flash_store.getLatestTimestamp();
  if (latest != 0xFFFFFFFFUL && latest >= start) {
    start = latest + 1;
  }
  return start;
}

static void waitForPowerButtonRelease() {
  while (board.isPowerButtonPressed()) {
    wdt_reset();
    delay(10);
  }
  delay(50);
}

static void sleepUntilPowerButton() {
  setWatchdogStage(WDT_STAGE_POWER);
  waitForPowerButtonRelease();
  radio_driver.powerOff();
  wdt_disable();
  board.sleep(0);
  wdt_enable(WDTO_4S);
  app.resetRadio();
  ui.redraw();
  waitForPowerButtonRelease();
}

static void halt() {
  while (1) {
  }
}

void setup() {
  wdt_disable();
  uint8_t last_wdt_stage = (mcusr_mirror & _BV(WDRF)) ? wdt_stage_mirror : '-';
  setWatchdogStage(WDT_STAGE_SETUP);
  delay(100);

  board.begin();
  flash_store.begin();
  rtc_clock.begin(initialClockTime());

  if (!radio_init()) {
    halt();
  }

  app.setResetCause(mcusr_mirror);
  app.setWatchdogStage(last_wdt_stage);
  app.begin();
  ui.begin();
  initStackCanary();
  setWatchdogStage(WDT_STAGE_IDLE);
  wdt_enable(WDTO_4S);
}

void loop() {
  setWatchdogStage(WDT_STAGE_IDLE);
  wdt_reset();
  rtc_clock.tick();
  setWatchdogStage(WDT_STAGE_APP);
  app.loop();
  wdt_reset();
  setWatchdogStage(WDT_STAGE_UI);
  ui.loop();
  if (board.isPowerButtonPressed()) {
    sleepUntilPowerButton();
  }
  updateStackMinFree();
  setWatchdogStage(WDT_STAGE_IDLE);
  board.idleSleep();
}