// pureboot — a serial bootloader on libavr, pure by constraint: one C++
// source with no inline assembly and no global register variables, built for
// every chip libavr targets, 512 bytes on each. The device speaks primitives
// — read/program flash, read/write EEPROM, fuse bytes, an info block, a jump
// — and everything composite (verify, erase, reset-vector surgery, updating
// the loader itself) lives in the host tool. Protocol reference: README.md
// next to this file.
//
// The image is position-independent: control flow is PC-relative, the write
// and read paths take wire addresses, the write guard refuses the 512-byte
// slot the code is *running* in (taken from the runtime return address), the
// info block is read relative to that same anchor, and the application jump
// is an indirect call to an absolute entry. The identical binary therefore
// runs from any 512-byte slot with every command intact: flashed one slot
// below the resident loader it becomes the staging loader that rewrites the
// resident — how pureboot updates itself, host-driven, with no other
// firmware involved.
//
// Entry: reset lands in avr::startup::entry below (BOOTRST on the
// boot-sectioned megas; the patched reset vector — or erased flash walking
// up into the loader — on the tinies and the boot-section-less m48s). A
// watchdog reset hands straight to the application. Otherwise the
// host has one activation window per awaited knock byte ("pb"); an idle line
// boots the application. A session then stays in the command loop until 'J'
// jumps away or the chip resets.

#include <libavr/libavr.hpp>

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

namespace pureboot {
namespace {

// Purely polled — interrupts stay off, every guard folds to nothing.
constexpr auto off = avr::irq::guard_policy::unused;

constexpr std::uint8_t ack = '+';

// Per-deployment personality, passed in by the build — pureboot_add_loader()
// (the CMake function next to this file) resolves the defaults: the clock the
// board actually runs, the wire baud, the serial backend and its pins. The
// device signature needs no configuring — it comes from the chip database
// (avr::hw::db.signature), the only universal source, since the tiny13A
// cannot even read its 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");
}

// Geometry: the resident loader owns the top slot of flash — 512 bytes,
// except on the >64 KiB chips whose own smallest boot sector is 1 KiB (the
// 1284s): there the slot is 1 KiB, matching the hardware boundary the
// 512-byte figure comes from everywhere else. The word below the slot is
// the trampoline (the application's relocated reset vector) on chips
// without a hardware boot section — the tinies and the m48s, whose SPM
// runs from anywhere (Atmel-8271 §26). A boot section also means the CPU
// runs on while the RWW section programs; everywhere else it halts through
// the operation. The m48s still carry RWWSRE as their temporary-buffer
// discard (§26.2), so the discard picks by that bit, not by the section.
constexpr std::uint16_t slot_bytes = spm::flash_bytes > 65536 ? 1024 : 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();
constexpr bool rww_discard = spm::detail::has_rww();

// Past 64 KiB a byte address no longer fits the wire's 16 bits, so on the
// large chips every flash address on the wire — and all slot arithmetic —
// is a word address instead ('J' always was one). A slot spans the same
// wire-high-byte pair in either unit (512 B = 2 x 256 bytes, 1 KiB =
// 2 x 256 words), so the slot index is the high byte with its low bit
// dropped everywhere.
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);
constexpr std::uint16_t wire_page_mask = word_flash ? (page / 2 - 1) : (page - 1);

// The activation window, in seconds, is a compile-time constant (the build
// may override it): the whole EEPROM belongs to the application, and
// re-timing the loader is a bootloader self-update with a re-timed binary.
#if !defined(PUREBOOT_TIMEOUT)
#define PUREBOOT_TIMEOUT 8
#endif
constexpr std::uint8_t timeout_seconds = PUREBOOT_TIMEOUT;

// The 12-byte info block the host reads with the 'b' command; flash-resident
// (there is no crt to copy a .data image), word-aligned so its wire (word)
// address is exact on the large chips. The page byte is the wire count
// convention: 0 means 256.
[[gnu::progmem]] alignas(2) inline constexpr std::array<std::uint8_t, 12> info_data = {
    'P',
    'B',
    1, // magic, protocol version
    avr::hw::db.signature[0],
    avr::hw::db.signature[1],
    avr::hw::db.signature[2],
    static_cast<std::uint8_t>(page),
    wire_base & 0xff,
    wire_base >> 8, // app flash ends here; resident loader base (a word address on large chips)
    avr::hw::db.mem.eeprom_size & 0xff,
    avr::hw::db.mem.eeprom_size >> 8,
    // bit 0: host must patch the reset vector (no hardware boot section);
    // bit 1: flash wire addresses are word addresses
    static_cast<std::uint8_t>((boot_section ? 0 : 1) | (word_flash ? 2 : 0)),
};

