/*
 * 620_data_link — end-to-end w6l support for writable globals.
 * Assemble a small .s using DATAW, link it with w6l, run the binary,
 * and verify both that the read returns the global's value and that
 * a store updates it. Also confirms the binary has two PT_LOADs with
 * the right R+X / R+W flags.
 *
 * Inputs without DATAW must still produce the single-PT_LOAD layout
 * — 992 (selfhost w6l .o diff) depends on that invariant.
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <unistd.h>
#include <sys/wait.h>

#define PT_LOAD	1
#define PF_X	1
#define PF_W	2
#define PF_R	4

#pragma pack(push, 1)
typedef struct {
	uint8_t  e_ident[16];
	uint16_t e_type, e_machine;
	uint32_t e_version;
	uint64_t e_entry, e_phoff, e_shoff;
	uint32_t e_flags;
	uint16_t e_ehsize, e_phentsize, e_phnum, e_shentsize, e_shnum, e_shstrndx;
} Ehdr;

typedef struct {
	uint32_t p_type, p_flags;
	uint64_t p_offset, p_vaddr, p_paddr;
	uint64_t p_filesz, p_memsz, p_align;
} Phdr;
#pragma pack(pop)

static int
write_file(const char *path, const char *body)
{
	FILE *f = fopen(path, "wb");
	if (!f) return -1;
	fputs(body, f);
	fclose(f);
	return 0;
}

static int
slurp(const char *path, uint8_t **out, size_t *outn)
{
	FILE *f = fopen(path, "rb");
	if (!f) return -1;
	fseek(f, 0, SEEK_END);
	long n = ftell(f);
	fseek(f, 0, SEEK_SET);
	uint8_t *b = malloc((size_t)n);
	if (fread(b, 1, (size_t)n, f) != (size_t)n) {
		free(b);
		fclose(f);
		return -1;
	}
	fclose(f);
	*out = b;
	*outn = (size_t)n;
	return 0;
}

static int
build(const char *bin, const char *body, const char *exe)
{
	char src[64], obj[64], cmd[1024];
	snprintf(src, sizeof src, "/tmp/wwt_dl_%d.s", getpid());
	snprintf(obj, sizeof obj, "/tmp/wwt_dl_%d.o", getpid());
	if (write_file(src, body) < 0) return -1;

	snprintf(cmd, sizeof cmd, "%s/w6a -o %s %s", bin, obj, src);
	if (system(cmd) != 0) { unlink(src); return -1; }
	snprintf(cmd, sizeof cmd, "%s/w6l -o %s %s", bin, exe, obj);
	int rc = system(cmd);
	unlink(src); unlink(obj);
	return rc;
}

static int
run_exit(const char *exe)
{
	int rc = system(exe);
	if (rc == -1) return -1;
	if (WIFEXITED(rc)) return WEXITSTATUS(rc);
	return -1;
}

/* Test 1: read a writable global. _start exits with the value stored
 * in `counter` (0x2a == 42). */
static int
test_read(const char *bin)
{
	char exe[64];
	snprintf(exe, sizeof exe, "/tmp/wwt_dl_%d.x1", getpid());
	const char *body =
	    "TEXT _start,$0\n"
	    "\tMOVQ\tcounter(SB), DI\n"
	    "\tMOVQ\t$60, AX\n"
	    "\tSYSCALL\n"
	    "DATAW counter(SB),\"\\x2a\\x00\\x00\\x00\\x00\\x00\\x00\\x00\"\n";
	if (build(bin, body, exe) != 0) {
		fprintf(stderr, "test_read: build failed\n");
		return -1;
	}
	int rc = run_exit(exe);
	unlink(exe);
	if (rc != 42) {
		fprintf(stderr, "test_read: exit=%d, want 42\n", rc);
		return -1;
	}
	return 0;
}

/* Test 2: store to a writable global, read it back, exit with the new
 * value. Verifies the R+W mapping really is writable. */
