#include <Arduino.h>
#include <hwdt_app_entry.h>
#include <user_interface.h>

extern "C" {

  struct WDTracksLastCall {
    struct {
      void *this_fn;
      void *call_site;
      void *sp;
    } enter;

    struct {
      void *this_fn;
      void *call_site;
      void *sp;
    } exit;

    bool in;
  };
  struct WDTracksLastCall wd_tracks;

  IRAM_ATTR void __cyg_profile_func_enter(void *this_fn, void *call_site) __attribute__((no_instrument_function));
  IRAM_ATTR void __cyg_profile_func_exit(void *this_fn, void *call_site) __attribute__((no_instrument_function));
  void print_wdtracks(void) __attribute__((no_instrument_function));

  void custom_crash_callback(struct rst_info * rst_info, uint32_t stack, uint32_t stack_end) __attribute__((no_instrument_function));

  /*
   * Printing from hwdt_pre_sdk_init() requires special considerations. Normal
   * "C" runtime initialization has not run at the time it is called. For
   * printing from this context, we use umm_info_safe_printf_P. It is capable of
   * handling this situation and is a function included in umm_malloc.
   */
  int umm_info_safe_printf_P(const char *fmt, ...);
  #define ETS_PRINTF(fmt, ...) umm_info_safe_printf_P(PSTR(fmt), ##__VA_ARGS__)

#if defined(DEBUG_ESP_HWDT) || defined(DEBUG_ESP_HWDT_NOEXTRA4K)
  void hwdt_pre_sdk_init(void) __attribute__((no_instrument_function));

  // Called from HWDT just before starting SDK
  // Serial speed has been initialized for us.
  void hwdt_pre_sdk_init(void) {
    if (REASON_WDT_RST == hwdt_info.reset_reason) {
      /*
        Note, we rely on the previous values of the wd_tracks structure, still
        being set from before the crash. At the time we are called here, the SDK
        has not been started. The "C" runtime code that will zero the structure
        does not run until later when the SDK calls user_init().
      */
      print_wdtracks();
    }
  }
#endif

  // Called from Postmortem
  void custom_crash_callback(struct rst_info * rst_info, uint32_t stack, uint32_t stack_end)
  {
    (void)stack;
    (void)stack_end;
    if (REASON_SOFT_WDT_RST == rst_info->reason) {
      print_wdtracks();
    }
  }


  /*
    Called at function entry after stackframe setup and registers are saved.
  */
  void __cyg_profile_func_enter(void *this_fn, void *call_site)
  {
    wd_tracks.enter.this_fn = this_fn;       // The address of the function called
    wd_tracks.enter.call_site = call_site;   // The saved return address "a0" of the callee
    asm volatile (
        "s32i a1, %[wd_tracks], %[enter_sp]\n\t"
        : // No output constraints
        : [wd_tracks]"a"(&wd_tracks), [enter_sp]"I"(offsetof(struct WDTracksLastCall, enter.sp))
        :);
    wd_tracks.in = true;
  }


  /*
    Called at function exit just before saved registers and stackframe are
    restored.
  */
  void __cyg_profile_func_exit(void *this_fn, void *call_site)
  {
    wd_tracks.exit.this_fn = this_fn;      // The address of the function called
    wd_tracks.exit.call_site = call_site;  // The saved return address "a0" of the callee
    asm volatile (
        "s32i a1, %[wd_tracks], %[exit_sp]\n\t"
        : // No output constraints
        : [wd_tracks]"a"(&wd_tracks), [exit_sp]"I"(offsetof(struct WDTracksLastCall, exit.sp))
        :);
    wd_tracks.in = false;
  }

  void print_wdtracks(void)
  {
      if (NULL == wd_tracks.enter.this_fn)  {
        // Compiler flags -finstrument-functions, etc. are most likely missing.
        return;
      }
      if (wd_tracks.in) {
        ETS_PRINTF("\nepc1=%p\nepc1=%p\n",
                   wd_tracks.enter.this_fn, wd_tracks.enter.call_site);

        if (wd_tracks.exit.this_fn != wd_tracks.enter.this_fn &&
            wd_tracks.exit.this_fn != wd_tracks.enter.call_site)
          ETS_PRINTF("epc1=%p\n", wd_tracks.exit.this_fn);
        if (wd_tracks.exit.call_site != wd_tracks.enter.this_fn &&
            wd_tracks.exit.call_site != wd_tracks.enter.call_site)
          ETS_PRINTF("epc1=%p\n",  wd_tracks.exit.call_site);

      } else {
        ETS_PRINTF("\nepc1=%p\nepc1=%p\n",
                   wd_tracks.exit.this_fn, wd_tracks.exit.call_site);

        if (wd_tracks.enter.this_fn != wd_tracks.exit.this_fn &&
            wd_tracks.enter.this_fn != wd_tracks.exit.call_site)
          ETS_PRINTF("epc1=%p\n", wd_tracks.enter.this_fn);
        if (wd_tracks.enter.call_site != wd_tracks.exit.this_fn &&
            wd_tracks.enter.call_site != wd_tracks.exit.call_site)
          ETS_PRINTF("epc1=%p\n",  wd_tracks.enter.call_site);
      }
      ETS_PRINTF("\n--------- OR FOR WDT, CUT HERE FOR EXCEPTION DECODER ---------\n");

      if (wd_tracks.in) {
        ETS_PRINTF("\nLast function entered:\n"
                   "  called function: %p\n"
                   "  callee return:   %p\n"
                   "  SP:              %p\n",
                   wd_tracks.enter.this_fn, wd_tracks.enter.call_site, wd_tracks.enter.sp);
      } else {
        ETS_PRINTF("\nLast function exited:\n"
                   "  called function: %p\n"
                   "  callee return:   %p\n"
                   "  SP:              %p\n",
                   wd_tracks.exit.this_fn, wd_tracks.exit.call_site, wd_tracks.exit.sp);
      }
  }
};