#include "system/firmware_guard.h"

#include <Arduino.h>
#include <esp_ota_ops.h>
#include <esp_system.h>
#include <esp_task_wdt.h>

#include "config.h"
#include "system/boot_target.h"
#include "system/system_store.h"

// Survives crashes and software resets, not a power loss
static const uint32_t RTC_MAGIC = 0xF11D0C0A;
RTC_NOINIT_ATTR static uint32_t rtc_magic;
RTC_NOINIT_ATTR static uint8_t rtc_crash_count;

static uint32_t main_start = 0; // millis() when the main code started
static bool watchdog_running = false;
static bool firmware_pending = false; // new OTA firmware, not yet marked as good
static bool crash_count_cleared = false;
static bool power_on_count_cleared = false;
static char status[80];

// The Arduino core marks a new OTA firmware as good right at boot. Returning true here
// leaves that to us, so the bootloader can roll back a firmware that fails early.
extern "C" bool verifyRollbackLater()
{
	return true;
}

static ResetKind resetKind()
{
	switch (esp_reset_reason())
	{
	case ESP_RST_POWERON:
		return ResetKind::PowerOn;
	case ESP_RST_SW:
		return ResetKind::Software;
	case ESP_RST_PANIC:
	case ESP_RST_INT_WDT:
	case ESP_RST_TASK_WDT:
	case ESP_RST_WDT:
		return ResetKind::Crash;
	default:
		return ResetKind::Other;
	}
}

static void setCrashCount(uint8_t count)
{
	rtc_magic = RTC_MAGIC;
	rtc_crash_count = count;
}

static const char *firmwareState()
{
	esp_ota_img_states_t state;
	if (esp_ota_get_state_partition(esp_ota_get_running_partition(), &state) != ESP_OK)
		return "usb";
	switch (state)
	{
	case ESP_OTA_IMG_PENDING_VERIFY:
		return "testing";
	case ESP_OTA_IMG_VALID:
		return "ok";
	case ESP_OTA_IMG_NEW:
		// The bootloader turns NEW into PENDING_VERIFY when it supports rollback
		return "ok, no rollback in bootloader";
	default:
		return "usb";
	}
}

static void updateStatus()
{
	snprintf(status, sizeof(status), "%s, %s, %u crashes, %s", FIRMWARE_VERSION, firmwareState(),
			 rtc_magic == RTC_MAGIC ? rtc_crash_count : 0,
			 recoveryAvailable() ? "recovery ok" : "no recovery app");
}

// Switching the power on and off quickly a few times starts the recovery app. Checked
// before anything else, so it works even when the rest of this firmware is broken.
static void checkPowerCycles(ResetKind reset)
{
	if (reset != ResetKind::PowerOn)
		return;
	uint8_t count = loadPowerOnCount() + 1;
	if (powerCyclesRequestRecovery(count) && recoveryAvailable())
	{
		savePowerOnCount(0);
		rebootToRecovery(RecoveryReason::PowerCycle);
	}
	savePowerOnCount(count);
}

BootMode firmwareGuardBegin()
{
	ResetKind reset = resetKind();
	checkPowerCycles(reset);

	uint8_t count = rtc_magic == RTC_MAGIC ? rtc_crash_count : 0;
	BootDecision decision = decideBoot(reset, count);
	setCrashCount(decision.crash_count);

	// Tells the recovery app which slot to start and which version is in it
	const char *label = runningPartitionLabel();
	if (strncmp(label, "ota_", 4) == 0)
		saveMainSlot(label, FIRMWARE_VERSION);

	esp_ota_img_states_t state;
	firmware_pending = esp_ota_get_state_partition(esp_ota_get_running_partition(), &state) == ESP_OK &&
					   state == ESP_OTA_IMG_PENDING_VERIFY;

	updateStatus();
	Serial.printf("Firmware %s, reset reason %d, crash count %u\n", status, esp_reset_reason(),
				  decision.crash_count);
	return decision.mode;
}

void firmwareGuardSafeModeEnded()
{
	setCrashCount(crashCountAfterSafeMode());
}

void firmwareGuardStartWatchdog()
{
	// Reconfigures the watchdog the ESP-IDF already started (default 5 s)
	esp_task_wdt_init(WATCHDOG_TIMEOUT_S, true);
	esp_task_wdt_add(NULL);
	watchdog_running = true;
	main_start = millis();
}

void firmwareGuardStopWatchdog()
{
	if (watchdog_running)
		esp_task_wdt_delete(NULL);
	watchdog_running = false;
}

bool firmwareGuardPending()
{
	return firmware_pending;
}

void firmwareGuardLoop(bool wifi_connected)
{
	if (watchdog_running)
		esp_task_wdt_reset();

	uint32_t running = millis() - main_start;

	if (firmware_pending && shouldMarkFirmwareValid(running, wifi_connected))
	{
		esp_ota_mark_app_valid_cancel_rollback();
		firmware_pending = false;
		updateStatus();
		Serial.println("New firmware marked as good");
	}

	// Running long enough after a power on: it was not a quick power cycle
	if (!power_on_count_cleared && millis() >= POWER_CYCLE_WINDOW)
	{
		savePowerOnCount(0);
		power_on_count_cleared = true;
	}

	if (!crash_count_cleared && shouldClearCrashCount(running))
	{
		setCrashCount(0);
		crash_count_cleared = true;
		updateStatus();
	}
}

const char *firmwareGuardResetReason()
{
	switch (esp_reset_reason())
	{
	case ESP_RST_POWERON:
		return "power on";
	case ESP_RST_EXT:
		return "reset pin";
	case ESP_RST_SW:
		return "restart";
	case ESP_RST_PANIC:
		return "crash";
	case ESP_RST_INT_WDT:
		return "watchdog (interrupt)";
	case ESP_RST_TASK_WDT:
		return "watchdog (task)";
	case ESP_RST_WDT:
		return "watchdog";
	case ESP_RST_DEEPSLEEP:
		return "deep sleep";
	case ESP_RST_BROWNOUT:
		return "brownout (supply voltage dropped)";
	case ESP_RST_SDIO:
		return "SDIO";
	default:
		return "unknown";
	}
}

const char *firmwareGuardStatus()
{
	return status;
}