// pureboot autobaud — the software-serial variant that measures the host's bit
// timing at runtime, so one binary runs at any F_CPU: the image carries no
// clock. THE REGISTER VERSION — keeps the running-slot write guard, at the cost
// of one global register variable (r4) holding the measured unit. That variable
// is pureboot's single, deliberate break from its no-global-register-variable
// rule, present only in this variant; everything else stays pure C++.
//
// Two source files exist for review (autobaud.md):
// pureboot_autobaud_pure.cpp, fully pure but dropping the write guard, and this
// one. They differ only in where the measured unit lives (a call-saved register
// here, GPIOR/RAM there) and whether the guard is present.
//
// The register buys ~32 B over a RAM home — an outlined rx/tx reads it with one
// move where a static costs an lds — and that is what lets the write guard stay
// while the image still fits 512 B (the 1284 at 512, exactly). The unit is
// written once through a noinline setter so the store lands immediately before a
// ret: GCC otherwise deletes a global-register store whose only readers are
// callees (autobaud.md, upstream bug 6).
//
// Simplifications shared with the pure version, each licensed: a slimmed info
// block (version + signature; the host derives geometry from its chip database)
// and a single-byte activation knock (the calibration pulse already proves a
// host). Position independence is kept: control flow is PC-relative and the
// write guard anchors on the runtime return address, as the fixed-baud loader.

#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;

// The measured per-bit delay (in _delay_loop_2 four-cycle iterations) in a
// call-saved register that the serial callees read directly, so no wire path
// threads it and rx/tx reach it with a move, not a load. Written only through
// set_unit() below.
register std::uint16_t g_unit asm("r4");

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; only the software-UART pins are a parameter.
#if !defined(PUREBOOT_RX)
#define PUREBOOT_RX pb0
#endif
#if !defined(PUREBOOT_TX)
#define PUREBOOT_TX pb1
#endif

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

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();
constexpr bool word_flash = spm::flash_bytes > 65536;

constexpr std::uint8_t version = 4;

#if !defined(PUREBOOT_AUTOBAUD_POLLS)
#define PUREBOOT_AUTOBAUD_POLLS 4000000
#endif
constexpr __uint24 autobaud_budget = PUREBOOT_AUTOBAUD_POLLS;

// The one store into g_unit, isolated so it lands right before the ret: a
// global-register store whose only later readers are callees is dropped
// otherwise (autobaud.md, upstream bug 6).
[[gnu::noinline]] void set_unit(std::uint16_t v)
{
	g_unit = v;
}

// The autobaud software link: bit-banged with cycle-counted delays, but the
// per-bit delay is g_unit, measured from the host's calibration pulse.
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 g_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 count >> 2 is the bit
	// period in _delay_loop_2's four-cycle iterations. 0 means the budget
	// expired.
	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);
		set_unit(u);
		return u;
	}

	static std::uint8_t rx()
	{
		while (in_t::read()) // await the start edge
			;
		_delay_loop_2(static_cast<std::uint16_t>(g_unit + (g_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(g_unit);
		}
		return value;
	}

	static void tx(std::uint8_t value)
	{
		out_t::clear(); // start bit
		_delay_loop_2(g_unit);
		for (std::uint8_t bit = 0; bit < 8; ++bit) {
			out_t::write(value & 1);
			value >>= 1;
			_delay_loop_2(g_unit);
		}
		out_t::set(); // stop bit
		_delay_loop_2(g_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);
}

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

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

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

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

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.
void program_flash(std::uint16_t wire_address, std::uint8_t slot_high)
{
	spm::flash_address_t address;
	std::uint8_t page_high;
	if constexpr (word_flash) {
		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) & ~static_cast<std::uint16_t>(page - 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 & ~static_cast<std::uint16_t>(page - 1));
		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) {
		spm::erase_page<off>(address);
		if constexpr (boot_section)
			spm::wait();
		spm::write_page<off>(address);
		if constexpr (boot_section)
			spm::wait();
	}
	if constexpr (boot_section)
		spm::rww_enable<off>();
}

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()
{
	if (hw::field_impl<wdrf_field()>::test())
		run_app();

	link::init();

	// The high byte of the slot this copy runs at, which the write guard
	// follows: the return address is a word address, so its high byte is the
	// 256-word slot index, doubled back into byte terms on a byte-addressed
	// chip. Taken as byteswap's low byte — the builtin already swaps the two
	// stacked bytes, and the double swap folds away.
	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);

	// Measure the calibration pulse into g_unit, then take one 'p' knock. An
	// expired budget (no host) boots the application; a pulse that decodes to
	// anything but 'p' re-measures.
	for (;;) {
		if (link::measure(autobaud_budget) == 0)
			run_app();
		if (link::rx() == '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 '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 the three, so the flash streamer
			// keeps a single call site and inlines into this never-returning
			// loop (autobaud.md).
			const std::uint16_t address = rx16();
			const std::uint8_t 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>;