// simavr "device" for the pureboot protocol tests, every chip. Loads the
// boot-linked ELF at the loader base, starts execution there (BOOTRST / the
// patched vector are not what is under test), and exposes the loader's
// serial link as a pty for the real host tool:
//
//   - Hardware USART builds: simavr's uart_pty on the selected instance.
//   - Software UART builds: an 8N1 bridge between a pty and the GPIO pins,
//     timed against the simulated cycle counter (drives the loader's RX,
//     decodes its TX).
//
// The link follows the chip's natural default (USART0 on the megas, the
// software UART on PB0/PB1 elsewhere) unless -l overrides it: `-l usart1`
// for the second instance, `-l sw:B5,B1` for a software build's RX,TX pins,
// and `-l sw:D0,D1@0` where those pins are a USART's own — see the pin
// ownership the bridge models below.
//
// simavr's tiny cores decode the SPM opcode but attach no NVM module — SPM
// is a silent no-op (the mega's boot section has one, avr_flash). The
// missing module is supplied here: the SPM ioctl reads SPMCSR/Z/r1:r0 and
// implements buffer fill, page erase, page write, and CTPB, completing
// instantly. RFLB's LPM diversion (fuse readout) stays unmodeled, so the
// 'F' command answers with flash bytes — the tests assert transport only.
//
// On exit (or SIGTERM) the flash and EEPROM are dumped to files for a
// ground-truth cross-check against what the host read back.
#include <csignal>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <print>
#include <string_view>

#include <fcntl.h>
#include <pty.h>
#include <termios.h>
#include <unistd.h>

// The parts headers (uart_pty.h) carry no C++ linkage guards of their own,
// unlike simavr's core headers — the block covers both harmlessly.
extern "C" {
#include "avr_eeprom.h"
#include "avr_flash.h"
#include "avr_ioport.h"
#include "avr_uart.h"
#include "sim_avr.h"
#include "sim_elf.h"
#include "sim_io.h"
#include "uart_pty.h"
}

