// pureboot — a serial bootloader on libavr: one C++ source, no inline
// assembly, no global register variables, a 512-byte slot on every chip
// libavr targets. The device speaks primitives; every composite (verify, erase,
// reset-vector surgery, self-update) lives in the host tool. Protocol,
// deployment and configuration: README.md next to this file.
//
// The image is position-independent — PC-relative control flow, wire
// addresses in, the write guard and the info block both anchored on the
// runtime return address — so the identical binary runs from any slot. That
// is what makes a copy one slot below able to rewrite the resident one, and
// every change here has to keep it (test/check_pi.py).

#include <chrono>

#include <libavr/libavr.hpp>

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

namespace pureboot {
namespace {

// Purely polled: every interrupt guard folds to nothing.
constexpr auto off = avr::irq::guard_policy::unused;

constexpr std::uint8_t ack = '+';

// Deployment parameters come from the build (pureboot_add_loader()). The
// signature is not one of them: the chip database is the only universal
// source — a tiny13A cannot read its own signature row from code. An autobaud
// build carries no clock and no baud at all; it measures both.
#if !defined(PUREBOOT_AUTOBAUD) && (!defined(PUREBOOT_CLOCK_HZ) || !defined(PUREBOOT_BAUD))
#error \
    "PUREBOOT_CLOCK_HZ and PUREBOOT_BAUD select this build's clock and baud — create loader targets with pureboot_add_loader(), or PUREBOOT_AUTOBAUD for a clock-free one (README.md)"
#endif

#if !defined(PUREBOOT_AUTOBAUD)
using dev = avr::device<{.clock = avr::hertz_t{PUREBOOT_CLOCK_HZ}}>;
constexpr avr::baud_t wire_baud{PUREBOOT_BAUD};
#endif

// The loader owns the top 512 bytes; a staging copy goes in the slot below.
// Chips without a hardware boot section — the tinies and the m48s, whose SPM
// runs from anywhere (Atmel-8271 §26) — keep the application's relocated
// reset vector in the word under the slot.
constexpr std::uint16_t slot_bytes = 512;
constexpr std::uint16_t page = spm::page_bytes;
constexpr bool boot_section = avr::hw::curated::has_boot_section();

// Past 64 KiB one bank of flash does not cover the chip, so a transfer's
// selector byte carries the bank and the wire address stays a byte address
// within it. 'J' is the exception: it is a word address everywhere, because
// that is what the hardware's own jump takes.
constexpr bool banked_flash = spm::flash_bytes > 65536;

// A compile-time window, so the whole EEPROM belongs to the application;
// re-timing a deployed loader is a self-update with a re-timed build. An
// autobaud build has no clock to convert seconds against and counts polls.
#if !defined(PUREBOOT_TIMEOUT)
#define PUREBOOT_TIMEOUT 8
#endif
constexpr std::uint8_t timeout_seconds = PUREBOOT_TIMEOUT;

#if !defined(PUREBOOT_AUTOBAUD_POLLS)
#define PUREBOOT_AUTOBAUD_POLLS 4000000
#endif
constexpr avr::uint24_t autobaud_budget = PUREBOOT_AUTOBAUD_POLLS;

// A build may bake a measured oscillator trim (README.md: the RC-oscillator
// deployment answer); the byte is applied at the top of run(). Orthogonal to
// the serial backend — an autobaud build may carry it for the application's
// benefit alone.
#if defined(PUREBOOT_OSCCAL)
static_assert(PUREBOOT_OSCCAL >= 0 && PUREBOOT_OSCCAL <= 0xff, "PUREBOOT_OSCCAL is one OSCCAL byte");
#endif

// The loader's one identity number. The protocol carries none of its own —
// a version implies it, and the host tool holds that map (README.md).
constexpr std::uint8_t version = 7;

// The image's identity stamp, for the host tool rather than for the wire: an
// update image is a bare 512-byte slot, and without this nothing in it says
// which chip it was built for. The tool refuses to install an image whose
// stamp does not match the device — flashing a foreign loader bricks the
// target, and the loader itself cannot check what has already replaced it.
//
// Never read from flash by the loader — 'b' answers out of this array, but at
// constant indices, so those fold to immediates and no runtime address of it
// is ever formed. `used` keeps the compiler from dropping the copy the host
// needs and `retain` keeps --gc-sections from collecting it.
// clang-format off
[[gnu::used, gnu::retain, gnu::section(".text.stamp")]]
inline constexpr std::uint8_t identity_stamp[]{
    'P', 'B',                      // the magic the host scans an image for
    version,                       // and from here on, exactly what 'b' answers
    avr::hw::db.signature[0],
    avr::hw::db.signature[1],
    avr::hw::db.signature[2],
};
// clang-format on
// Where the identity proper starts: past the magic the host scans for.
constexpr std::uint8_t stamp_identity = 2;

// The address spaces a transfer can name, in a selector byte's low nibble.
// Flash is 0 so it is the cheapest to select.
//
// spm_ops is the one that is not memory: a write there hands its byte to
// SPMCSR and fires the instruction at the transfer's address, which is how
// page erase, page write and RWW re-enable reach the wire without the loader
// carrying a command for each. The hardware's four-cycle store-to-SPM window
// is why this is one fused primitive and not a poke of SPMCSR — no host can
// hit that window across a serial link.
enum : std::uint8_t { sp_flash = 0, sp_eeprom = 1, sp_data = 2, sp_fuse = 3, sp_spm = 4 };

// A selector's high nibble is the flash bank — the address bits above the
// 16-bit wire address, RAMPZ on the chips that have one. Keeping it here
// rather than widening the wire address is what lets one 16-bit cursor serve
// every space: a 24-bit cursor would pay its extra byte on EEPROM and data
// reads that can never need it.
[[gnu::always_inline]] inline std::uint8_t space_of(std::uint8_t selector)
{
	return selector & 0x0f;
}

[[gnu::always_inline]] inline std::uint8_t bank_of(std::uint8_t selector)
{
	return static_cast<std::uint8_t>(selector >> 4);
}

// The slot a flash address falls in, as one byte. A slot is half as many words
// as bytes, so the word address's high byte is exactly this index — which is
// what lets the write guard compare a single byte, and what the running copy's
// own return address yields for free.
constexpr std::uint8_t slot_shift = std::countr_zero(slot_bytes);
constexpr std::uint8_t bank_shift = 16 - slot_shift;

[[gnu::always_inline]] inline std::uint8_t slot_of([[maybe_unused]] std::uint8_t bank, std::uint16_t at)
{
	const auto within = static_cast<std::uint8_t>(at >> slot_shift);
	if constexpr (banked_flash)
		return static_cast<std::uint8_t>((bank << bank_shift) | within);
	else
		return within;
}

// The serial link, per the build's PUREBOOT_USART / PUREBOOT_SOFT_SERIAL /
// PUREBOOT_AUTOBAUD, defaulting to the chip's USART0 where it has one. The
// software receiver is the polled one: the vector table belongs to the
// application. Templates on the clock, so only the selected backend
// instantiates. pending() is the cheap line test the activation window polls;
// drain() holds until the last frame is off the wire, so a hand-over cannot
// let the target's re-init clip the ack.
#if defined(PUREBOOT_SOFT_SERIAL) && defined(PUREBOOT_USART)
#error "PUREBOOT_SOFT_SERIAL and PUREBOOT_USART select opposing serial backends"
#endif
#if defined(PUREBOOT_AUTOBAUD) && defined(PUREBOOT_USART)
#error "PUREBOOT_AUTOBAUD measures a software link; it cannot drive a hardware USART"
#endif
#if !defined(PUREBOOT_RX)
#define PUREBOOT_RX pb0
#endif
#if !defined(PUREBOOT_TX)
#define PUREBOOT_TX pb1
#endif
#if defined(PUREBOOT_USART)
constexpr int usart_unit = PUREBOOT_USART;
#else
constexpr int usart_unit = 0;
#endif

template <avr::hertz_t C, avr::baud_t B>
struct hardware_link {
	using uart = avr::uart::usart<usart_unit, C, {.baud = B, .max_baud_error = 2.5_pct}>;

