// pureboot autobaud, the unified-primitive version — VARIANT A, an explicit
// space byte. One read command and one write command carry a space selector, so
// flash, EEPROM, RAM and the fuses share a single cursor, a single transfer loop
// and a single argument decode instead of one command body each.
//
// Strictly pure: no inline assembly, no global register variables, and no GPIOR
// either — the measured unit lives in a plain static, so the loader claims no
// chip resource an application might want. (PUREBOOT_UNIT_GPIOR=1 puts it back
// in the I/O scratch registers, kept only as a measurement axis.)
//
// Against pureboot_autobaud_pure.cpp this version:
//   - adds RAM read and write, which the loader has never had. Because AVR maps
//     the register file and the whole I/O space into the data address space,
//     that one space also gives the host arbitrary peripheral access for free;
//   - collapses 'R' (read flash), 'r' (read EEPROM), 'w' (write EEPROM) and 'F'
//     (fuses) — four bodies, four loops — into 'G' and 'P' over four spaces;
//   - fixes the activation hang: a lone calibration pulse used to leave the
//     loader blocked forever in the knock's rx(), so a stray edge on an
//     unattended device wedged it in the loader and the application never ran.
//
// Position independence is kept and is total: control flow is PC-relative, the
// wire carries addresses, and nothing anchors on the runtime address.

#include <libavr/libavr.hpp>

#include <util/delay_basic.h>

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

namespace pureboot {
namespace {

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

constexpr std::uint8_t ack = '+';

// Autobaud carries no clock, so PUREBOOT_CLOCK_HZ / PUREBOOT_BAUD are not
// consulted; only the software-UART pins are a deployment parameter.
#if !defined(PUREBOOT_RX)
#define PUREBOOT_RX pb0
#endif
#if !defined(PUREBOOT_TX)
#define PUREBOOT_TX pb1
#endif

// The unit's home. A plain static by default — pureboot claims no GPIOR, so the
// application keeps both scratch registers. The GPIOR spelling is retained
// behind a macro purely so the two can be measured against each other.
#if !defined(PUREBOOT_UNIT_GPIOR)
#define PUREBOOT_UNIT_GPIOR 0
#endif

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

// The loader owns the top 512 bytes; a staging copy goes in the slot below.
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 = 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).
constexpr bool word_flash = spm::flash_bytes > 65536;

// The loader's one identity number (README.md); the slimmed info block carries
// it and the signature, and the host maps a version to its protocol.
constexpr std::uint8_t version = 5;

// The activation window as a fixed poll budget: with no clock, whole seconds
// cannot be timed. A __uint24 (AVR's three-byte type) holds it — a fourth byte
// would cost two words at each countdown step for range never used.
#if !defined(PUREBOOT_AUTOBAUD_POLLS)
#define PUREBOOT_AUTOBAUD_POLLS 4000000
#endif
constexpr __uint24 autobaud_budget = PUREBOOT_AUTOBAUD_POLLS;

// The two SPM commands the loader still has to recognise by value, on the chips
// where it cannot issue a runtime one generically. Taken from the chip's own
// definitions rather than spelled 3 and 5 — though every part pureboot targets
// agrees on those, which is what lets the host send the raw SPMCSR byte.
constexpr std::uint8_t spm_erase = __BOOT_PAGE_ERASE;
constexpr std::uint8_t spm_write = __BOOT_PAGE_WRITE;

// The spaces a transfer can name. Flash is 0 so it is the cheap default.
//
// sp_spm is the one that is not memory: a write there hands its data byte to
// SPMCSR and fires the instruction at the given flash address, so page erase,
// page write and RWW re-enable become host-issued commands instead of a
// hardcoded tail inside 'W'. The store side already owns an address, a data
// byte and an ack, so the whole sequence costs only the fused out/spm pair. It
// also lets the host reach every other SPM operation — lock bits included —
// which the loader previously had no way to expose.
enum : std::uint8_t { sp_flash = 0, sp_eeprom = 1, sp_ram = 2, sp_fuse = 3, sp_spm = 4 };

// A transfer's selector byte is `space | bank << 4`: the low nibble names the
// space, the high nibble carries flash's third address byte (RAMPZ) on the
// chips that have one. Putting the bank here rather than widening the address
// keeps the shared cursor sixteen bits for every space — a three-byte cursor
// costs its extra increment on EEPROM and RAM reads too, which never need it.
// The host must not span a bank boundary in one transfer; it already chunks by
// page, so nothing it does today comes close.

// .noinit, not .bss: the unit is always measured before it is read, so it needs
// no zeroing — and a zeroed .bss would drag in __do_clear_bss, 18 bytes of
// startup code for a variable that is written before its first use.
[[gnu::section(".noinit")]] std::uint16_t unit_backing;

[[gnu::always_inline]] inline std::uint16_t get_unit()
{
#if PUREBOOT_UNIT_GPIOR
	return static_cast<std::uint16_t>(hw::reg_impl<hw::db.reg_index("GPIOR1")>::read() |
	                                  (hw::reg_impl<hw::db.reg_index("GPIOR2")>::read() << 8));
#else
	return unit_backing;
#endif
}

[[gnu::always_inline]] inline void put_unit(std::uint16_t u)
{
#if PUREBOOT_UNIT_GPIOR
	hw::reg_impl<hw::db.reg_index("GPIOR1")>::write(static_cast<std::uint8_t>(u));
	hw::reg_impl<hw::db.reg_index("GPIOR2")>::write(static_cast<std::uint8_t>(u >> 8));
#else
	unit_backing = u;
#endif
}

// The autobaud software link: bit-banged with cycle-counted delays like the
// fixed-baud software backend, but the per-bit delay is the measured unit, not
// a consteval constant. rx/tx load the unit once into a local, so each bit
// spins from a register with no per-bit reload.
struct link {
	using in_t = avr::io::input<avr::PUREBOOT_RX, avr::io::pull::up>;
	using out_t = avr::io::output<avr::PUREBOOT_TX>;

