// 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 and every SPM/EEPROM lock
// are libavr's to handle; the only attribute is the naked reset entry that
// stands in for the absent C runtime.

#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}>;
// One-wire: RX and TX share the line, exactly as the native-UART TSB expects.
using serial_t = dev::uart0<{.baud = 115200_Bd, .max_baud_error = 3_pct, .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;

// 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 (the LASTPAGE holding the app-jump vector, activation timeout and
// password), 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;

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

// The 16-byte device-info block the host reads on activation. A flash_table
// keeps it in progmem with no .data image (there is no crt to copy one).
// clang-format off
inline constexpr std::array<std::uint8_t, 16> info_data = {
    'T', 'S', 'B',
    build_date & 0xFF, build_date >> 8,
    0xF3,             // status byte (native-UART fixed-baud lineage)
    0x1E, 0x95, 0x0F, // ATmega328P signature
    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
using info = avr::flash_table<info_data>;

// Blocking byte read/write over the one-wire line: read() releases the line to
// the receiver, write() takes it and holds it until the frame is out.
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 `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(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)
{
	for (std::uint16_t i = 0; i < page; i += 2) {
		std::uint8_t lo = rx();
		std::uint8_t hi = rx();
		spm::fill<off>(addr + i, static_cast<std::uint16_t>(lo | (hi << 8)));
	}
	spm::write_page<off>(addr);
	spm::wait();
}

// 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>(addr + i, rx());
}

// Erase one flash page and wait it out — 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<off>(addr);
	spm::wait();
}

// Erase the whole application, one page at a time (unwritten pages stay erased).
void erase_application()
{
	for (std::uint16_t a = 0; a < app_end; a += page)
		erase_page(a);
	spm::rww_enable<off>();
}

// Run the application: reset vector at 0x0000. 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
	reinterpret_cast<void (*)()>(0)();
	__builtin_unreachable();
}

// 'f': 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);
	}
}

// 'e': 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);
	}
}

// 'F': 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);
}

// 'E': 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);
}

// 'C': 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>(a, 0xff);
	erase_page(app_end);
	spm::rww_enable<off>();
}

// The password gate. The config page holds the password at app_end+3,
// 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 + 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()
{
	// 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 three '@' 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.
	__uint24 idle = static_cast<__uint24>(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:
		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 'f':
			read_flash();
			break;
		case 'F':
			write_flash();
			break;
		case 'e':
			read_eeprom();
			break;
		case 'E':
			write_eeprom();
			break;
		case 'c':
			send_flash(app_end, page);
			break;
		case 'C':
			write_config();
			break;
		default:
			appjump(); // 'q' or any other byte runs the application
		}
	}
}

} // namespace
} // namespace tsb

// Reset lands here: BOOTRST vectors to the boot section base and .vectors is
// laid first, so this is the first instruction executed. No crt ran, so set the
// stack pointer before anything is called.
extern "C" [[gnu::naked, gnu::used, gnu::section(".vectors")]] void __boot_entry()
{
	SP = RAMEND;
	tsb::run();
}