static int
test_write(const char *bin)
{
	char exe[64];
	snprintf(exe, sizeof exe, "/tmp/wwt_dl_%d.x2", getpid());
	const char *body =
	    "TEXT _start,$0\n"
	    "\tMOVQ\t$99, AX\n"
	    "\tMOVQ\tAX, counter(SB)\n"
	    "\tMOVQ\tcounter(SB), DI\n"
	    "\tMOVQ\t$60, AX\n"
	    "\tSYSCALL\n"
	    "DATAW counter(SB),\"\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\"\n";
	if (build(bin, body, exe) != 0) {
		fprintf(stderr, "test_write: build failed\n");
		return -1;
	}
	int rc = run_exit(exe);
	unlink(exe);
	if (rc != 99) {
		fprintf(stderr, "test_write: exit=%d, want 99\n", rc);
		return -1;
	}
	return 0;
}

/* Test 3: structurally inspect the linked ELF — two PT_LOADs, the
 * first R+X and the second R+W. Catches regressions where the layout
 * is right but the flags or count are off. */
static int
test_two_loads(const char *bin)
{
	char exe[64];
	snprintf(exe, sizeof exe, "/tmp/wwt_dl_%d.x3", getpid());
	const char *body =
	    "TEXT _start,$0\n"
	    "\tMOVQ\tcounter(SB), AX\n"
	    "\tMOVQ\t$60, AX\n"
	    "\tMOVQ\t$0, DI\n"
	    "\tSYSCALL\n"
	    "DATAW counter(SB),\"\\x01\\x00\\x00\\x00\\x00\\x00\\x00\\x00\"\n";
	if (build(bin, body, exe) != 0) {
		fprintf(stderr, "test_two_loads: build failed\n");
		return -1;
	}
	uint8_t *buf = NULL;
	size_t n = 0;
	if (slurp(exe, &buf, &n) < 0) {
		fprintf(stderr, "test_two_loads: cannot read exe\n");
		unlink(exe);
		return -1;
	}
	unlink(exe);

	const Ehdr *eh = (const Ehdr *)buf;
	int rc = 0;
	if (eh->e_phnum != 2) {
		fprintf(stderr, "test_two_loads: e_phnum=%u, want 2\n", eh->e_phnum);
		rc = -1;
		goto out;
	}
	const Phdr *ph = (const Phdr *)(buf + eh->e_phoff);
	int saw_rx = 0, saw_rw = 0;
	for (int i = 0; i < 2; i++) {
		if (ph[i].p_type != PT_LOAD) {
			fprintf(stderr, "test_two_loads: ph[%d] not PT_LOAD\n", i);
			rc = -1;
			continue;
		}
		if (ph[i].p_flags == (PF_R | PF_X)) saw_rx = 1;
		else if (ph[i].p_flags == (PF_R | PF_W)) saw_rw = 1;
		else {
			fprintf(stderr, "test_two_loads: ph[%d] flags=0x%x\n",
			    i, ph[i].p_flags);
			rc = -1;
		}
	}
	if (!saw_rx) { fprintf(stderr, "test_two_loads: missing R+X\n"); rc = -1; }
	if (!saw_rw) { fprintf(stderr, "test_two_loads: missing R+W\n"); rc = -1; }
out:
	free(buf);
	return rc;
}

/* Test 4: an input without DATAW still produces a single PT_LOAD.
 * 992 (selfhost w6l .o diff) depends on this byte-level invariant. */
static int
test_no_data_single_load(const char *bin)
{
	char exe[64];
	snprintf(exe, sizeof exe, "/tmp/wwt_dl_%d.x4", getpid());
	const char *body =
	    "TEXT _start,$0\n"
	    "\tMOVQ\t$60, AX\n"
	    "\tMOVQ\t$0, DI\n"
	    "\tSYSCALL\n";
	if (build(bin, body, exe) != 0) {
		fprintf(stderr, "test_no_data_single_load: build failed\n");
		return -1;
	}
	uint8_t *buf = NULL;
	size_t n = 0;
	if (slurp(exe, &buf, &n) < 0) {
		fprintf(stderr, "test_no_data_single_load: cannot read exe\n");
		unlink(exe);
		return -1;
	}
	unlink(exe);
	int rc = 0;
	const Ehdr *eh = (const Ehdr *)buf;
	if (eh->e_phnum != 1) {
		fprintf(stderr, "test_no_data_single_load: e_phnum=%u, want 1\n",
		    eh->e_phnum);
		rc = -1;
	}
	free(buf);
	return rc;
}