	// The compiled idle poll around the window's narrow (uint24_t) countdown:
	// the RXC test, then sbiw + sbci + brne (5). The test's cost follows the
	// status register's home — a 2-cycle bit-skip where UCSRnA sits in
	// bit-addressable I/O (the classic megas), lds + skip (4) in extended
	// I/O. A uint32_t countdown pays one more sbci — window_polls() adds it
	// where the count forces the wide type. Held by the pureboot.window gate.
	// The lookup rides the baud parameter so it stays dependent: the trait is
	// an incomplete type on the USART-less chips, which parse this template
	// without ever instantiating it.
	template <avr::baud_t Baud, typename U = avr::hw::usart_of<usart_unit>>
	static consteval std::uint8_t poll_cost()
	{
		return U::ucsra::addr < 0x40 ? 7 : 9;
	}
	static constexpr std::uint8_t poll_cycles = poll_cost<B>();

	static void init()
	{
		avr::init<uart>();
	}

	static bool pending()
	{
		return uart::rx_ready();
	}

	static std::uint8_t rx()
	{
		return uart::read_blocking();
	}

	static void tx(std::uint8_t byte)
	{
		uart::write(byte);
	}

	static void drain()
	{
		// A drain here always follows this link's own write — the frame is
		// in flight by construction, so the completion the wait needs is
		// guaranteed and the bounded default's countdown would be dead bytes.
		uart::drain_unbounded();
	}
};

template <avr::hertz_t C, avr::baud_t B>
struct software_link {
	using rx_t = avr::uart::software_rx_polled<C, avr::PUREBOOT_RX, B>;
	using tx_t = avr::uart::software_tx<C, avr::PUREBOOT_TX, B>;

