// TinySafeBoot on libavr - the policy floor: pureboot's rules, measured.
//
// The full TinySafeBoot feature set - watchdog bail, one-wire half-duplex,
// config-page activation timeout, password gate, emergency erase, and
// config/flash/EEPROM read-write - under philosophy #5 exactly as pureboot
// obeys it: no assembly, no register variables; code, attributes, and flags
// only. Every lesson pureboot's development produced is applied - the
// library's half-duplex serial and startup entry, lean bring-up from reset
// state, one merged send loop over both memories, oracle-shaped loop bounds,
// locals threaded through noinline primitives, pureboot's codegen flags  -
// and the result sits below the idiomatic tier and above the 512 B boot
// section the tricks/asm tiers reach with the banned mechanisms
// (oracle/README.md holds all four). This tier exists to keep that gap an
// artifact
// rather than a claim: the gap to 512 is the rent of policy-clean C++  -
// helpers that hold a cursor across rx()/tx() pay push/pop and argument
// threading where a global-register protocol pays nothing, and both
// control-flow merges tried (a parametrized paged session, a merged store
// loop) measured larger than the split cases they replaced. TSB's wire fixes
// the per-command loop shapes on the device, so pureboot 5's one-transfer-
// loop collapse has no purchase here.
//
// The wire protocol is strict request/response, which is what makes the
// shared line safe: the device drives it only between a received command and
// its reply, and releases it (the library's half-duplex choreography)
// whenever it waits.

#include <libavr/libavr.hpp>

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

using dev = avr::device<{.clock = 16_MHz}>;
// One-wire: RX and TX share the line, exactly as the native-UART TSB expects.
// 115200 at 16 MHz lands +2.1 % off, past the receiver-tolerance table the
// solver holds rates to - the oracle's own deployment has run there for a
// decade, so the override states that it is meant.
using serial_t = dev::uart0<{
    .baud = 115200_Bd,
    .allow_baud_error = true,
    .half_duplex = true,
}>;
inline constexpr serial_t serial{};