/* Test 5: BSS optimisation — trailing zero bytes in .data are dropped
 * from the file (p_filesz < p_memsz) and the loader zero-fills the gap.
 * Asserts that a trailing zero-init slot doesn't bloat the binary. */
static int
test_bss_filesz(const char *bin)
{
	char exe[64];
	snprintf(exe, sizeof exe, "/tmp/wwt_dl_%d.x5", getpid());
	/* Two DATAW slots: a 1-byte non-zero followed by a 31-byte zero
	 * run. Memsz must cover all 32B; filesz should stop after the 1B
	 * non-zero (so file_data == 1, memsz_data == 32). */
	const char *body =
	    "TEXT _start,$0\n"
	    "\tMOVQ\tnz(SB), DI\n"
	    "\tMOVQ\t$60, AX\n"
	    "\tSYSCALL\n"
	    "DATAW nz(SB),\"\\x07\\x00\\x00\\x00\\x00\\x00\\x00\\x00\"\n"
	    "DATAW zlong(SB),\"\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00"
	                    "\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00"
	                    "\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00\"\n";
	if (build(bin, body, exe) != 0) {
		fprintf(stderr, "test_bss_filesz: build failed\n");
		return -1;
	}
	uint8_t *buf = NULL;
	size_t n = 0;
	if (slurp(exe, &buf, &n) < 0) {
		fprintf(stderr, "test_bss_filesz: cannot read exe\n");
		unlink(exe);
		return -1;
	}
	int rc = 0;
	const Ehdr *eh = (const Ehdr *)buf;
	const Phdr *ph = (const Phdr *)(buf + eh->e_phoff);
	int saw_rw = 0;
	for (int i = 0; i < eh->e_phnum; i++) {
		if (ph[i].p_type != PT_LOAD) continue;
		if (ph[i].p_flags != (PF_R | PF_W)) continue;
		saw_rw = 1;
		if (ph[i].p_filesz >= ph[i].p_memsz) {
			fprintf(stderr, "test_bss_filesz: filesz=%lu, "
			    "memsz=%lu (want filesz<memsz)\n",
			    (unsigned long)ph[i].p_filesz,
			    (unsigned long)ph[i].p_memsz);
			rc = -1;
		}
		/* The scan trims every trailing zero — it doesn't know
		 * symbol boundaries. nz is \x07 then 7 zero bytes, so
		 * the file ends at byte 1 (the 0x07). Memsz still covers
		 * all 32B and the loader zero-fills. */
		if (ph[i].p_filesz != 1) {
			fprintf(stderr, "test_bss_filesz: filesz=%lu, "
			    "want 1\n", (unsigned long)ph[i].p_filesz);
			rc = -1;
		}
		if (ph[i].p_memsz != 32) {
			fprintf(stderr, "test_bss_filesz: memsz=%lu, "
			    "want 32\n", (unsigned long)ph[i].p_memsz);
			rc = -1;
		}
	}
	if (!saw_rw) {
		fprintf(stderr, "test_bss_filesz: no R+W PT_LOAD found\n");
		rc = -1;
	}
	free(buf);
	if (rc != 0) { unlink(exe); return rc; }
	/* Behaviour check: the zero-init slot must still read as zero
	 * at runtime — the loader has to zero-fill the BSS gap. */
	int er = run_exit(exe);
	unlink(exe);
	if (er != 7) {
		fprintf(stderr, "test_bss_filesz: exit=%d, want 7\n", er);
		return -1;
	}
	return 0;
}

int
main(void)
{
	const char *bin = getenv("BIN");
	if (!bin) bin = "out/bin";
	int fail = 0;
	if (test_read(bin)                != 0) fail++;
	if (test_write(bin)               != 0) fail++;
	if (test_two_loads(bin)           != 0) fail++;
	if (test_no_data_single_load(bin) != 0) fail++;
	if (test_bss_filesz(bin)          != 0) fail++;
	if (fail) {
		fprintf(stderr, "data_link: %d/5 subtests failed\n", fail);
		return 1;
	}
	printf("data_link: 5/5 ok\n");
	return 0;
}