// pureboot — a serial bootloader on libavr: one C++ source, no inline
// assembly, no global register variables, 512 bytes 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 <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.
#if !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() (README.md)"
#endif

using dev = avr::device<{.clock = avr::hertz_t{PUREBOOT_CLOCK_HZ}}>;
constexpr avr::baud_t wire_baud{PUREBOOT_BAUD};

// The watchdog reset flag's home: MCUSR, or the classic megas' MCUCSR.
consteval std::int16_t wdrf_field()
{
	auto reg = std::string_view{avr::hw::db.regs[static_cast<std::size_t>(avr::power::detail::reset_reg())].name};
	return avr::hw::db.field_index(reg, "WDRF");
}

// 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::uint32_t base = spm::flash_bytes - slot_bytes;
constexpr std::uint16_t page = spm::page_bytes;
constexpr bool boot_section = avr::hw::curated::has_boot_section();

// Past 64 KiB a byte address no longer fits the wire's 16 bits, so flash
// addresses there are word addresses ('J' always was one). A slot is 256 of
// those — one value of a wire address's high byte, where 512 bytes span two.
constexpr bool word_flash = spm::flash_bytes > 65536;
constexpr std::uint16_t wire_base =
    word_flash ? static_cast<std::uint16_t>(base / 2) : static_cast<std::uint16_t>(base);

// 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.
#if !defined(PUREBOOT_TIMEOUT)
#define PUREBOOT_TIMEOUT 8
#endif
constexpr std::uint8_t timeout_seconds = PUREBOOT_TIMEOUT;

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

// The 'b' reply, byte for byte (layout: README.md). Flash-resident because
// no crt copies a .data image — and flash_table's storage carries the word
// alignment 'b' needs to halve the address on the large chips.
inline constexpr avr::flash_table<std::array<std::uint8_t, 12>{
    'P', 'B', version, avr::hw::db.signature[0], avr::hw::db.signature[1], avr::hw::db.signature[2],
    static_cast<std::uint8_t>(page), // 0 means 256
    wire_base & 0xff, wire_base >> 8, avr::hw::db.mem.eeprom_size & 0xff, avr::hw::db.mem.eeprom_size >> 8,
    static_cast<std::uint8_t>((boot_section ? 0 : 1) | (word_flash ? 2 : 0)), // patch-vector, word-addressed
}>
    info_data;

// The serial link, per the build's PUREBOOT_USART / PUREBOOT_SOFT_SERIAL,
// 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_RX)
#define PUREBOOT_RX pb0
#endif
#if !defined(PUREBOOT_TX)
#define PUREBOOT_TX pb1
#endif
#if defined(PUREBOOT_USART)
constexpr char usart_digit = '0' + PUREBOOT_USART;
#else
constexpr char usart_digit = '0';
#endif

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

	// The compiled idle poll: lds UCSR0A (2), sbrc skipping the exit (2),
	// sbiw + sbci + sbci + brne (6).
	static constexpr std::uint8_t poll_cycles = 10;

	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()
	{
		uart::drain();
	}
};

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

	// The compiled idle poll: sbis skipping the exit (2), sbiw + sbci +
	// sbci + brne (6).
	static constexpr std::uint8_t poll_cycles = 8;

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

#if defined(PUREBOOT_USART)
static_assert(avr::uart::has_usart<usart_digit>(), "PUREBOOT_USART selects a hardware USART this chip does not have");
using link = hardware_link<dev::clock>;
#elif defined(PUREBOOT_SOFT_SERIAL)
using link = software_link<dev::clock>;
#else
using link =
    std::conditional_t<avr::uart::has_usart<usart_digit>(), hardware_link<dev::clock>, software_link<dev::clock>>;
#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();
}

[[gnu::noinline, noreturn]] void run_app()
{
	jump(pureboot_app);
}