	static void init()
	{
		avr::init<in_t, out_t>();
		out_t::set(); // idle high
	}

	// Time the calibration pulse into the unit. The host sends 0xC0 — a start
	// bit plus six zero data bits are one low pulse of seven bit-times — and the
	// counted poll loop is seven cycles an iteration, so the count is the pulse
	// length in cycles ÷ 7 × 7 = one bit period in cycles, and count >> 2 is that
	// period in _delay_loop_2's four-cycle iterations. Waits for the start edge
	// under the poll budget; 0 (returned, and stored) means the budget expired.
	//
	// The shape of the counting loop is load-bearing, not incidental: the >> 2
	// is exact only while the pulse's bit-count equals the loop's cycles per
	// iteration. Both are 7 here (sbis 1 + rjmp 2 + adiw 2 + rjmp 2). Reshaping
	// this loop silently changes the lock; test/pbautobaud.py is what pins it.
	static std::uint16_t measure(__uint24 budget)
	{
		while (in_t::read())
			if (!--budget)
				return 0;
		std::uint16_t count = 0;
		while (!in_t::read())
			++count;
		const std::uint16_t u = static_cast<std::uint16_t>(count >> 2);
		put_unit(u);
		return u;
	}

	// The eight data bits, entered just past the falling edge of the start bit.
	// Split out of rx() so the activation path can wait for that edge under a
	// budget while the command loop waits for it indefinitely.
	static std::uint8_t sample()
	{
		const std::uint16_t unit = get_unit();
		_delay_loop_2(static_cast<std::uint16_t>(unit + (unit >> 1))); // 1.5 bits to the LSB centre
		std::uint8_t value = 0;
		for (std::uint8_t bit = 0; bit < 8; ++bit) {
			value >>= 1;
			if (in_t::read())
				value |= 0x80;
			_delay_loop_2(unit);
		}
		return value;
	}

	static std::uint8_t rx()
	{
		while (in_t::read()) // await the start edge
			;
		return sample();
	}

	// A byte under the poll budget, for the activation knock. An expired budget
	// returns 0, which is not the knock, so the caller falls back into the
	// budgeted measure() — and a line that stays idle boots the application
	// there. That is the whole hang fix: no wait during activation is unbounded,
	// so a stray calibration pulse can no longer wedge the loader.
	static std::uint8_t rx_bounded(__uint24 budget)
	{
		while (in_t::read())
			if (!--budget)
				return 0;
		return sample();
	}