	// The compiled idle poll around the window's narrow (uint24_t) countdown:
	// sbis skipping the exit (2), sbiw + sbci + brne (5). A uint32_t
	// countdown pays one more sbci — window_polls() adds it where the count
	// forces the wide type. Held by the pureboot.window gate.
	static constexpr std::uint8_t poll_cycles = 7;

	static void init()
	{
		avr::init<rx_t, tx_t>();
	}

	static bool pending()
	{
		return rx_t::start_pending();
	}

	static std::uint8_t rx()
	{
		return rx_t::template read_blocking<off>();
	}

	static void tx(std::uint8_t byte)
	{
		tx_t::template write<off>(byte);
	}

	static void drain()
	{
		// The software transmitter returns only after the stop bit.
	}
};

// The clock-free link: the bit period is measured from the host's calibration
// pulse instead of derived from a clock, so one image serves every F_CPU and
// every rate. Activation differs in kind from the other two — there is no
// clock to time a window against — so this backend brings its own, below.
struct autobaud_link {
	// The unit in GPIOR2:GPIOR1 where the chip has them: the loader owns the
	// whole chip while it runs, and the pair costs one word per access where
	// the RAM word costs two — six words across the image.
	using uart = avr::uart::software_autobaud<avr::PUREBOOT_RX, avr::PUREBOOT_TX, avr::uart::unit_home::gpior>;

	static void init()
	{
		avr::init<uart>();
	}

	static std::uint8_t rx()
	{
		return uart::template read<off>();
	}

	static void tx(std::uint8_t byte)
	{
		uart::template write<off>(byte);
	}

