// TinySafeBoot on libavr — tier 2: C++ with compiler trickery, no assembly.
//
// 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 — in pure C++, 526 bytes: 14 over the 512-byte
// boot section the hand-written oracle fits, from 168 over at this tier's first
// floor. The structure mirrors the oracle's: a handful of tiny noinline
// primitives sharing one whole-loader register allocation, expressed as global
// register variables so no helper ever saves, spills, or reloads any of it.
//
// The register protocol (all call-saved, so calls preserve them by ABI):
//   Y (r28:r29)  g_addr      the walked flash/EEPROM address — adiw-able
//   r16          g_cnt       byte countdown of the running block — ldi-able
//   r7           g_window    rx timeout, roughly 30 ms units at 16 MHz
//   r6           g_receiving one-wire direction latch, cleared at bring-up
//                            (power-on registers are undefined)
//
// GCC 16.1 miscompiles stores into global register variables: an update whose
// remaining uses all hide inside callees is deleted whenever a CALL follows it
// before any jump/ret (the backend's liveness walk lumps fixed registers with
// call-clobbered ones — minimal repro in libavr's
// local/scratch/probes/gcc-avr-globalreg-repro.cpp, lessons.md entry). Every
// g_* update below therefore sits where a *local* read or a jump/ret follows
// it — the helpers advance g_addr immediately before returning, and rx()
// re-floors the window on every call instead of storing the floored value
// once. The layout is load-bearing; do not "simplify" it.
//
// 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 (RXEN0 only) whenever it waits.

#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;
namespace hw = avr::hw;

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;

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).
constexpr std::uint8_t act_min = 16;
// Post-activation window: the host gets seconds, not milliseconds, mid-session.
constexpr std::uint8_t comm_window = 200;

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

// Fixed 115200 8N1; the library solves UBRR + U2X from clock and baud.
constexpr auto baud = avr::uart::detail::solve_baud(16_MHz, 115200_Bd);

// The 16-byte device-info block, streamed out on activation.
// clang-format off
[[gnu::progmem]] constexpr std::uint8_t info[16] = {
    'T', 'S', 'B',
    build_date & 0xFF, build_date >> 8,
    0xF3,             // status: native-UART fixed-baud lineage
    0x1E, 0x95, 0x0F, // ATmega328P signature
    page / 2,         // page size in words
    (app_end / 2) & 0xFF, (app_end / 2) >> 8,
    eeprom_end & 0xFF, eeprom_end >> 8,
    0xAA, 0xAA,
};
// clang-format on

register std::uint16_t g_addr asm("r28");
register std::uint8_t g_cnt asm("r16");
register std::uint8_t g_window asm("r7");
register std::uint8_t g_receiving asm("r6");

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

// Bounded byte receive, the oracle's shape: release the one-wire line on a
// direction change, poll RXC0 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.
[[gnu::noinline, gnu::noclone]] std::uint8_t rx()
{
	if (!g_receiving) {
		g_receiving = 1;
		hw::ucsr0b::write(hw::ucsr0b::rxen0(1)); // RXEN0 alone: release and listen
	}
	// act_min ORs in here, per call, not once into g_window at setup — the
	// one placement the global-register-store miscompile cannot delete.
	std::uint16_t outer = static_cast<std::uint16_t>(g_window | act_min) << 8;
	do {
		std::uint8_t fine = 0;
		do {
			auto status = hw::ucsr0a::read();
			if (status & hw::ucsr0a::rxc0(1).value)
				return hw::udr0::read();
		} while (--fine);
	} while (--outer);
	return 0;
}

// One-wire transmit: take the line (TXEN0 alone — the receiver must be off
// while driving) on a direction change, with a turn-around guard so a shorted
// peer can switch first; then hold the line until the whole frame is out
// (TXC0, not UDRE0 — the stop bit must be on the wire before a caller may
// release the line), and W1C TXC0 by storing the sampled status back, which
// keeps U2X0.
[[gnu::noinline, gnu::noclone]] void tx(std::uint8_t byte)
{
	if (g_receiving) {
		g_receiving = 0;
		hw::ucsr0b::write(hw::ucsr0b::txen0(1));
		for (std::uint8_t guard = 46; guard; --guard)
			;
	}
	hw::udr0::write(byte);
	std::uint8_t status;
	do {
		status = hw::ucsr0a::read();
	} while (!(status & hw::ucsr0a::txc0(1).value));
	hw::ucsr0a::write(status);
}

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

// One flash byte ← [g_addr++] (the advance right before ret — see header).
[[gnu::noinline, gnu::noclone]] std::uint8_t sflash()
{
	std::uint8_t byte = avr::flash_load(flash_ptr(g_addr));
	++g_addr;
	return byte;
}

// One EEPROM byte ← [g_addr++].
[[gnu::noinline, gnu::noclone]] std::uint8_t eerd()
{
	std::uint8_t byte = ee::read(g_addr);
	++g_addr;
	return byte;
}

// One EEPROM byte → [g_addr++].
[[gnu::noinline, gnu::noclone]] void eewr(std::uint8_t byte)
{
	ee::write<off>(g_addr, byte);
	++g_addr;
}