namespace tsb {
namespace {

// The loader is purely polled - it never enables interrupts - so every SPM and
// EEPROM lock folds to nothing under this posture.
constexpr auto off = avr::irq::guard_policy::unused;

// Strict request/response: every SPM operation is waited out before the next
// byte moves, so no flash operation is ever in flight at an EEPROM access -
// the write procedure's step 2 has nothing to guard, the omission the
// datasheet grants (DS40002061B section 8.6.3).
constexpr auto no_spm = ee::spm_interlock::omitted;

// The handshake bytes, identical across every TSB host.
constexpr std::uint8_t confirm = '!';
constexpr std::uint8_t request = '?';
constexpr std::uint8_t knock = '@';

// Boot geometry for the 1 KB boot section (BOOTSZ=10); the page size and the
// flash/EEPROM extents are the chip database's to know. app_end is the config
// page (TSB's LASTPAGE), one page below the boot section.
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;

// Lockout-proof floor for the activation window: the oracle's F_CPU/1MHz, so
// it follows the clock rather than restating it (rule 41).
constexpr auto act_min = static_cast<std::uint8_t>(dev::clock.hz / 1'000'000);
// Post-activation window: the host gets seconds, not milliseconds, mid-session.
constexpr std::uint8_t comm_window = 200;

// Firmware version stamp: YY*512 + MM*32 + DD, the encoding the host decodes.
constexpr std::uint16_t build_date = 26 * 512 + 7 * 32 + 27;

// The 16-byte device-info block, streamed out on activation.
// clang-format off
[[gnu::progmem]] constexpr auto info = std::to_array<std::uint8_t>({
    'T', 'S', 'B',
    build_date & 0xFF, build_date >> 8,
    0xF3,             // status: native-UART fixed-baud lineage
    avr::hw::db.signature[0], avr::hw::db.signature[1], avr::hw::db.signature[2],
    page / 2,         // page size in words
    (app_end / 2) & 0xFF, (app_end / 2) >> 8, // app-flash boundary, words
    eeprom_end & 0xFF, eeprom_end >> 8,
    0xAA, 0xAA,       // ATmega processor-type marker (bytes 14 == 15)
});
// clang-format on

// The receive window, pre-floored where it is set. In .noinit: there is no
// crt to clear a .bss image, and run() stores it before the first receive.
[[gnu::section(".noinit")]] std::uint8_t window;

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

// Bounded byte receive: poll under nested countdowns, 0 on silence. The 0
// then falls through every compare - not a knock, not a confirm, not a
// command - so a silent host unwinds the loader to the application from
// anywhere, and a mid-session cable pull cannot wedge it. The line release on
// a direction change is the serial backend's.
[[gnu::noinline]] std::uint8_t rx()
{
	std::uint16_t outer = static_cast<std::uint16_t>(window) << 8;
	do {
		std::uint8_t fine = 0;
		do {
			if (auto byte = serial.read()) {
				return *byte;
			}
		} while (--fine);
	} while (--outer);
	return 0;
}

// One-wire transmit: the backend takes the line with a turn-around guard and
// holds it until the whole frame is out.
[[gnu::noinline]] void tx(std::uint8_t byte)
{
	serial.write(byte);
}

// '?', then hand back the host's reply for the callers' one-byte compare.
[[gnu::noinline]] std::uint8_t rcnf()
{
	tx(request);
	return rx();
}

// The one send loop: the info block, the config page, application flash and
// EEPROM pages all stream through here.
[[gnu::noinline]] void send_block(bool eep, std::uint16_t at, std::uint8_t count)
{
	do {
		tx(eep ? ee::read<no_spm>(at) : avr::flash_load(flash_ptr(at)));
		++at;
	} while (--count);
}

// One EEPROM byte in - shared by the emergency wipe and the 'E' stream.
[[gnu::noinline]] void eeput(std::uint16_t at, std::uint8_t value)
{
	ee::write<off, no_spm>(at, value);
}

// Wait out a running SPM op, then re-open the RWW section - after every page
// op and before handing over, as the oracle does.
[[gnu::noinline]] void settle()
{
	spm::wait();
	spm::rww_enable<off>();
}

// One host page straight into the erased flash page at `at` - through the SPM
// word buffer (low byte then high), no SRAM staging - then committed. `at`
// names a page base, so the cursor's low byte reaching the boundary ends the
// walk.
[[gnu::noinline]] void store_flash_page(std::uint16_t at)
{
	const auto open = spm::page::begin<spm::from::boot_section, off>(at);
	do {
		std::uint8_t low = rx();
		std::uint8_t high = rx();
		spm::fill<off>(open, at, std::bit_cast<std::uint16_t>(std::array{low, high}));
		at += 2;
	} while (static_cast<std::uint8_t>(at) & (page - 1));
	spm::command<off>(spm::op::write, at - page);
	settle();
}

extern "C" [[noreturn]] void tsb_app(); // the application's reset vector: --defsym=tsb_app=0

[[noreturn]] void appjump()
{
	settle();
	tsb_app();
}

// Step one page down and erase it - the erase shared by the whole-app walk,
// the config rewrite and the emergency wipe; hands the stepped address back.
[[gnu::noinline]] std::uint16_t erase_below(std::uint16_t at)
{
	at -= page;
	spm::command<off>(spm::op::erase, at);
	settle();
	return at;
}

// Erase the whole application, top-down like the oracle: the loop bound is a
// compare with zero, and the returned 0 is the address every caller wants
// next.
[[gnu::noinline]] std::uint16_t erase_application()
{
	std::uint16_t at = app_end;
	do {
		at = erase_below(at);
	} while (at != 0);
	return at;
}

[[noreturn]] void run()
{
	// A watchdog reset hands straight back to the application, as the
	// reference loader does, rather than re-entering the bootloader.
	if (avr::hw::mcusr::wdrf.test()) {
		appjump();
	}

	// Lean bring-up from reset state: UCSR0C already reads 8N1, UBRR0H reads
	// 0, and the half-duplex write()/read() raise TXEN0/RXEN0 on first use  -
	// only the divisor low byte and U2X0 need a store. The solver still does
	// the datasheet work; the asserts pin the reset-state assumptions.
	{
		constexpr auto sol = avr::uart::solve_baud(dev::clock, 115200_Bd, 8, avr::uart::parity::none);
		static_assert(sol.u2x && sol.ubrr < 256, "lean bring-up writes UBRR0L only, with U2X0");
		avr::hw::ubrr0::write(static_cast<std::uint8_t>(sol.ubrr));
		avr::hw::ucsr0a::write(avr::hw::ucsr0a::u2x0(1));
	}

	// Activation: 3x'@', each inside the config page's timeout window
	// (floored so a corrupt page cannot lock the loader out); anything else  -
	// including silence - hands over.
	window = avr::flash_load(flash_ptr(app_end + 2)) | act_min;
	for (std::uint8_t k = 3; k; --k) {
		if (rx() != knock) {
			appjump();
		}
	}
	window = comm_window;

	// Password gate (config page from app_end+3, 0xff-terminated; a blank
	// page is no password). A wrong byte blanks the comparison and drains the
	// line forever, so a wrong password can never fall through; a 0 requests
	// emergency erase behind two confirms. On pass the info block goes out;
	// the emergency path skips it and drops into the command loop.
	std::uint16_t at = app_end + 3;
	std::uint8_t mask = 0xff;
	for (;;) {
		std::uint8_t expected = avr::flash_load(flash_ptr(at)) & mask;
		++at;
		if (expected == 0xff) {
			send_block(false, reinterpret_cast<std::uint16_t>(info.data()), info.size());
			break;
		}
		std::uint8_t got = rx();
		if (got == 0) {
			if (mask == 0) {
				continue;
			}
			if (rcnf() != confirm || rcnf() != confirm) {
				appjump();
			}
			std::uint16_t a = erase_application();
			do {
				eeput(a, 0xff);
			} while (++a <= eeprom_end);
			erase_below(app_end + page);
			break;
		}
		if (got != expected) {
			mask = 0;
		}
	}

	for (;;) {
		tx(confirm); // Mainloop ready
		const std::uint8_t command = rx();
		switch (command) {
		case 'f': // read application flash, one page per host '!'
			for (std::uint16_t a = 0; a < app_end; a += page) {
				if (rx() != confirm) {
					break;
				}
				send_block(false, a, page);
			}
			break;
		case 'e': // read EEPROM, one page per host '!', until the host stops
			for (std::uint16_t a = 0;; a += page) {
				if (rx() != confirm) {
					break;
				}
				send_block(true, a, page);
			}
			break;
		case 'F': { // erase the application, then take pages behind '?'
			std::uint16_t a = erase_application();
			for (; rcnf() == confirm; a += page) {
				store_flash_page(a);
			}
			break;
		}
		case 'E': // take EEPROM pages behind '?', each write host-paced
			for (std::uint16_t a = 0; rcnf() == confirm;) {
				std::uint8_t count = page;
				do {
					eeput(a, rx());
					++a;
				} while (--count);
			}
			break;
		case 'c': // read the config page
		read_config:
			send_block(false, app_end, page);
			break;
		case 'C': // replace the config page, then echo it back to verify
			if (rcnf() != confirm) {
				break;
			}
			store_flash_page(erase_below(app_end + page));
			goto read_config;
		default: // 'q' or any other byte runs the application
			appjump();
		}
	}
}

} // namespace
} // namespace tsb

// Reset lands at the boot section base (BOOTRST): the entry stub in .vectors
// is laid first and does the one line of crt a crt-less image needs.
template struct avr::startup::entry<tsb::run, avr::startup::stack::hardware>;