// The window as one 32-bit countdown, divided by the backend's counted
// poll-loop cycles. Whole seconds is all it promises.
consteval std::uint32_t window_polls()
{
	return timeout_seconds * static_cast<std::uint32_t>(dev::clock.hz / link::poll_cycles);
}

bool pending_before_deadline()
{
	std::uint32_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();
}

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

// Counts arrive in the wire's 8-bit form: 0 means 256. Both streamers fold
// into the one command that reads flash, which is what lets the far one's
// 24-bit cursor sit in the command loop's own call-saved registers.
[[maybe_unused, gnu::always_inline]] inline void send_flash_near(std::uint16_t address, std::uint8_t count)
{
	do
		link::tx(avr::flash_load(reinterpret_cast<const std::uint8_t *>(address++)));
	while (--count);
}

// The 24-bit cursor as the machine holds it — the RAMPZ byte and a 16-bit Z,
// carried apart; the reassembled address folds away inside the far load.
[[maybe_unused, gnu::always_inline]] inline void send_flash_far(std::uint16_t address, std::uint8_t count)
{
	std::uint8_t rampz = static_cast<std::uint8_t>(address >> 15);
	std::uint16_t z = static_cast<std::uint16_t>(address << 1);
	do {
		link::tx(avr::flash_load_far<std::uint8_t>((static_cast<std::uint32_t>(rampz) << 16) | z));
		// Carrying the wrap is smaller than the flat 32-bit cursor GCC
		// builds without it.
		if (++z == 0)
			++rampz;
	} while (--count);
}

[[gnu::always_inline]] inline void send_flash(std::uint16_t address, std::uint8_t count)
{
	if constexpr (word_flash)
		send_flash_far(address, count);
	else
		send_flash_near(address, count);
}

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

void send_eeprom(std::uint16_t address, std::uint8_t count)
{
	do
		link::tx(ee::read(address++));
	while (--count);
}

// 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.
void store_eeprom(std::uint16_t address, std::uint8_t count)
{
	do {
		ee::write<off>(address++, link::rx());
		tx_ack();
	} while (--count);
}

// One page into the SPM buffer, then erase and program — except the slot
// this code is running in (`slot_high`, from run()), which is drained and
// left alone. A broken host therefore cannot brick the running loader, and a
// copy one slot lower may rewrite the resident one.
//
// 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 program_flash(std::uint16_t wire_address, std::uint8_t slot_high)
{
	// One induction either way: a byte-addressed wire address walks the page
	// itself (aligned, so the offset bits wrap to zero), while a word one
	// becomes a byte cursor once. The slot index is the wire address's high
	// byte — on byte-addressed chips the byte address's, with the low bit
	// dropped, since a slot is two of those.
	spm::flash_address_t address;
	std::uint8_t page_high;
	if constexpr (word_flash) {
		// A page is aligned, so it never crosses 64 KiB: RAMPZ is a per-page
		// constant and the 16-bit Z's low byte is the whole in-page offset.
		const std::uint8_t rampz = static_cast<std::uint8_t>(wire_address >> 15);
		const std::uint16_t z0 = static_cast<std::uint16_t>(wire_address << 1);
		std::uint16_t z = z0;
		do {
			std::uint8_t low = link::rx();
			std::uint8_t high = link::rx();
			spm::fill<off>((static_cast<spm::flash_address_t>(rampz) << 16) | z, word_of({low, high}));
			z += 2;
		} while (static_cast<std::uint8_t>(z));
		address = (static_cast<spm::flash_address_t>(rampz) << 16) | z0;
		page_high = static_cast<std::uint8_t>(wire_address >> 8);
	} else {
		address = static_cast<spm::flash_address_t>(wire_address);
		do {
			std::uint8_t low = link::rx();
			std::uint8_t high = link::rx();
			spm::fill<off>(address, word_of({low, high}));
			address += 2;
		} while (static_cast<std::uint8_t>(address) & (page - 1));
		address -= 2; // back inside the page — erase and write ignore the word bits
		page_high = static_cast<std::uint8_t>(address >> 8) & 0xfe;
	}
	if (page_high != slot_high) {
		// Only a boot-sectioned mega runs on while its RWW section programs;
		// everywhere else the CPU halts through erase and write.
		spm::erase_page<off>(address);
		if constexpr (boot_section)
			spm::wait();
		spm::write_page<off>(address);
		if constexpr (boot_section)
			spm::wait();
	}
	// Programming leaves the RWW section disabled; reads need it back on. The
	// same store discards the buffer (§26.2.2), so a boot-sectioned mega never
	// meets the stale-word case above.
	if constexpr (boot_section)
		spm::rww_enable<off>();
}