namespace {

avr_t *avr;
uart_pty_t uart_pty;
bool link_software;
char uart_digit = '0';
char sw_rx_port = 'B', sw_tx_port = 'B';
int sw_rx_bit = 0, sw_tx_bit = 1;
char sw_tx_owner = 0; // the USART whose TXD the software link sits on
const char *dump_path;
std::uint32_t reset_pc;
volatile std::sig_atomic_t reset_requested;

// -w: report the cycle of the first transmit activity, once. What the
// activation-window gate reads — with an idle line and an application
// installed, the first thing that ever talks is the application's banner,
// so this cycle *is* the loader's window plus a banner lead measured in
// microseconds. Idle pacing is skipped in this mode: there is no real-time
// host in the loop, and a paced multi-second window would take hours.
bool window_report;
bool window_tx_seen;

void window_first_tx()
{
	if (!window_report || window_tx_seen)
		return;
	window_tx_seen = true;
	std::println("PB_WINDOW_TX {}", avr->cycle);
	std::fflush(stdout);
}

void window_uart_hook(avr_irq_t *, std::uint32_t, void *)
{
	window_first_tx();
}

int parse_link(std::string_view spec)
{
	if (spec == "usart0" || spec == "usart1") {
		link_software = false;
		uart_digit = spec[5];
		return 0;
	}
	if (spec.starts_with("sw")) {
		link_software = true;
		if (spec.size() == 2)
			return 0;
		char owner = 0;
		int fields =
		    std::sscanf(spec.data() + 2, ":%c%d,%c%d@%c", &sw_rx_port, &sw_rx_bit, &sw_tx_port, &sw_tx_bit, &owner);
		if (fields == 4 || fields == 5) {
			sw_tx_owner = owner;
			return 0;
		}
	}
	return -1;
}

// simavr 1.6's avr_flash PGERS handler erases spm_pagesize bytes starting at
// Z & ~1 instead of the page containing Z (its PGWRT path masks correctly) —
// hardware ignores the in-page bits (§26.8.1), so an erase issued with Z
// anywhere inside the page wipes half the neighbouring page in simulation
// only. Wrap the mega's registered flash ioctl and re-dispatch page erases
// with Z forced to the page boundary; everything else passes through.
//
// A second gap on the boot-section-less m48s: their RWWSRE bit is the
// temporary-buffer discard (Atmel-8271 §26.2/§26.3.1), but the stock model
// gates its RWWSRE branch on AVR_SELFPROG_HAVE_RWW — absent on the m48
// core — so the discard store falls through into the buffer-fill branch and
// plants whatever Z/R1:R0 happen to hold. Perform the silicon's discard
// here instead.
avr_flash_t *mega_flash;
int (*mega_flash_ioctl)(avr_io_t *io, std::uint32_t ctl, void *param);

int fixed_flash_ioctl(avr_io_t *io, std::uint32_t ctl, void *param)
{
	if (ctl == AVR_IOCTL_FLASH_SPM && avr_regbit_get(io->avr, mega_flash->pgers)) {
		auto z = static_cast<std::uint16_t>(io->avr->data[30] | (io->avr->data[31] << 8));
		auto masked = static_cast<std::uint16_t>(z & ~(mega_flash->spm_pagesize - 1));
		io->avr->data[30] = static_cast<std::uint8_t>(masked);
		io->avr->data[31] = static_cast<std::uint8_t>(masked >> 8);
		int result = mega_flash_ioctl(io, ctl, param);
		io->avr->data[30] = static_cast<std::uint8_t>(z);
		io->avr->data[31] = static_cast<std::uint8_t>(z >> 8);
		return result;
	}
	if (ctl == AVR_IOCTL_FLASH_SPM && !(mega_flash->flags & AVR_SELFPROG_HAVE_RWW) &&
	    (io->avr->data[mega_flash->r_spm] & 0x11) == 0x11) { // RWWSRE|SELFPRGEN: the m48 buffer discard
		for (int i = 0; i < mega_flash->spm_pagesize / 2; i++) {
			mega_flash->tmppage[i] = 0xffff;
			mega_flash->tmppage_used[i] = 0;
		}
		avr_regbit_clear(io->avr, mega_flash->selfprgen);
		return 0;
	}
	return mega_flash_ioctl(io, ctl, param);
}

void fix_mega_flash_erase()
{
	for (avr_io_t *io = avr->io_port; io; io = io->next) {
		if (io->kind && std::string_view{io->kind} == "flash") {
			mega_flash = reinterpret_cast<avr_flash_t *>(io);
			mega_flash_ioctl = io->ioctl;
			io->ioctl = fixed_flash_ioctl;
			return;
		}
	}
	std::println(stderr, "device: no flash module to fix — SPM page erases may misalign");
}

void request_reset(int)
{
	reset_requested = 1;
}

// ------------------------------------------------------------- tiny NVM ---

struct tiny_nvm_t {
	avr_io_t io;
	std::uint8_t buffer[128];
	std::uint8_t used[128]; // a buffer word loads once until erased — like silicon
	unsigned page;
};

tiny_nvm_t nvm;

int nvm_ioctl(avr_io_t *io, std::uint32_t ctl, void *)
{
	if (ctl != AVR_IOCTL_FLASH_SPM)
		return -1;
	auto *n = reinterpret_cast<tiny_nvm_t *>(io);
	avr_t *mcu = io->avr;
	std::uint8_t command = mcu->data[0x57] & 0x1f; // SPMCSR, both tinies
	auto z = static_cast<std::uint16_t>(mcu->data[30] | (mcu->data[31] << 8));
	std::uint32_t page_base = static_cast<std::uint32_t>(z & ~(n->page - 1)) % (mcu->flashend + 1);
	if (command == 0x01) { // SPMEN alone: buffer fill from r1:r0
		unsigned offset = z & (n->page - 1) & ~1u;
		if (!n->used[offset]) { // first write wins until the buffer clears
			n->buffer[offset] = mcu->data[0];
			n->buffer[offset + 1] = mcu->data[1];
			n->used[offset] = 1;
		}
	} else if (command == 0x03) { // PGERS
		std::memset(mcu->flash + page_base, 0xff, n->page);
	} else if (command == 0x05) { // PGWRT: programming only clears bits
		for (unsigned i = 0; i < n->page; i++)
			mcu->flash[page_base + i] &= n->buffer[i];
		std::memset(n->buffer, 0xff, n->page);
		std::memset(n->used, 0, n->page);
	} else if (command == 0x11) { // CTPB
		std::memset(n->buffer, 0xff, n->page);
		std::memset(n->used, 0, n->page);
	}
	mcu->data[0x57] &= static_cast<std::uint8_t>(~0x1f); // the operation completes instantly
	return 0;
}

// ----------------------------------------------------------- GPIO bridge ---

int pty_master = -1;
avr_irq_t *rx_pin; // the loader's RX (PB0), driven from the pty
avr_cycle_count_t bit_cycles;

int tx_level = 1, tx_active, tx_bit;
std::uint8_t tx_shift;

avr_cycle_count_t tx_sample(avr_t *, avr_cycle_count_t when, void *)
{
	if (tx_bit < 0) {
		// Half a bit into the start bit: a real receiver re-samples here and
		// abandons a false start. The device's own init produces one — DDR
		// drives the pin low for the instructions until the idle level is
		// written — and without this check that glitch decodes as a stray
		// byte (and would read as first transmit activity under -w).
		if (tx_level) {
			tx_active = 0;
			return 0;
		}
		window_first_tx();
		tx_bit = 0;
		return when + bit_cycles;
	}
	if (tx_bit < 8) {
		tx_shift = static_cast<std::uint8_t>((tx_shift >> 1) | (tx_level ? 0x80 : 0));
		if (++tx_bit < 8)
			return when + bit_cycles;
		// The byte is delivered at the stop bit's sampling point (9.5 bit
		// times), where a hardware receiver raises its RXC — not sooner: a
		// host answering before the stop bit would put its start bit on the
		// wire while the device is still driving, which the device,
		// transmitting, is not watching for.
		return when + bit_cycles;
	}
	if (write(pty_master, &tx_shift, 1) != 1)
		std::println(stderr, "device: pty write lost a byte");
	tx_active = 0;
	return 0;
}

// A USART owns its TxD pin whenever its transmitter is enabled, and the port
// register cannot drive it (§20.2 / Atmel-8271 §19.2) — which is why a
// bit-banged link deployed on those pins is mute until it clears UCSRnB.
// simavr wires a USART entirely through IRQs and never touches the port pin
// model, so the ownership does not exist there and the mute cannot happen:
// supply it, or the very state this models is untestable. The link spec's
// trailing @n names the USART; without one the pins are nobody's.
avr_uart_t *tx_owner;

bool tx_pin_taken()
{
	return tx_owner && avr_regbit_get(avr, tx_owner->txen);
}

// simavr leaves TXEN set in UCSRnB out of reset, where silicon clears the
// whole register (§20.11.3) — which would hand the pin to a USART no code has
// enabled, making a freshly reset chip mute for reasons hardware does not
// have. Reset it the way the datasheet does, so the ownership starts from
// nobody's and only an application that really enables the USART takes it.
void reset_tx_owner()
{
	if (tx_owner)
		avr_regbit_clear(avr, tx_owner->txen);
}

void find_tx_owner()
{
	for (avr_io_t *io = avr->io_port; io; io = io->next)
		if (io->kind && std::string_view{io->kind} == "uart" &&
		    reinterpret_cast<avr_uart_t *>(io)->name == sw_tx_owner) {
			tx_owner = reinterpret_cast<avr_uart_t *>(io);
			reset_tx_owner();
			return;
		}
	std::println(stderr, "device: no USART{} to own the software link's TX pin", sw_tx_owner);
}

void tx_hook(avr_irq_t *, std::uint32_t value, void *)
{
	if (tx_pin_taken()) { // the USART holds the line; the port write goes nowhere
		tx_level = 1;
		return;
	}
	int level = value & 1;
	if (!tx_active && tx_level == 1 && level == 0) { // start edge, confirmed mid-bit
		tx_active = 1;
		tx_bit = -1;
		avr_cycle_timer_register(avr, bit_cycles / 2, tx_sample, nullptr);
	}
	tx_level = level;
}

std::uint8_t rx_queue[8192];
unsigned rx_head, rx_tail; // ring: head = next to send
int rx_active, rx_bit;
std::uint8_t rx_byte;

void rx_start_next();

avr_cycle_count_t rx_step(avr_t *, avr_cycle_count_t when, void *)
{
	if (rx_bit < 8) {
		avr_raise_irq(rx_pin, (rx_byte >> rx_bit) & 1);
		rx_bit++;
		return when + bit_cycles;
	}
	if (rx_bit == 8) { // stop bit, plus one idle bit of margin
		avr_raise_irq(rx_pin, 1);
		rx_bit++;
		return when + 2 * bit_cycles;
	}
	rx_active = 0;
	rx_start_next();
	return 0;
}

void rx_start_next()
{
	if (rx_active || rx_head == rx_tail)
		return;
	rx_byte = rx_queue[rx_head];
	rx_head = (rx_head + 1) % sizeof(rx_queue);
	rx_active = 1;
	rx_bit = 0;
	avr_raise_irq(rx_pin, 0); // start bit
	avr_cycle_timer_register(avr, bit_cycles, rx_step, nullptr);
}

// A reset abandons whatever the bridge was mid-transfer: bytes still queued
// for a chip that no longer has the context to receive them meaningfully,
// and a decode in progress on a TX line the reset may have already changed.
// The pending cycle timers must go with the state: avr_reset drops the TX
// output latch, whose falling edge starts a spurious decode before this
// runs, and a stale tx_sample would then interleave with the loader's first
// real answer through the shared shift state, corrupting it.
void bridge_reset()
{
	avr_cycle_timer_cancel(avr, tx_sample, nullptr);
	avr_cycle_timer_cancel(avr, rx_step, nullptr);
	rx_head = rx_tail = 0;
	rx_active = 0;
	tx_active = 0;
	tx_level = 1;
	// Re-drive the idle line through a forced transition: ioport pin irqs are
	// IRQ_FLAG_FILTERED, and avr_reset zeroes the port latch while the irq
	// keeps its pre-reset cached value — so a plain raise(1) against a cached
	// 1 is dropped and the device reads the line stuck low. A loader entering
	// calibration on that line measures reset-to-first-edge as one giant
	// pulse and mis-locks or boots the application on the first real knock.
	// No cycles run between the two raises, so the device only ever sees the
	// final idle-high.
	avr_raise_irq(rx_pin, 0);
	avr_raise_irq(rx_pin, 1);
}

void poll_pty()
{
	std::uint8_t chunk[256];
	ssize_t got = read(pty_master, chunk, sizeof(chunk));
	for (ssize_t i = 0; i < got; i++) {
		unsigned next = (rx_tail + 1) % sizeof(rx_queue);
		if (next == rx_head)
			break; // full: the host will retry on timeout
		rx_queue[rx_tail] = chunk[i];
		rx_tail = next;
	}
	if (got > 0)
		rx_start_next();
}

// ------------------------------------------------------------------ main ---

[[noreturn]] void finish(int)
{
	if (dump_path) {
		std::FILE *f = std::fopen(dump_path, "wb");
		if (f) {
			std::fwrite(avr->flash, 1, avr->flashend + 1, f);
			std::fclose(f);
		}
		avr_eeprom_desc_t ee = {.ee = nullptr, .offset = 0, .size = 0};
		if (avr_ioctl(avr, AVR_IOCTL_EEPROM_GET, &ee) == 0 && ee.ee && ee.size) {
			char path[512];
			std::snprintf(path, sizeof(path), "%s.eeprom", dump_path);
			f = std::fopen(path, "wb");
			if (f) {
				std::fwrite(ee.ee, 1, ee.size, f);
				std::fclose(f);
			}
		}
	}
	if (!link_software)
		uart_pty_stop(&uart_pty);
	_exit(0);
}

} // namespace