// The serial link. PUREBOOT_USART forces a hardware USART instance,
// PUREBOOT_SOFT_SERIAL the polled software UART (no vector — the table
// belongs to the application) on PUREBOOT_RX/PUREBOOT_TX; with neither, the
// chip's first USART where it has one and the software UART elsewhere. Both
// are class templates on the clock so only the selected backend is ever
// instantiated. pending() is the cheap line test the activation window
// polls; rx() then picks the byte up; drain() holds until the last
// transmitted frame is fully on the wire (the jump hand-over must not 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

// Whether the chip carries the selected USART: the suffixed instance name,
// or — for instance 0 — the classic megas' un-numbered block.
consteval bool usart_exists()
{
	const char name[] = {'U', 'S', 'A', 'R', 'T', usart_digit};
	if (avr::hw::db.has_instance(std::string_view{name, sizeof(name)}))
		return true;
	return usart_digit == '0' && avr::hw::db.has_instance("USART");
}

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(usart_exists(), "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<usart_exists(), hardware_link<dev::clock>, software_link<dev::clock>>;
#endif

// The application's entry, an absolute address the linker pins (--defsym in
// CMakeLists.txt): 0x0000 on the mega (word 0 stays the application's own
// vector — BOOTRST re-vectors a reset into the loader in hardware) and the
// trampoline word at base - 2 on the tinies. Reaching it must not depend on
// where this copy runs, so the jump goes through a pointer: [[gnu::noipa]]
// keeps the constant from folding back into a PC-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);
}

// One activation window is a single 32-bit poll countdown. The divisor is
// the backend's counted poll-loop cycles (its own comment reads them off the
// compiled loop); whole-second precision is all the window promises, so the
// nearest cycle count is plenty.
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 activation deadline: an idle line means no host is
// there, and the application runs.
std::uint8_t rx_deadline()
{
	if (!pending_before_deadline())
		run_app();
	return link::rx();
}

std::uint16_t rx16()
{
	std::uint16_t low = link::rx();
	return static_cast<std::uint16_t>(low | (link::rx() << 8));
}

// The streamers take the count in the wire's 8-bit form: 0 means 256.
// send_flash stays out of line: its two callers ('b' and 'R') otherwise each
// inline a private copy of the loop. On the large chips the address is a
// word address and the read goes through ELPM (flash_load_far).
[[gnu::noinline]] void send_flash(std::uint16_t address, std::uint8_t count)
{
	if constexpr (word_flash) {
		// The 24-bit cursor as the machine holds it: the RAMPZ byte and a
		// 16-bit Z, carried explicitly (the reassembled 32-bit address
		// folds away inside the inlined far load). A single read never
		// crosses a 64 KiB boundary — the protocol forbids it and the host
		// splits its chunks there — so RAMPZ holds for the whole run.
		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));
			if (++z == 0)
				++rampz; // robustness for a host that reads across 64 KiB
		} while (--count);
	} else {
		do
			link::tx(avr::flash_load(reinterpret_cast<const std::uint8_t *>(address++)));
		while (--count);
	}
}

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

// EEPROM write, host-paced: each ack goes out once the byte's write has
// begun, so the next byte arrives while it completes and the following
// write's own ready-wait sees an idle line. Nothing is ever missed, on
// either serial backend, without a buffer.
void store_eeprom(std::uint16_t address, std::uint8_t count)
{
	do {
		ee::write<off>(address++, link::rx());
		link::tx(ack);
	} while (--count);
}