// The four fuse and lock bytes in the hardware's own Z order: low, lock,
// extended, high.
void send_fuses()
{
	std::uint8_t which = 0;
	do
		link::tx(spm::read_fuse<off>(static_cast<spm::fuse>(which)));
	while (++which != 4);
}

[[noreturn]] void run()
{
	// A watchdog reset belongs to the application, whose watchdog stays forced
	// on until it clears WDRF — no activation window in its way.
	if (avr::hw::field_impl<wdrf_field()>::test())
		run_app();

	link::init();

	// The high byte of the slot this copy runs at, which the write guard and
	// the info block both follow: the return address is a word address, so its
	// high byte is the 256-word slot index, doubled back into byte terms where
	// the wire counts bytes. Taken as byteswap's low byte — the builtin already
	// swaps the two stacked bytes, and the double swap folds away, where `>> 8`
	// would leave the swap materialized.
	const std::uint16_t ra_words = reinterpret_cast<std::uint16_t>(__builtin_return_address(0));
	const std::uint8_t ra_high = static_cast<std::uint8_t>(std::byteswap(ra_words));
	const std::uint8_t slot_high = word_flash ? ra_high : static_cast<std::uint8_t>(ra_high << 1);

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

	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 'J': { // jump to a wire word address: hand-over and staging transfer
			auto target = reinterpret_cast<void (*)()>(rx16());
			tx_ack();
			link::drain();
			jump(target);
		}
		case 'b':   // info block, read relative to the running slot
		case 'R':   // read flash: addr16, n8 (0 = 256)
		case 'r':   // read EEPROM: addr16, n8
		case 'w': { // write EEPROM: addr16, n8, then n bytes each acked
			// One address-and-count path for all four: 'b' is a flash read
			// whose arguments the loader already knows, so it joins the
			// wire-argument three rather than streaming from a call site of its
			// own. That leaves one flash streamer in the image, and lets its
			// cursor live in this never-returning loop's own call-saved
			// registers instead of being saved and restored around a call.
			std::uint16_t address;
			std::uint8_t count;
			if (command == 'b') {
				// The block sits in the image's first 256 bytes (check_pi.py
				// asserts it) and slots are 512-aligned, so the low byte of its
				// link address is its offset in any slot — halved where wire
				// units are words. The high byte is runtime data, so no
				// absolute address is ever materialized.
				const auto link_byte =
				    static_cast<std::uint8_t>(reinterpret_cast<std::uint16_t>(info_data.storage.data()));
				const std::uint8_t low = word_flash ? static_cast<std::uint8_t>(link_byte >> 1) : link_byte;
				address = static_cast<std::uint16_t>(low | (slot_high << 8));
				count = static_cast<std::uint8_t>(info_data.size());
			} else {
				address = rx16();
				count = link::rx();
			}
			if (command == 'r')
				send_eeprom(address, count);
			else if (command == 'w')
				store_eeprom(address, count);
			else
				send_flash(address, count);
			break;
		}
		case 'W': // program one flash page: addr16, page bytes
			program_flash(rx16(), slot_high);
			break;
		case 'F': // fuse and lock bytes
			send_fuses();
			break;
		default: // unknown bytes are ignored; the loop re-acks
			break;
		}
	}
}

} // namespace
} // namespace pureboot

template struct avr::startup::entry<pureboot::run>;