// 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 - with no assembly and no register variables:
// code, attributes and flags only. It applies pureboot's size techniques - 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
// lands between the idiomatic tier and the 512 B section the tricks/asm tiers
// reach with the banned mechanisms (oracle/README.md holds all four). 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.
//
// 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 "tsb.hpp"

using dev = avr::device<{.clock = tsb::clock}>;
// 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.
constexpr avr::uart::config link{
    .baud = tsb::baud_rate,
    .allow_baud_error = true,
    .half_duplex = true,
};
using serial_t = dev::uart0<link>;
inline constexpr serial_t serial{};

namespace tsb {
namespace {

// The boot section: 1 KB, BOOTSZ=10.
constexpr std::uint16_t boot_bytes = 1024;
constexpr std::uint16_t app_end = config_page(boot_bytes);

[[gnu::progmem]] constexpr auto info = info_block(date_stamp(26, 7, 27), app_end);

// 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;

// 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 {
		const auto low = rx();
		const auto 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, link.baud, link.data_bits, link.par);
		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: activation_knocks '@', 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 + timeout_at)) | act_min;
	for (std::uint8_t k = activation_knocks; k; --k) {
		if (rx() != knock) {
			appjump();
		}
	}
	window = comm_window;

	// Password gate (config page from password_at, 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 + password_at;
	std::uint8_t mask = 0xff;
	for (;;) {
		const 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;
		}
		const auto 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
		switch (rx()) {
		case command::read_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 command::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 command::write_flash: { // erase the application, then take pages behind '?'
			std::uint16_t a = erase_application();
			for (; rcnf() == confirm; a += page) {
				store_flash_page(a);
			}
			break;
		}
		case command::write_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 command::read_config:
		read_config:
			send_block(false, app_end, page);
			break;
		case command::write_config: // 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>;