// Stream g_cnt flash bytes from g_addr to the host.
[[gnu::noinline, gnu::noclone]] void sendf()
{
	do {
		tx(sflash());
	} while (--g_cnt);
}

// 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, gnu::noclone]] void settle()
{
	spm::wait();
	spm::rww_enable<off>();
}

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

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

// Step g_addr one page down and erase that page. The decrement lives in here,
// before the erase's own use of it, not in the caller's loop where a following
// call would get it deleted (see header).
[[gnu::noinline, gnu::noclone]] void erase_below()
{
	g_addr -= page;
	spm::erase_page<off>(g_addr);
	settle();
}

// Erase the whole application, top-down like the oracle: the loop bound is a
// compare with zero, and g_addr = 0 — the value every caller wants next — is
// handed back for free.
[[gnu::noinline, gnu::noclone]] void erase_application()
{
	g_addr = app_end;
	do {
		erase_below();
	} while (g_addr != 0);
}

// Stream one host page into the erased flash page at g_addr (SPM word buffer,
// low byte then high) — no SRAM staging, receive and program are one loop.
// g_addr is left at the next page base.
[[gnu::noinline, gnu::noclone]] void store_flash()
{
	g_cnt = page / 2;
	do {
		std::uint16_t word = rx();
		word |= static_cast<std::uint16_t>(rx()) << 8;
		spm::fill<off>(g_addr, word);
		g_addr += 2;
	} while (--g_cnt);
	spm::write_page<off>(g_addr - page);
	settle();
}

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

	// Lean bring-up from reset state: UCSR0C already reads 8N1, UBRR0H reads
	// 0, and rx()/tx() raise RXEN0/TXEN0 on first use — only the divisor low
	// byte and U2X0 need a store. The library still does the datasheet work.
	static_assert(baud.u2x && baud.ubrr < 256, "lean bring-up writes UBRR0L only, with U2X0");
	hw::reg<"UBRR0">::write(static_cast<std::uint8_t>(baud.ubrr));
	hw::ucsr0a::write(hw::ucsr0a::u2x0(1));
	// General-purpose registers are undefined at power-on (no crt zeroes them);
	// the direction latch must start "not receiving" so the first rx() enables
	// the receiver. The reference loader clears its shadow register for the
	// same reason.
	g_receiving = 0;

	// Activation: 3×'@', each inside the config page's timeout window (rx
	// floors it so a corrupt page cannot lock the loader out); anything else —
	// including silence — hands over.
	g_window = avr::flash_load(flash_ptr(app_end + 2));
	for (std::uint8_t k = 3; k; --k)
		if (rx() != knock)
			appjump();
	g_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.
	g_addr = app_end + 3;
	std::uint8_t mask = 0xff;
	for (;;) {
		std::uint8_t expected = avr::flash_load(flash_ptr(g_addr)) & mask;
		++g_addr;
		if (expected == 0xff) {
			g_addr = reinterpret_cast<std::uint16_t>(&info[0]);
			g_cnt = sizeof(info);
			sendf();
			break;
		}
		std::uint8_t got = rx();
		if (got == 0) {
			if (mask == 0)
				continue;
			if (rcnf() != confirm || rcnf() != confirm)
				appjump();
			erase_application(); // leaves g_addr = 0 for the EEPROM walk
			do {
				eewr(0xff);
			} while (g_addr <= eeprom_end);
			g_addr = app_end + page;
			erase_below();
			break;
		}
		if (got != expected)
			mask = 0;
	}

	for (;;) {
		tx(confirm); // Mainloop ready
		g_addr = 0;
		switch (rx()) {
		case 'f': // read application flash, one page per host '!'
			for (;;) {
				if (rx() != confirm)
					break;
				g_cnt = page;
				sendf();
				if (g_addr >= app_end)
					break;
			}
			break;
		case 'F':                // erase the application, then take pages behind '?'
			erase_application(); // leaves g_addr = 0, the write start
			while (rcnf() == confirm)
				store_flash();
			break;
		case 'e': // read EEPROM, one page per host '!', until the host stops
			for (;;) {
				if (rx() != confirm)
					break;
				g_cnt = page;
				do {
					tx(eerd());
				} while (--g_cnt);
			}
			break;
		case 'E': // take EEPROM pages behind '?'
			while (rcnf() == confirm) {
				g_cnt = page;
				do {
					eewr(rx());
				} while (--g_cnt);
			}
			break;
		case 'c': // read the config page
		read_config:
			g_addr = app_end;
			g_cnt = page;
			sendf();
			break;
		case 'C': // replace the config page, then echo it back to verify
			if (rcnf() != confirm)
				break;
			g_addr = app_end + page;
			erase_below(); // leaves g_addr = app_end, the store target
			store_flash();
			goto read_config;
		default: // 'q' or any other byte runs the application
			appjump();
		}
	}
}

} // 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;
	// The one line of crt this loader needs: compiled code assumes
	// __zero_reg__ (r1) is 0, and power-on registers are undefined.
	asm volatile("clr __zero_reg__");
	tsb::run();
}