// TinySafeBoot on libavr — tier 2: C++ with compiler trickery.
//
// Same protocol, libavr surface and full feature set as the pure variant
// (tsb_pure.cpp) — watchdog bail, one-wire, config-page timeout, password gate,
// emergency erase, config/flash/EEPROM read-write — but the readable
// one-handler-per-command shape is traded for size. Flash and EEPROM share a
// single code path selected by a *runtime* flag decoded from the command byte,
// so the compiler cannot constant-propagate it into two clones; attributes
// (noinline/noclone) pin that sharing down; the hot page address and byte
// counter live in call-saved global registers to erase the prologue push/pop
// that C++ function decomposition otherwise pays; and pages stream straight to
// SPM/EEPROM with no SRAM staging. No inline assembly.

#include <libavr/libavr.hpp>

#include <avr/io.h> // SP / RAMEND for the crt-free boot entry

using namespace avr::literals;
namespace spm = avr::spm;
namespace ee = avr::eeprom;

using dev = avr::device<{.clock = 16_MHz}>;
using serial_t = dev::uart0<{.baud = 115200_Bd, .max_baud_error = 3_pct, .half_duplex = true}>;
inline constexpr serial_t serial{};

namespace tsb {

constexpr auto off = avr::irq::guard_policy::unused;

constexpr std::uint8_t confirm = '!';
constexpr std::uint8_t request = '?';
constexpr std::uint8_t knock = '@';

constexpr std::uint16_t page = spm::page_bytes;
constexpr std::uint16_t boot_bytes = 1024;
constexpr std::uint16_t app_end = spm::flash_bytes - boot_bytes - page;
constexpr std::uint16_t eeprom_end = avr::hw::db.mem.eeprom_size - 1;

constexpr std::uint16_t build_date = 26 * 512 + 7 * 32 + 19;

// clang-format off
[[gnu::progmem]] constexpr std::uint8_t info[16] = {
    'T', 'S', 'B',
    build_date & 0xFF, build_date >> 8,
    0xF3,
    0x1E, 0x95, 0x0F,
    page / 2,
    (app_end / 2) & 0xFF, (app_end / 2) >> 8,
    eeprom_end & 0xFF, eeprom_end >> 8,
    0xAA, 0xAA,
};
// clang-format on

// The hot page walk lives in call-saved global registers, TSB-style: g_addr is
// the running flash/EEPROM byte address, g_cnt the byte countdown. Being global
// they are never spilled around the rx/tx/spm calls the way a local would be —
// r4-r7 are call-saved, so the library's UART and SPM helpers preserve them.
register std::uint16_t g_addr asm("r4");
register std::uint8_t g_cnt asm("r6");

std::uint8_t rx()
{
	return serial.read_blocking();
}

void tx(std::uint8_t byte)
{
	serial.write(byte);
}

const std::uint8_t *flash_ptr(std::uint16_t addr)
{
	return reinterpret_cast<const std::uint8_t *>(addr);
}

// Stream g_cnt bytes to the host from flash (LPM) or EEPROM, memory chosen at
// run time so the optimiser cannot split the loop into two clones.
[[gnu::noinline, gnu::noclone]] void send(bool flash)
{
	do {
		tx(flash ? avr::flash_load(flash_ptr(g_addr)) : ee::read(g_addr));
		++g_addr;
	} while (--g_cnt);
}

[[gnu::noinline]] bool request_confirm()
{
	tx(request);
	return rx() == confirm;
}

// Stream one page straight from the host into the already-erased flash page at
// g_addr (SPM word buffer, low byte then high) or into EEPROM — no SRAM staging,
// so receive and store are one loop. The memory is a run-time flag.
[[gnu::noinline, gnu::noclone]] void store_page(bool flash)
{
	g_cnt = 0;
	if (flash) {
		do {
			std::uint8_t lo = rx();
			std::uint8_t hi = rx();
			spm::fill<off>(g_addr + g_cnt, static_cast<std::uint16_t>(lo | (hi << 8)));
			g_cnt += 2;
		} while (g_cnt != page);
		spm::write_page<off>(g_addr);
		spm::wait();
	} else {
		do {
			ee::write<off>(g_addr + g_cnt, rx());
		} while (++g_cnt != page);
	}
}

[[noreturn]] void appjump()
{
	spm::wait();
	reinterpret_cast<void (*)()>(0)();
	__builtin_unreachable();
}

// One flash page erased and waited out. Shared, so the SPM erase command
// sequence is emitted once for the whole-app erase, the config page and the
// emergency wipe rather than three times.
[[gnu::noinline]] void erase1(std::uint16_t addr)
{
	spm::erase_page<off>(addr);
	spm::wait();
}

// Re-enable RWW flash reads after programming, shared for the same reason.
[[gnu::noinline]] void rww()
{
	spm::rww_enable<off>();
}

// Erase the whole application, one page at a time.
[[gnu::noinline]] void erase_application()
{
	g_addr = 0;
	do {
		erase1(g_addr);
		g_addr += page;
	} while (g_addr < app_end);
	rww();
}

// 'f'/'e': stream memory back one page per host '!'. send advances g_addr, so
// flash self-terminates at the application boundary; EEPROM runs until the host
// stops.
[[gnu::noinline]] void read_mem(bool flash)
{
	g_addr = 0;
	for (;;) {
		if (rx() != confirm)
			return;
		g_cnt = page;
		send(flash);
		if (flash && g_addr >= app_end)
			return;
	}
}

// 'F'/'E': flash erases the whole application first, then both take the pages
// the host offers behind '?'.
[[gnu::noinline]] void write_mem(bool flash)
{
	if (flash)
		erase_application();
	g_addr = 0;
	while (request_confirm()) {
		store_page(flash);
		g_addr += page;
	}
}

// 'C': replace the config page, then echo it back for the host to verify.
void write_config()
{
	if (!request_confirm())
		return;
	g_addr = app_end;
	erase1(g_addr);
	store_page(true);
	rww();
	g_addr = app_end;
	g_cnt = page;
	send(true);
}

// Emergency erase: wipe the application flash, the EEPROM and the config page.
[[gnu::noinline]] void emergency_erase()
{
	erase_application();
	g_addr = 0;
	do {
		ee::write<off>(g_addr, 0xff);
	} while (++g_addr <= eeprom_end);
	erase1(app_end);
	rww();
}

// The password gate. A byte of 0 requests emergency erase; a wrong byte hangs
// the loader (still draining the line), so it can never fall through to erase.
enum class gate : std::uint8_t { pass, emergency };

[[gnu::noinline]] gate password_gate()
{
	for (const std::uint8_t *pw = flash_ptr(app_end + 3);; ++pw) {
		std::uint8_t expected = avr::flash_load(pw);
		if (expected == 0xff)
			return gate::pass;
		std::uint8_t got = rx();
		if (got == 0)
			return gate::emergency;
		if (got != expected)
			for (;;)
				rx();
	}
}

[[noreturn]] void run()
{
	if (avr::hw::mcusr::wdrf.test())
		appjump();

	avr::init<serial_t>();

	std::uint32_t idle = static_cast<std::uint32_t>(avr::flash_load(flash_ptr(app_end + 2)) | 16) << 16;
	std::uint8_t knocks = 0;
	while (knocks < 3) {
		if (auto byte = serial.read())
			knocks = *byte == knock ? knocks + 1 : 0;
		else if (--idle == 0)
			appjump();
	}

	switch (password_gate()) {
	case gate::pass:
		g_addr = reinterpret_cast<std::uint16_t>(&info[0]);
		g_cnt = sizeof(info);
		send(true);
		break;
	case gate::emergency:
		if (!request_confirm() || !request_confirm())
			appjump();
		emergency_erase();
		break;
	}

	for (;;) {
		tx(confirm); // Mainloop ready
		// Decode the command arithmetically so flash/write stay run-time values:
		// bit 5 is the case bit (upper = write), the folded-lower letter picks the
		// memory. A single unified path serves f/F/e/E.
		std::uint8_t cmd = rx();
		std::uint8_t lower = cmd | 0x20;
		bool write = (cmd & 0x20) == 0;
		if (lower == 'f' || lower == 'e') {
			bool flash = lower == 'f';
			if (write)
				write_mem(flash);
			else
				read_mem(flash);
		} else if (lower == 'c') {
			if (write) {
				write_config();
			} else {
				g_addr = app_end;
				g_cnt = page;
				send(true);
			}
		} else {
			appjump();
		}
	}
}

} // namespace tsb

extern "C" [[gnu::naked, gnu::used, gnu::section(".vectors")]] void __boot_entry()
{
	SP = RAMEND;
	tsb::run();
}