// 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.
//
// 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 <fcntl.h>
#include <pty.h>
#include <signal.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <termios.h>
#include <unistd.h>

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

static avr_t *avr;
static uart_pty_t uart_pty;
static int link_software;
static char uart_digit = '0';
static char sw_rx_port = 'B', sw_tx_port = 'B';
static int sw_rx_bit = 0, sw_tx_bit = 1;
static const char *dump_path;
static uint32_t reset_pc;
static volatile sig_atomic_t reset_requested;

static int parse_link(const char *spec)
{
	if (strcmp(spec, "usart0") == 0 || strcmp(spec, "usart1") == 0) {
		link_software = 0;
		uart_digit = spec[5];
		return 0;
	}
	if (strncmp(spec, "sw", 2) == 0) {
		link_software = 1;
		if (spec[2] == '\0')
			return 0;
		if (sscanf(spec + 2, ":%c%d,%c%d", &sw_rx_port, &sw_rx_bit, &sw_tx_port, &sw_tx_bit) == 4)
			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.
static avr_flash_t *mega_flash;
static int (*mega_flash_ioctl)(avr_io_t *io, uint32_t ctl, void *param);

static int fixed_flash_ioctl(avr_io_t *io, uint32_t ctl, void *param)
{
	if (ctl == AVR_IOCTL_FLASH_SPM && avr_regbit_get(io->avr, mega_flash->pgers)) {
		uint16_t z = (uint16_t)(io->avr->data[30] | (io->avr->data[31] << 8));
		uint16_t masked = (uint16_t)(z & ~(mega_flash->spm_pagesize - 1));
		io->avr->data[30] = (uint8_t)masked;
		io->avr->data[31] = (uint8_t)(masked >> 8);
		int result = mega_flash_ioctl(io, ctl, param);
		io->avr->data[30] = (uint8_t)z;
		io->avr->data[31] = (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);
}

static void fix_mega_flash_erase(void)
{
	for (avr_io_t *io = avr->io_port; io; io = io->next) {
		if (io->kind && strcmp(io->kind, "flash") == 0) {
			mega_flash = (avr_flash_t *)io;
			mega_flash_ioctl = io->ioctl;
			io->ioctl = fixed_flash_ioctl;
			return;
		}
	}
	fprintf(stderr, "device: no flash module to fix — SPM page erases may misalign\n");
}

static void request_reset(int sig)
{
	(void)sig;
	reset_requested = 1;
}

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

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

static tiny_nvm_t nvm;

static int nvm_ioctl(avr_io_t *io, uint32_t ctl, void *param)
{
	(void)param;
	if (ctl != AVR_IOCTL_FLASH_SPM)
		return -1;
	tiny_nvm_t *n = (tiny_nvm_t *)io;
	avr_t *mcu = io->avr;
	uint8_t command = mcu->data[0x57] & 0x1f; // SPMCSR, both tinies
	uint16_t z = (uint16_t)(mcu->data[30] | (mcu->data[31] << 8));
	uint32_t page_base = (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
		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];
		memset(n->buffer, 0xff, n->page);
		memset(n->used, 0, n->page);
	} else if (command == 0x11) { // CTPB
		memset(n->buffer, 0xff, n->page);
		memset(n->used, 0, n->page);
	}
	mcu->data[0x57] &= (uint8_t)~0x1f; // the operation completes instantly
	return 0;
}

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

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

static int tx_level = 1, tx_active, tx_bit;
static uint8_t tx_shift;

static avr_cycle_count_t tx_sample(avr_t *mcu, avr_cycle_count_t when, void *param)
{
	(void)mcu;
	(void)param;
	if (tx_bit < 8) {
		tx_shift = (uint8_t)((tx_shift >> 1) | (tx_level ? 0x80 : 0));
		if (++tx_bit < 8)
			return when + bit_cycles;
		/* The byte is not delivered until its stop bit has passed. A real
		 * receiver cannot answer sooner, and a host that did would put its
		 * start bit on the wire while the device is still driving the stop
		 * bit — which the device, transmitting, is not watching for. */
		return when + bit_cycles;
	}
	if (write(pty_master, &tx_shift, 1) != 1)
		fprintf(stderr, "device: pty write lost a byte\n");
	tx_active = 0;
	return 0;
}

static void tx_hook(avr_irq_t *irq, uint32_t value, void *param)
{
	(void)irq;
	(void)param;
	int level = value & 1;
	if (!tx_active && tx_level == 1 && level == 0) { // start edge
		tx_active = 1;
		tx_bit = 0;
		avr_cycle_timer_register(avr, bit_cycles + bit_cycles / 2, tx_sample, NULL);
	}
	tx_level = level;
}

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

static void rx_start_next(void);

static avr_cycle_count_t rx_step(avr_t *mcu, avr_cycle_count_t when, void *param)
{
	(void)mcu;
	(void)param;
	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;
}

static void rx_start_next(void)
{
	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, NULL);
}

// 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.
static void bridge_reset(void)
{
	avr_cycle_timer_cancel(avr, tx_sample, NULL);
	avr_cycle_timer_cancel(avr, rx_step, NULL);
	rx_head = rx_tail = 0;
	rx_active = 0;
	tx_active = 0;
	tx_level = 1;
	avr_raise_irq(rx_pin, 1); // idle line
}

static void poll_pty(void)
{
	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 ---

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

int main(int argc, char *argv[])
{
	int link_given = 0;
	for (int opt; (opt = getopt(argc, argv, "l:")) != -1;) {
		if (opt != 'l' || parse_link(optarg) != 0) {
			fprintf(stderr, "device: bad link spec (usart0, usart1, sw, or sw:B0,B1 as RX,TX)\n");
			return 2;
		}
		link_given = 1;
	}
	int args = argc - optind;
	if (args < 7 || args > 9) {
		fprintf(stderr,
		        "usage: %s [-l link] <pureboot.elf> <mcu> <hz> <base_hex> <page> <baud> <flash_dump>"
		        " [reset_hex] [resume_flash]\n"
		        "  -l link:      usart0 | usart1 | sw[:B0,B1] (RX,TX); default: the chip's own\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 char *mcu_name = argv[2];
	uint32_t base = (uint32_t)strtoul(argv[4], NULL, 0);
	unsigned page = (unsigned)atoi(argv[5]);
	unsigned baud = (unsigned)atoi(argv[6]);
	dump_path = argv[7];
	int is_mega = strncmp(mcu_name, "atmega", 6) == 0;
	if (!link_given)
		link_software = !is_mega; // the chips' natural links: USART0, or PB0/PB1

	avr = avr_make_mcu_by_name(mcu_name);
	if (!avr) {
		fprintf(stderr, "device: no %s core\n", mcu_name);
		return 1;
	}
	avr_init(avr);
	avr->frequency = (uint32_t)strtoul(argv[3], NULL, 0);
	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.
		FILE *f = fopen(argv[9], "rb");
		if (!f || fread(avr->flash, 1, avr->flashend + 1, f) == 0) {
			fprintf(stderr, "device: cannot read %s\n", argv[9]);
			return 1;
		}
		fclose(f);
	} else {
		elf_firmware_t fw = {0};
		if (elf_read_firmware(argv[1], &fw) != 0) {
			fprintf(stderr, "device: cannot read %s\n", argv[1]);
			return 1;
		}
		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.
	int boot_section = is_mega && strncmp(mcu_name, "atmega48", 8) != 0;
	reset_pc = args > 7 ? (uint32_t)strtoul(argv[8], NULL, 0) : (boot_section ? base : 0);
	avr->pc = reset_pc;
	avr->codeend = avr->flashend;

	// Erased EEPROM, as hardware powers up (simavr zeroes it).
	uint8_t blank[1024];
	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;
		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.
		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);
		printf("PB_PTY %s\n", uart_pty.pty.slavename);
	} else {
		bit_cycles = (avr->frequency + baud / 2) / baud; // matches uart.hpp's own rounding exactly
		rx_pin = avr_io_getirq(avr, AVR_IOCTL_IOPORT_GETIRQ(sw_rx_port), (unsigned)sw_rx_bit);
		avr_irq_register_notify(avr_io_getirq(avr, AVR_IOCTL_IOPORT_GETIRQ(sw_tx_port), (unsigned)sw_tx_bit), tx_hook,
		                        NULL);
		avr_raise_irq(rx_pin, 1); // idle line

		int slave;
		struct termios raw;
		cfmakeraw(&raw);
		if (openpty(&pty_master, &slave, NULL, &raw, NULL) != 0) {
			fprintf(stderr, "device: openpty failed\n");
			return 1;
		}
		fcntl(pty_master, F_SETFL, O_NONBLOCK);
		printf("PB_PTY %s\n", ttyname(slave));
	}
	fflush(stdout);

	signal(SIGTERM, finish);
	signal(SIGINT, finish);
	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
				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();
			}
		}
		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 (!rx_active && !tx_active && rx_head == rx_tail)
				usleep(200);
		}
	}
	finish(0);
	return 0;
}