	static void tx(std::uint8_t value)
	{
		const std::uint16_t unit = get_unit();
		out_t::clear(); // start bit
		_delay_loop_2(unit);
		for (std::uint8_t bit = 0; bit < 8; ++bit) {
			out_t::write(value & 1);
			value >>= 1;
			_delay_loop_2(unit);
		}
		out_t::set(); // stop bit
		_delay_loop_2(unit);
	}

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

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

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

[[gnu::always_inline]] inline std::uint16_t word_of(std::array<std::uint8_t, 2> pair)
{
	return std::bit_cast<std::uint16_t>(pair);
}

[[gnu::noinline]] void tx_ack()
{
	link::tx(ack);
}

// One byte out of any space. The four accessors share the cursor, the loop and
// the call site — the whole point of the unified commands — so each costs only
// its own instruction rather than a body, a loop and a dispatch arm.
[[gnu::always_inline]] inline std::uint8_t load(std::uint8_t space, std::uint8_t bank, std::uint16_t at)
{
	if (space == sp_eeprom)
		return ee::read(at);
	if (space == sp_ram)
		return *reinterpret_cast<volatile std::uint8_t *>(at);
	if (space == sp_fuse)
		return spm::read_fuse<off>(static_cast<spm::fuse>(at));
	if constexpr (word_flash)
		return avr::flash_load_far<std::uint8_t>((static_cast<std::uint32_t>(bank) << 16) | 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 the fuses are not writable at all.
[[gnu::always_inline]] inline void store(std::uint8_t space, [[maybe_unused]] std::uint8_t bank, std::uint16_t at,
                                         std::uint8_t value)
{
	if (space == sp_ram) {
		*reinterpret_cast<volatile std::uint8_t *>(at) = value;
		return;
	}
	if (space == sp_spm) {
		// The fused store-and-fire: SPMCSR takes the data byte and the SPM
		// issues against Z in the same unscheduled pair the hardware's
		// four-cycle window demands, which is exactly why this is one
		// primitive and not a poke of SPMCSR followed by a poke of anything
		// else. A host cannot hit that window across a serial link.
		// The preprocessor rather than `if constexpr` only because
		// spm::detail::page_command does not exist at all where there is no
		// RAMPZ, and a discarded constexpr branch outside a template is still
		// name-checked. A generic spm::command() in libavr would let the
		// runtime command through on every chip and retire the dispatch below.
#if defined(RAMPZ)
		spm::detail::page_command(value, (static_cast<spm::flash_address_t>(bank) << 16) | at);
#else
		if (value == spm_erase)
			spm::erase_page<off>(at);
		else if (value == spm_write)
			spm::write_page<off>(at);
		else if constexpr (boot_section)
			spm::rww_enable<off>();
#endif
		if constexpr (boot_section)
			spm::wait();
		return;
	}
	ee::write<off>(at, value);
}

// One page into the SPM buffer, and only that: the erase, the write and the RWW
// re-enable that used to follow are now three host-issued writes to sp_spm,
// which reach the same fused out/spm pair through the store path's own address
// and data. No running-slot write guard: the host guarantees it never targets
// the loader's own slot (licensed — README.md).
void program_flash([[maybe_unused]] std::uint8_t bank, std::uint16_t at)
{
	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();
#if defined(RAMPZ)
		spm::fill<off>((static_cast<spm::flash_address_t>(bank) << 16) | z, word_of({low, high}));
#else
		spm::fill<off>(z, word_of({low, high}));
#endif
		z += 2;
	} while (static_cast<std::uint8_t>(z) & (page - 1));
}

[[noreturn]] void run()
{
	if (hw::field_impl<wdrf_field()>::test())
		run_app();

	link::init();

	// Measure the calibration pulse into the unit, then take one 'p' knock —
	// both under the poll budget. An expired budget (no host) boots the
	// application; anything but 'p', including the knock timing out, re-measures
	// and so returns to the budgeted wait that boots it.
	for (;;) {
		if (link::measure(autobaud_budget) == 0)
			run_app();
		if (link::rx_bounded(autobaud_budget) == 'p')
			break;
	}

	for (;;) {
		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': // chip identity: version then the three signature bytes
			link::tx(version);
			link::tx(hw::db.signature[0]);
			link::tx(hw::db.signature[1]);
			link::tx(hw::db.signature[2]);
			break;
		case 'W':   // program one flash page: sel8, addr16, then page bytes
		case 'G':   // read:  sel8, addr16, n8 (0 = 256)
		case 'g': { // write: sel8, addr16, n8, then n bytes each acked
			// The unified transfer. One decode, one cursor, one loop for every
			// space and both directions — the four command bodies this replaces
			// each carried their own copy of all three. Read and write are the
			// same letter in the two cases, so the direction is bit 5 of the
			// command and the loop tests it with a one-word skip. 'W' joins the
			// same selector-and-address decode rather than keeping a word
			// address of its own, which makes flash addressing uniform across
			// every command that names it and costs nothing to share.
			const std::uint8_t sel = link::rx();
			const std::uint8_t space = sel & 0x0f;
			const std::uint8_t bank = static_cast<std::uint8_t>(sel >> 4);
			std::uint16_t at = rx16();
			if (command == 'W') {
				program_flash(bank, at);
				break;
			}
			std::uint8_t count = link::rx();
			do {
				if (command & 0x20) {
					store(space, bank, at, link::rx());
					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

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