int main(int argc, char *argv[])
{
	bool link_given = false;
	for (int opt; (opt = getopt(argc, argv, "l:w")) != -1;) {
		if (opt == 'w') {
			window_report = true;
			continue;
		}
		if (opt != 'l' || parse_link(optarg) != 0) {
			std::println(stderr, "device: bad link spec (usart0, usart1, sw, or sw:B0,B1 as RX,TX)");
			return 2;
		}
		link_given = true;
	}
	int args = argc - optind;
	if (args < 7 || args > 9) {
		std::print(stderr,
		           "usage: {} [-l link] [-w] <pureboot.elf> <mcu> <hz> <base_hex> <page> <baud> <flash_dump>"
		           " [reset_hex] [resume_flash]\n"
		           "  -l link:      usart0 | usart1 | sw[:B0,B1[@0]] (RX,TX, then the USART owning\n"
		           "                them); default: the chip's own\n"
		           "  -w:           print PB_WINDOW_TX <cycle> at the first transmit activity and\n"
		           "                free-run idle time (window measurement mode)\n"
		           "  reset_hex:    reset vector (default: base with a boot section, else 0)\n"
		           "  resume_flash: raw full-flash image loaded instead of the ELF — a prior\n"
		           "                run's dump, for power-fail resume tests\n",
		           argv[0]);
		return 2;
	}
	argv += optind - 1; // argv[1] is the ELF again, whatever was parsed
	const std::string_view mcu_name = argv[2];
	auto base = static_cast<std::uint32_t>(std::strtoul(argv[4], nullptr, 0));
	auto page = static_cast<unsigned>(std::atoi(argv[5]));
	auto baud = static_cast<unsigned>(std::atoi(argv[6]));
	dump_path = argv[7];
	const bool is_mega = mcu_name.starts_with("atmega");
	if (!link_given)
		link_software = !is_mega; // the chips' natural links: USART0, or PB0/PB1

	avr = avr_make_mcu_by_name(mcu_name.data());
	if (!avr) {
		std::println(stderr, "device: no {} core", mcu_name);
		return 1;
	}
	avr_init(avr);
	avr->frequency = static_cast<std::uint32_t>(std::strtoul(argv[3], nullptr, 0));
	std::memset(avr->flash, 0xff, avr->flashend + 1); // real flash powers up erased

	if (args > 8) {
		// Resume: the full flash image of an interrupted prior run.
		std::FILE *f = std::fopen(argv[9], "rb");
		if (!f || std::fread(avr->flash, 1, avr->flashend + 1, f) == 0) {
			std::println(stderr, "device: cannot read {}", argv[9]);
			return 1;
		}
		std::fclose(f);
	} else {
		elf_firmware_t fw{};
		if (elf_read_firmware(argv[1], &fw) != 0) {
			std::println(stderr, "device: cannot read {}", argv[1]);
			return 1;
		}
		// An image past flash end would smash the simulator's heap and turn
		// into phantom peripheral behavior (lessons: believe the size gate
		// first) — refuse it loudly instead.
		if (base + fw.flashsize > avr->flashend + 1) {
			std::println(stderr, "device: {} B at {:#x} runs past flash end {:#x} — image does not fit its slot",
			             fw.flashsize, base, avr->flashend);
			return 1;
		}
		std::memcpy(avr->flash + base, fw.flash, fw.flashsize);
	}
	// The boot-sectioned megas enter the loader in hardware (BOOTRST, not
	// modeled — the argument picks the modeled fuse's target); the tinies
	// and the boot-section-less m48s reset to word 0 like silicon — erased
	// flash walks up into the loader, and after the host's surgery the
	// patched vector routes there.
	const bool boot_section = is_mega && !mcu_name.starts_with("atmega48");
	reset_pc = args > 7 ? static_cast<std::uint32_t>(std::strtoul(argv[8], nullptr, 0)) : (boot_section ? base : 0);
	avr->pc = reset_pc;
	avr->codeend = avr->flashend;

	// Erased EEPROM, as hardware powers up (simavr zeroes it).
	std::uint8_t blank[1024];
	std::memset(blank, 0xff, sizeof(blank));
	avr_eeprom_desc_t seed = {.ee = blank, .offset = 0, .size = 0};
	if (avr_ioctl(avr, AVR_IOCTL_EEPROM_GET, &seed) == 0 && seed.size <= sizeof(blank)) {
		seed.ee = blank;
		avr_ioctl(avr, AVR_IOCTL_EEPROM_SET, &seed);
	}

	// The megas carry simavr's avr_flash module (and its two gaps the wrap
	// above fixes); the tinies get the NVM module simavr lacks. Which serial
	// bridge runs is the link's business, not the chip class's.
	if (is_mega) {
		fix_mega_flash_erase();
	} else {
		nvm.page = page;
		std::memset(nvm.buffer, 0xff, sizeof(nvm.buffer));
		nvm.io.kind = "tiny_nvm";
		nvm.io.ioctl = nvm_ioctl;
		avr_register_io(avr, &nvm.io);
	}

	if (!link_software) {
		// POLL_SLEEP paces an idle-polling loader in host real time (a
		// no-hardware CPU-saving hack); clear it so cycles run free.
		std::uint32_t flags = 0;
		avr_ioctl(avr, AVR_IOCTL_UART_GET_FLAGS(uart_digit), &flags);
		flags &= ~AVR_UART_FLAG_POLL_SLEEP;
		avr_ioctl(avr, AVR_IOCTL_UART_SET_FLAGS(uart_digit), &flags);
		uart_pty_init(avr, &uart_pty);
		uart_pty_connect(&uart_pty, uart_digit);
		if (window_report)
			avr_irq_register_notify(avr_io_getirq(avr, AVR_IOCTL_UART_GETIRQ(uart_digit), UART_IRQ_OUTPUT),
			                        window_uart_hook, nullptr);
		std::println("PB_PTY {}", uart_pty.pty.slavename);
	} else {
		bit_cycles = (avr->frequency + baud / 2) / baud; // matches uart.hpp's own rounding exactly
		if (sw_tx_owner)
			find_tx_owner();
		rx_pin = avr_io_getirq(avr, AVR_IOCTL_IOPORT_GETIRQ(sw_rx_port), static_cast<unsigned>(sw_rx_bit));
		avr_irq_register_notify(
		    avr_io_getirq(avr, AVR_IOCTL_IOPORT_GETIRQ(sw_tx_port), static_cast<unsigned>(sw_tx_bit)), tx_hook,
		    nullptr);
		avr_raise_irq(rx_pin, 1); // idle line

		int slave;
		struct termios raw;
		cfmakeraw(&raw);
		if (openpty(&pty_master, &slave, nullptr, &raw, nullptr) != 0) {
			std::println(stderr, "device: openpty failed");
			return 1;
		}
		fcntl(pty_master, F_SETFL, O_NONBLOCK);
		std::println("PB_PTY {}", ttyname(slave));
	}
	std::fflush(stdout);

	std::signal(SIGTERM, finish);
	std::signal(SIGINT, finish);
	std::signal(SIGUSR1, request_reset); // an external reset line, for the tests

	long since_poll = 0;
	for (;;) {
		int state = avr_run(avr);
		if (state == cpu_Done || state == cpu_Crashed)
			break;
		if (reset_requested) {
			reset_requested = 0;
			avr_reset(avr);
			avr->pc = reset_pc;
			if (!link_software) { // reset restores the pacing hack; re-clear it
				std::uint32_t flags = 0;
				avr_ioctl(avr, AVR_IOCTL_UART_GET_FLAGS(uart_digit), &flags);
				flags &= ~AVR_UART_FLAG_POLL_SLEEP;
				avr_ioctl(avr, AVR_IOCTL_UART_SET_FLAGS(uart_digit), &flags);
			} else {
				bridge_reset();
				reset_tx_owner();
			}
		}
		if (link_software && ++since_poll >= 2000) {
			since_poll = 0;
			poll_pty();
			// An unthrottled idle simulation runs the activation window out
			// from under the host's real-time knock cadence: a 1 MHz build's
			// 8 s window is 8 M cycles — tens of wall milliseconds — so a
			// first knock lost to an in-flight reset misses the window
			// entirely. Pace the simulation only while the bridge is fully
			// quiet (nothing decoding, nothing queued); transfers keep full
			// speed, and a quiet window stretches toward real time.
			if (!window_report && !rx_active && !tx_active && rx_head == rx_tail)
				usleep(200);
		}
	}
	finish(0);
}