// One flash page: stream the bytes into the SPM buffer as little-endian
// words, then erase and program — except the 512-byte slot this code runs
// in, which is drained but never programmed, so a copy can never erase
// itself. `slot_high` is the high byte of that running slot's base (run()
// derives it); a broken host thus cannot brick the running loader, and a
// copy flashed one slot lower may rewrite the slot above it — how pureboot
// updates itself. On the mega the RWW section is re-enabled so reads work
// immediately.
void program_flash(std::uint16_t wire_address, std::uint8_t slot_high)
{
	// A buffer word cannot be loaded twice without an erase (§26.2.1), so a
	// refused page's drained data must not linger for the next write:
	// discard the buffer up front — CTPB on the tinies; on the megas
	// writing RWWSRE aborts a pending load (§26.2.2 — on the m48s that
	// flush is the bit's whole documented job).
	if constexpr (rww_discard)
		spm::rww_enable<off>();
	else
		spm::clear_buffer<off>();
	// One induction either way. On the byte-addressed chips the wire address
	// itself walks the page (aligned, so the offset bits wrap to zero); on
	// the word-addressed large chips the wire word address becomes a 32-bit
	// byte cursor once, and their 256-byte page makes its low byte the whole
	// in-page offset. The slot index is one high byte of the wire address —
	// two values on byte-addressed chips (the & ~1), bits 16:9 re-packed on
	// the large ones.
	spm::flash_address_t address;
	std::uint8_t page_high;
	if constexpr (word_flash) {
		// Pages are aligned, so one page never crosses a 64 KiB boundary:
		// RAMPZ is a per-page constant and the fill cursor is a 16-bit Z
		// whose low byte is the whole in-page offset (256-byte pages). The
		// slot index is simply the wire word address's high byte.
		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,
			               static_cast<std::uint16_t>(low | (high << 8)));
			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) & 0xfe;
	} 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, static_cast<std::uint16_t>(low | (high << 8)));
			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) {
		// The tinies and the m48s halt the CPU through the erase and the
		// write, so only the boot-sectioned megas — running on while their
		// RWW section programs — wait.
		spm::erase_page<off>(address);
		if constexpr (boot_section)
			spm::wait();
		spm::write_page<off>(address);
		if constexpr (boot_section) {
			spm::wait();
			spm::rww_enable<off>();
		}
	}
}

// The four fuse/lock bytes in the hardware's own Z order: low, lock,
// extended, high. Writing fuses is not a thing self-programming can do on
// AVR — SPM reaches flash (and boot lock bits) only.
void send_fuses()
{
	std::uint8_t which = 0;
	do
		link::tx(spm::read_fuse<off>(static_cast<spm::fuse>(which)));
	while (++which & 3);
}

[[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.
	// The flag register is MCUSR, or the classic megas' MCUCSR.
	if (avr::hw::field_impl<wdrf_field()>::test())
		run_app();

	link::init();

	// The high byte of the 512-byte-aligned base this copy runs at: the
	// return address is a word address, whose high byte is the 256-word slot
	// index — on byte-addressed chips doubled back into byte terms.
	// program_flash refuses this one slot and the info block is addressed
	// from it, so both follow wherever the code was flashed. The high byte is
	// spelled as byteswap's low byte: the builtin's value is itself built by
	// swapping the two stacked bytes, and the double swap folds to the single
	// byte pick a hand assembler writes — `>> 8` leaves 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 & 0xfe : static_cast<std::uint8_t>(ra_high << 1);

	// The knock: 'p' then 'b', each under a fresh window; any other byte is
	// line noise and waits again. Falling out of a window runs the app.
	while (rx_deadline() != 'p' || rx_deadline() != 'b') {
	}

	for (;;) {
		// No prompt while an EEPROM write runs: a pending write blocks SPM
		// and fuse reads (§26.2.1), and the ack tells the host all is done.
		ee::wait();
		link::tx(ack);
		const std::uint8_t command = link::rx();
		switch (command) {
		case 'b': { // info block, read relative to the running slot
			// The block sits in the image's first 256 bytes (the build lint
			// asserts it), and slots are 512-aligned — so the low byte of its
			// link address (in wire units: bytes, or words on the large
			// chips) is its offset in any slot, and the high byte of its
			// runtime address is the running slot's. Composed from the two
			// bytes — the high half is runtime data, so no absolute address
			// is ever materialized.
			const auto link_low = reinterpret_cast<std::uint16_t>(info_data.data());
			const std::uint8_t low =
			    word_flash ? static_cast<std::uint8_t>(link_low >> 1) : static_cast<std::uint8_t>(link_low);
			send_flash(static_cast<std::uint16_t>(low | (slot_high << 8)), static_cast<std::uint8_t>(info_data.size()));
			break;
		}
		case 'J': { // jump to a wire word address: hand-over and staging transfer
			auto target = reinterpret_cast<void (*)()>(rx16());
			link::tx(ack);
			link::drain();
			jump(target);
		}
		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
			std::uint16_t address = rx16();
			std::uint8_t count = link::rx();
			if (command == 'R')
				send_flash(address, count);
			else if (command == 'r')
				send_eeprom(address, count);
			else
				store_eeprom(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>;