// TinySafeBoot on libavr - tier 1: pure, idiomatic C++.
//
// A serial flash bootloader for the ATmega328P boot section, reimplementing the
// TinySafeBoot native-UART fixed-baud protocol on libavr with the full feature
// set of the hand-written oracle: a watchdog-reset bail, one-wire half-duplex,
// a config-page activation timeout, the password gate, emergency erase, and
// config/flash/EEPROM read-write. This variant is written for clarity -
// well-factored functions, no compiler-specific size hacks, no inline assembly.
// The one-wire wiring, the flash-resident info block, every SPM/EEPROM lock and
// the crt-free reset entry are libavr's to handle.

#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.
using serial_t = dev::uart0<{
    .baud = tsb::baud_rate,
    .allow_baud_error = true,
    .half_duplex = true,
}>;
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);

// A flash_table keeps the info block in progmem with no .data image (there is
// no crt to copy one).
inline constexpr auto info_data = info_block(date_stamp(26, 7, 20), app_end);
using info = avr::flash_table<info_data>;

// 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 read over the one-wire line - read() releases the line to the
// receiver - answering 0 on silence. That 0 falls through every compare below:
// 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 oracle lists that timeout among its own fixes, and a blocking
// read is how a tier loses it.
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;
}

// write() takes the line and holds it until the frame is out.
void tx(std::uint8_t byte)
{
	serial.write(byte);
}

// Stream `count` bytes to the host, from flash (LPM) or from EEPROM.
void send_flash(std::uint16_t addr, std::uint8_t count)
{
	while (count--) {
		tx(avr::flash_load(flash_ptr(addr++)));
	}
}

void send_eeprom(std::uint16_t addr, std::uint8_t count)
{
	while (count--) {
		tx(ee::read<no_spm>(addr++));
	}
}

// Prompt the host with '?' and report whether it answered '!'.
bool request_confirm()
{
	tx(request);
	return rx() == confirm;
}

// Stream one page from the host straight into the already-erased flash page at
// `addr`, filling the SPM word buffer low byte then high - no SRAM staging, so
// receiving and programming are the same loop.
void store_flash_page(std::uint16_t addr)
{
	const auto open = spm::page::begin<spm::from::boot_section, off>(addr);
	for (std::uint16_t i = 0; i < page; i += 2) {
		const auto lo = rx();
		const auto hi = rx();
		spm::fill<off>(open, addr + i, static_cast<std::uint16_t>(lo | (hi << 8)));
	}
	spm::write_page<spm::from::boot_section, off>(addr); // blocking: waits the write out
}

// Stream one page from the host straight into EEPROM, byte by byte.
void store_eeprom_page(std::uint16_t addr)
{
	for (std::uint16_t i = 0; i < page; ++i) {
		ee::write<off, no_spm>(addr + i, rx());
	}
}

// Erase one flash page, waited out by the blocking spelling - the erase step
// shared by the whole-app erase, the config-page rewrite and the emergency
// wipe.
void erase_page(std::uint16_t addr)
{
	spm::erase_page<spm::from::boot_section, off>(addr);
}

// Erase the whole application, one page at a time, top-down as the reference
// loader does (unwritten pages stay erased and the host cannot observe the
// order; the loop bound becomes a compare with zero).
void erase_application()
{
	for (std::uint16_t a = app_end; a != 0;) {
		a -= page;
		erase_page(a);
	}
	spm::rww_enable<off>();
}

// The application's reset vector; the linker pins it to 0x0000 (--defsym).
extern "C" [[noreturn]] void tsb_app();

// Run the application. Any non-command byte, a wrong password, or an idle
// programmer port lands here.
[[noreturn]] void appjump()
{
	spm::wait(); // make sure any pending SPM finished before handing over
	tsb_app();
}

// Stream the application flash back, one page per host '!'. Self-terminates at
// the application boundary; the host normally stops earlier with a non-'!'.
void read_flash()
{
	for (std::uint16_t a = 0; a < app_end; a += page) {
		if (rx() != confirm) {
			return;
		}
		send_flash(a, page);
	}
}

// Stream EEPROM back, one page per host '!', until the host stops.
void read_eeprom()
{
	for (std::uint16_t a = 0;; a += page) {
		if (rx() != confirm) {
			return;
		}
		send_eeprom(a, page);
	}
}

// Erase the whole application first, then take pages the host offers behind
// '?'.
void write_flash()
{
	erase_application();
	for (std::uint16_t a = 0; request_confirm(); a += page) {
		store_flash_page(a);
	}
}

// Take pages the host offers behind '?' into EEPROM.
void write_eeprom()
{
	for (std::uint16_t a = 0; request_confirm(); a += page) {
		store_eeprom_page(a);
	}
}

// Replace the config page, then echo it back for the host to verify.
void write_config()
{
	if (!request_confirm()) {
		return;
	}
	erase_page(app_end);
	store_flash_page(app_end);
	spm::rww_enable<off>();
	send_flash(app_end, page);
}

// Emergency erase: wipe the application flash, the EEPROM and the config page.
// Reachable only from the password gate (a wrong byte can never reach it), so a
// blank config still leaves the loader recoverable.
void emergency_erase()
{
	erase_application();
	for (std::uint16_t a = 0; a <= eeprom_end; ++a) {
		ee::write<off, no_spm>(a, 0xff);
	}
	erase_page(app_end);
	spm::rww_enable<off>();
}

// The password gate. The config page holds the password at password_at,
// terminated by 0xff (a blank page means no password). A byte of 0 requests
// emergency erase; a wrong byte hangs the loader, still draining the line, so a
// wrong password can never fall through to the erase.
enum class gate : std::uint8_t { pass, emergency };

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

[[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();
	}

	avr::init<serial_t>();

	// Activation: the host knocks activation_knocks times inside a window whose length is the
	// config page's timeout byte, floored so a corrupt page can never lock the
	// loader out. An idle port times out and boots the application; the same
	// window then bounds every receive of the session.
	window = avr::flash_load(flash_ptr(app_end + timeout_at)) | act_min;
	__uint24 idle = static_cast<__uint24>(window) << 16;
	std::uint8_t knocks = 0;
	while (knocks < activation_knocks) {
		if (auto byte = serial.read()) {
			knocks = *byte == knock ? knocks + 1 : 0;
		} else if (--idle == 0) {
			appjump();
		}
	}

	switch (password_gate()) {
	case gate::pass:
		send_flash(reinterpret_cast<std::uint16_t>(info::storage.data()), info::size());
		break;
	case gate::emergency:
		if (!request_confirm() || !request_confirm()) {
			appjump();
		}
		emergency_erase();
		break;
	}

	for (;;) {
		tx(confirm); // Mainloop ready
		switch (rx()) {
		case command::read_flash:
			read_flash();
			break;
		case command::write_flash:
			write_flash();
			break;
		case command::read_eeprom:
			read_eeprom();
			break;
		case command::write_eeprom:
			write_eeprom();
			break;
		case command::read_config:
			send_flash(app_end, page);
			break;
		case command::write_config:
			write_config();
			break;
		default:
			appjump(); // 'q' or any other byte runs the application
		}
	}
}

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