	static void drain()
	{
		// A drain here always follows this link's own write — the frame is
		// in flight by construction, so the completion the wait needs is
		// guaranteed and the bounded default's countdown would be dead bytes.
		uart::drain_unbounded();
	}
};

#if defined(PUREBOOT_AUTOBAUD)
using link = autobaud_link;
#elif defined(PUREBOOT_USART)
static_assert(avr::uart::has_usart<usart_unit>(), "PUREBOOT_USART selects a hardware USART this chip does not have");
using link = hardware_link<dev::clock, wire_baud>;
#elif defined(PUREBOOT_SOFT_SERIAL)
using link = software_link<dev::clock, wire_baud>;
#else
using link = std::conditional_t<avr::uart::has_usart<usart_unit>(), hardware_link<dev::clock, wire_baud>,
                                software_link<dev::clock, wire_baud>>;
#endif

// The application's entry, pinned by the linker (--defsym): word 0 on a
// boot-sectioned mega, the trampoline at base − 2 elsewhere. Reaching it must
// not depend on where this copy runs, so the jump goes through a pointer, and
// [[gnu::noipa]] keeps the constant from folding back into a relative call.
extern "C" [[noreturn]] void pureboot_app();

[[gnu::noipa, noreturn]] void jump(void (*target)())
{
	target();
	__builtin_unreachable();
}

[[noreturn]] void run_app()
{
	jump(pureboot_app);
}

// Activation: a bounded wait for the host, then the knock. Both forms boot the
// application when the window closes on an idle line, and both bound *every*
// wait — a knock awaited without a deadline would let one stray edge hold an
// unattended device in the loader forever.
#if defined(PUREBOOT_AUTOBAUD)
// The window is a fixed poll budget: with no clock, whole seconds cannot be
// timed. A uint24_t holds it — a fourth byte would cost two words at every
// countdown step for range never used.
void await_host()
{
	for (;;) {
		if (!link::uart::calibrate(autobaud_budget))
			run_app();
		// The calibration pulse has already proven a host is there, so one
		// byte activates. A knock that never arrives falls back to calibrate(),
		// whose own budget then boots the application.
		if (link::uart::template read<off>(autobaud_budget) == 'p')
			return;
	}
}
#else
// The window as one countdown, divided by the backend's counted poll-loop
// cycles. Whole seconds is all it promises. The per-poll cost depends on the
// countdown's own width (a uint32_t decrement chain is one sbci longer), and
// the width depends on the poll count — solved narrow-first: a count that
// fits 24 bits at the narrow cost keeps the narrow loop, anything else takes
// the wide loop at its own cost. A count fitting 24 bits only at the wide
// cost stays wide, so the choice cannot oscillate on the boundary.
consteval std::uint32_t polls_at(std::uint32_t per_poll)
{
	// Whole-window cycles first, then the per-poll division: one truncation
	// instead of one per second. Same instructions either way — only the
	// countdown's immediate moves.
	return static_cast<std::uint32_t>(dev::cycles_for<std::chrono::seconds{timeout_seconds}>() / per_poll);
}

consteval bool narrow_window()
{
	return polls_at(link::poll_cycles) <= 0xffffff;
}

consteval std::uint32_t window_polls()
{
	return polls_at(narrow_window() ? link::poll_cycles : link::poll_cycles + 1u);
}

// The countdown in the narrowest type that holds it: a fourth byte would
// cost a wider decrement chain at every poll for range most windows never
// use (the autobaud budget makes the same choice).
using window_t = std::conditional_t<narrow_window(), avr::uint24_t, std::uint32_t>;

bool pending_before_deadline()
{
	window_t polls = window_polls();
	do {
		if (link::pending())
			return true;
	} while (--polls);
	return false;
}

// A knock byte under the deadline: an idle window means no host, so the
// application runs.
std::uint8_t rx_deadline()
{
	if (!pending_before_deadline())
		run_app();
	return link::rx();
}

void await_host()
{
	// 'p' then 'b', each under a fresh window; anything else is line noise.
	while (rx_deadline() != 'p' || rx_deadline() != 'b') {
	}
}
#endif

// Inlined: read across a call, the first byte strands in a call-saved
// register the caller has to push and pop.
[[gnu::always_inline]] inline std::uint16_t rx16()
{
	std::uint16_t low = link::rx();
	return static_cast<std::uint16_t>(low | (link::rx() << 8));
}

// The wire's byte pair as the word it is — AVR is little-endian too, so the
// cast is the identity a shift-and-or spelling makes the compiler rediscover.
// Callers read into named variables first: the wire order is a sequence of
// reads, not an argument order.
[[gnu::always_inline]] inline std::uint16_t word_of(std::array<std::uint8_t, 2> pair)
{
	return std::bit_cast<std::uint16_t>(pair);
}

// Out of line: several sites send it, and a call is shorter than a
// load-immediate at each.
[[gnu::noinline]] void tx_ack()
{
	link::tx(ack);
}

// A wire address and its selector's bank as the flash address they name.
[[gnu::always_inline]] inline spm::flash_address_t flash_address([[maybe_unused]] std::uint8_t bank, std::uint16_t at)
{
	if constexpr (banked_flash)
		return (static_cast<spm::flash_address_t>(bank) << 16) | at;
	else
		return at;
}

// One byte out of any space. Every accessor shares the transfer's cursor, its
// loop and its call site, so a space costs only its own instruction rather
// than a body, a loop and a dispatch arm of its own.
[[gnu::always_inline]] inline std::uint8_t load(std::uint8_t space, [[maybe_unused]] std::uint8_t bank,
                                                std::uint16_t at)
{
	if (space == sp_eeprom)
		return ee::read(at);
	if (space == sp_data)
		return *reinterpret_cast<volatile std::uint8_t *>(at);
	if (space == sp_fuse)
		return spm::read_fuse<off>(static_cast<spm::fuse>(at));
	if constexpr (banked_flash)
		return avr::flash_load_far<std::uint8_t>(flash_address(bank, at));
	else
		return avr::flash_load(reinterpret_cast<const std::uint8_t *>(at));
}

// One byte into a writable space. Flash is not one of them — it arrives a
// page at a time through 'W' and is committed through sp_spm — and the fuses
// are not writable at all: SPM reaches flash and boot lock bits only.
[[gnu::always_inline]] inline void store(std::uint8_t space, std::uint8_t bank, std::uint16_t at, std::uint8_t value,
                                         std::uint8_t slot_high)
{
	if (space == sp_data) {
		*reinterpret_cast<volatile std::uint8_t *>(at) = value;
		return;
	}
	if (space == sp_spm) {
		// The running-slot write guard. An SPM command aimed at the slot this
		// code executes from is dropped, so a broken host cannot brick the
		// running loader — while a copy one slot lower may still rewrite the
		// resident one, which is what a self-update is. Guarding the commit
		// rather than the page fill covers erase and write both, and leaves a
		// refused page's words in the buffer: harmless, since the next page
		// write auto-erases it (§26.2.1).
		if (slot_of(bank, at) != slot_high)
			spm::command<off>(value, flash_address(bank, at));
		// Only a boot-sectioned mega runs on while its RWW section programs;
		// everywhere else the CPU halts through erase and write, so the wait
		// is already over by the time it returns.
		if constexpr (boot_section)
			spm::wait();
		return;
	}
	// Host-paced: the ack goes out once the write has begun, so the next byte
	// arrives while it completes and nothing is missed without a buffer.
	ee::write<off>(at, value);
}

// One page into the SPM buffer, and only that: the erase and the write that
// commit it are host-issued sp_spm stores, which reach the same fused
// store-and-SPM pair through the transfer path's own address and data.
//
// Nothing discards the buffer first: it is write-once per word (§26.2.1), so
// filling over a refused page or an application's leavings programs stale
// words — but a page write auto-erases it (§26.2.1; §19.2 on the tinies), so
// that write clears the condition and the host's read-back rewrites the page.
void fill_page(std::uint8_t bank, std::uint16_t at)
{
	// The address names a page, so its in-page bits are dropped and the walk
	// starts at the page base; the low byte of the cursor is the whole in-page
	// offset, since a page is aligned and never crosses a bank.
	std::uint16_t z = at & ~static_cast<std::uint16_t>(page - 1);
	do {
		std::uint8_t low = link::rx();
		std::uint8_t high = link::rx();
		spm::fill<off>(flash_address(bank, z), word_of({low, high}));
		z += 2;
	} while (static_cast<std::uint8_t>(z) & (page - 1));
}

[[noreturn]] void run()
{
#if defined(PUREBOOT_OSCCAL)
	// The build's oscillator trim, ahead of everything — the WDRF bail
	// included — so every path out of reset, the watchdog hand-over to the
	// application first among them, runs on the corrected clock.
	avr::clock::calibrate(PUREBOOT_OSCCAL);
#endif
	// A watchdog reset belongs to the application, whose watchdog stays forced
	// on until it clears WDRF — no activation window in its way.
	if (avr::power::peek_reset_cause().watchdog)
		run_app();

	link::init();

	// The slot this copy runs in, which the write guard follows: the return
	// address is a word address and a slot is half as many words as bytes, so
	// its high byte is the slot index outright. No absolute address is ever
	// formed, so the image stays position-independent.
	const auto slot_high = avr::startup::caller_page();

	await_host();

	for (;;) {
		// No prompt while an EEPROM write runs: it blocks SPM and fuse reads
		// (§26.2.1), and the prompt is the previous command's completion ack.
		ee::wait();
		tx_ack();
		const std::uint8_t command = link::rx();
		switch (command) {
		case 'b': // identity: the version, then the three signature bytes
			// Straight out of the stamp, so the wire and the image can never
			// disagree about what this loader is. The indices are constant and
			// the array is constexpr, so these are immediates, not flash reads:
			// nothing here needs the stamp's runtime address.
			for (std::uint8_t at = stamp_identity; at != sizeof identity_stamp; ++at)
				link::tx(identity_stamp[at]);
			break;
		case 'J':   // jump:  sel8 (reserved), addr16 as a wire word address
		case 'W':   // fill one flash page buffer: sel8, addr16, then page bytes
		case 'G':   // read:  sel8, addr16, n8 (0 = 256)
		case 'g': { // write: sel8, addr16, n8, then n bytes, each acked
			// One decode, one cursor and one loop for every space, both
			// directions and the jump: a command per memory would carry a copy
			// of all three each. 'J' — the hand-over and staging transfer —
			// carries a selector it ignores so its address rides the same two
			// reads as everything else; 'W' joins the same decode rather than
			// keeping an address form of its own, so flash addressing is
			// uniform across every command that names it.
			const std::uint8_t selector = link::rx();
			const std::uint8_t space = space_of(selector);
			const std::uint8_t bank = bank_of(selector);
			std::uint16_t at = rx16();
			if (command == 'J') {
				tx_ack();
				link::drain();
				jump(reinterpret_cast<void (*)()>(at));
			}
			if (command == 'W') {
				fill_page(bank, at);
				break;
			}
			std::uint8_t count = link::rx();
			do {
				// Read and write are one letter apart in case, so the direction
				// is a single bit and the loop picks it with a one-word skip.
				if (command & 0x20) {
					store(space, bank, at, link::rx(), slot_high);
					tx_ack();
				} else
					link::tx(load(space, bank, at));
				++at;
			} while (--count);
			break;
		}
		default: // unknown bytes are ignored; the loop re-acks
			break;
		}
	}
}

} // namespace
} // namespace pureboot

// stack::hardware: activation is reset-only, so the reset logic's own
// SP = RAMEND stands wherever the datasheet guarantees it (the classic
// megas still get the write); a 'J' entry runs on the caller's live stack.
template struct avr::startup::entry<pureboot::run, avr::startup::stack::hardware>;