// simavr "device" for the pureboot protocol tests, all three chips. 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:
//
//   - Megas: the hardware USART through simavr's uart_pty.
//   - Tinies: an 8N1 bridge between a pty and the GPIO software UART
//     (drives PB0, the loader's RX; decodes PB1, its TX), timed against the
//     simulated cycle counter.
//
// 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 use_uart_pty;
static const char *dump_path;
static uint32_t reset_pc;
static volatile sig_atomic_t reset_requested;

// 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.
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;
	}
	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;
	tx_shift = (uint8_t)((tx_shift >> 1) | (tx_level ? 0x80 : 0));
	if (++tx_bit < 8)
		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 (use_uart_pty)
		uart_pty_stop(&uart_pty);
	_exit(0);
}

int main(int argc, char *argv[])
{
	if (argc < 8 || argc > 10) {
		fprintf(stderr,
		        "usage: %s <pureboot.elf> <mcu> <hz> <base_hex> <page> <baud> <flash_dump>"
		        " [reset_hex] [resume_flash]\n"
		        "  reset_hex:    reset vector (default: base on the mega, 0 on the tinies)\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;
	}
	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];
	use_uart_pty = strncmp(mcu_name, "atmega", 6) == 0; // every mega links over its hardware USART

	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 (argc > 9) {
		// 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 mega enters the loader in hardware (BOOTRST, not modeled — the
	// argument picks the modeled fuse's target); the tinies reset to word 0
	// like silicon — erased flash walks up into the loader, and after the
	// host's surgery the patched vector routes there.
	reset_pc = argc > 8 ? (uint32_t)strtoul(argv[8], NULL, 0) : (use_uart_pty ? 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);
	}

	if (use_uart_pty) {
		fix_mega_flash_erase();
		// 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('0'), &flags);
		flags &= ~AVR_UART_FLAG_POLL_SLEEP;
		avr_ioctl(avr, AVR_IOCTL_UART_SET_FLAGS('0'), &flags);
		uart_pty_init(avr, &uart_pty);
		uart_pty_connect(&uart_pty, '0');
		printf("PB_PTY %s\n", uart_pty.pty.slavename);
	} 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);

		bit_cycles = (avr->frequency + baud / 2) / baud; // matches uart.hpp's own rounding exactly
		rx_pin = avr_io_getirq(avr, AVR_IOCTL_IOPORT_GETIRQ('B'), 0);
		avr_irq_register_notify(avr_io_getirq(avr, AVR_IOCTL_IOPORT_GETIRQ('B'), 1), 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 (use_uart_pty) { // reset restores the pacing hack; re-clear it
				uint32_t flags = 0;
				avr_ioctl(avr, AVR_IOCTL_UART_GET_FLAGS('0'), &flags);
				flags &= ~AVR_UART_FLAG_POLL_SLEEP;
				avr_ioctl(avr, AVR_IOCTL_UART_SET_FLAGS('0'), &flags);
			} else {
				bridge_reset();
			}
		}
		if (!use_uart_pty && ++since_poll >= 2000) {
			since_poll = 0;
			poll_pty();
		}
	}
	finish(0);
	return 0;
}