/*
 * 510_dataw — DATAW directive smoke. Assemble a tiny .s that uses
 * DATAW and verify the resulting .o has a writable .data section
 * containing the right bytes, with a symbol that points to it.
 *
 * Also confirms the no-DATAW path stays byte-identical to the
 * pre-DATAW format (6 sections, no .data); the selfhost .o diff in
 * 991 depends on that invariant.
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <unistd.h>

#define SHT_PROGBITS	1
#define SHT_SYMTAB	2
#define SHF_WRITE	0x1
#define SHF_ALLOC	0x2
#define SHF_EXECINSTR	0x4

#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 sh_name, sh_type;
	uint64_t sh_flags, sh_addr, sh_offset, sh_size;
	uint32_t sh_link, sh_info;
	uint64_t sh_addralign, sh_entsize;
} Shdr;

typedef struct {
	uint32_t st_name;
	uint8_t  st_info, st_other;
	uint16_t st_shndx;
	uint64_t st_value, st_size;
} Sym64;
#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;
}

/* Find the section with the given name. Returns index, or -1. */
static int
find_section(const uint8_t *buf, const char *name)
{
	const Ehdr *eh = (const Ehdr *)buf;
	const Shdr *sh = (const Shdr *)(buf + eh->e_shoff);
	const char *shstr = (const char *)(buf + sh[eh->e_shstrndx].sh_offset);
	for (uint16_t i = 0; i < eh->e_shnum; i++) {
		if (strcmp(shstr + sh[i].sh_name, name) == 0)
			return (int)i;
	}
	return -1;
}

/* Find a symbol by name in the .symtab. Returns pointer, or NULL. */
static const Sym64 *
find_symbol(const uint8_t *buf, const char *name)
{
	const Ehdr *eh = (const Ehdr *)buf;
	const Shdr *sh = (const Shdr *)(buf + eh->e_shoff);
	int idx_symtab = -1;
	for (uint16_t i = 0; i < eh->e_shnum; i++)
		if (sh[i].sh_type == SHT_SYMTAB) { idx_symtab = i; break; }
	if (idx_symtab < 0) return NULL;
	int idx_strtab = (int)sh[idx_symtab].sh_link;
	const Sym64 *syms = (const Sym64 *)(buf + sh[idx_symtab].sh_offset);
	uint64_t nsyms = sh[idx_symtab].sh_size / sizeof(Sym64);
	const char *str = (const char *)(buf + sh[idx_strtab].sh_offset);
	for (uint64_t i = 1; i < nsyms; i++) {
		if (strcmp(str + syms[i].st_name, name) == 0)
			return &syms[i];
	}
	return NULL;
}

static int
test_dataw_emits_writable_section(const char *bin)
{
	char src[64], obj[64];
	snprintf(src, sizeof src, "/tmp/wwt_dataw_%d.s", getpid());
	snprintf(obj, sizeof obj, "/tmp/wwt_dataw_%d.o", getpid());

	/* Hand-crafted .s: one writable data symbol holding 8 bytes (the
	 * little-endian encoding of the i64 42) plus a trivial TEXT so
	 * the assembler has something in .text to anchor. */
	const char *body =
	    "TEXT _start(SB),$0\n"
	    "\tMOVQ\t$0, AX\n"
	    "\tRET\n"
	    "DATAW counter(SB),\"\\x2a\\x00\\x00\\x00\\x00\\x00\\x00\\x00\"\n";
	if (write_file(src, body) < 0) {
		fprintf(stderr, "dataw: cannot write %s\n", src);
		return -1;
	}

	char cmd[512];
	snprintf(cmd, sizeof cmd, "%s/w6a -o %s %s", bin, obj, src);
	if (system(cmd) != 0) {
		fprintf(stderr, "dataw: w6a failed\n");
		unlink(src);
		return -1;
	}

	uint8_t *buf = NULL;
	size_t n = 0;
	if (slurp(obj, &buf, &n) < 0) {
		fprintf(stderr, "dataw: cannot read .o\n");
		unlink(src); unlink(obj);
		return -1;
	}

	const Ehdr *eh = (const Ehdr *)buf;
	int rc = 0;

	/* Expect 7 sections: NULL, .text, .rela.text, .data, .symtab, .strtab, .shstrtab. */
	if (eh->e_shnum != 7) {
		fprintf(stderr, "dataw: e_shnum=%u, want 7\n", eh->e_shnum);
		rc = -1;
	}

	int idx_data = find_section(buf, ".data");
	if (idx_data < 0) {
		fprintf(stderr, "dataw: .data section not found\n");
		rc = -1;
		goto out;
	}
	const Shdr *sh = (const Shdr *)(buf + eh->e_shoff);
	const Shdr *sd = &sh[idx_data];
	if (sd->sh_type != SHT_PROGBITS) {
		fprintf(stderr, "dataw: .data sh_type=%u, want PROGBITS\n", sd->sh_type);
		rc = -1;
	}
	if ((sd->sh_flags & (SHF_ALLOC | SHF_WRITE)) != (SHF_ALLOC | SHF_WRITE)) {
		fprintf(stderr, "dataw: .data flags=0x%lx, want ALLOC|WRITE\n",
		    (unsigned long)sd->sh_flags);
		rc = -1;
	}
	if ((sd->sh_flags & SHF_EXECINSTR) != 0) {
		fprintf(stderr, "dataw: .data is executable, want non-X\n");
		rc = -1;
	}
	if (sd->sh_size != 8) {
		fprintf(stderr, "dataw: .data size=%lu, want 8\n",
		    (unsigned long)sd->sh_size);
		rc = -1;
	}

	/* Verify the bytes themselves match the LE i64 42 we wrote. */
	const uint8_t *db = buf + sd->sh_offset;
	uint8_t want[8] = { 0x2a, 0, 0, 0, 0, 0, 0, 0 };
	if (memcmp(db, want, 8) != 0) {
		fprintf(stderr, "dataw: .data contents mismatch\n");
		rc = -1;
	}

	const Sym64 *s = find_symbol(buf, "counter");
	if (s == NULL) {
		fprintf(stderr, "dataw: symbol 'counter' not in .symtab\n");
		rc = -1;
		goto out;
	}
	if ((int)s->st_shndx != idx_data) {
		fprintf(stderr, "dataw: 'counter'.st_shndx=%u, want %d (.data)\n",
		    s->st_shndx, idx_data);
		rc = -1;
	}
	if (s->st_value != 0) {
		fprintf(stderr, "dataw: 'counter'.st_value=%lu, want 0\n",
		    (unsigned long)s->st_value);
		rc = -1;
	}

out:
	free(buf);
	unlink(src);
	unlink(obj);
	return rc;
}

static int
test_no_dataw_keeps_legacy_layout(const char *bin)
{
	char src[64], obj[64];
	snprintf(src, sizeof src, "/tmp/wwt_nodataw_%d.s", getpid());
	snprintf(obj, sizeof obj, "/tmp/wwt_nodataw_%d.o", getpid());

	/* A .s without any DATAW must still produce the legacy 6-section
	 * layout. Test 991 (selfhost .o byte-diff) relies on this. */
	const char *body =
	    "TEXT _start(SB),$0\n"
	    "\tMOVQ\t$0, AX\n"
	    "\tRET\n";
	if (write_file(src, body) < 0) {
		fprintf(stderr, "nodataw: cannot write %s\n", src);
		return -1;
	}

	char cmd[512];
	snprintf(cmd, sizeof cmd, "%s/w6a -o %s %s", bin, obj, src);
	if (system(cmd) != 0) {
		fprintf(stderr, "nodataw: w6a failed\n");
		unlink(src);
		return -1;
	}

	uint8_t *buf = NULL;
	size_t n = 0;
	if (slurp(obj, &buf, &n) < 0) {
		fprintf(stderr, "nodataw: cannot read .o\n");
		unlink(src); unlink(obj);
		return -1;
	}

	const Ehdr *eh = (const Ehdr *)buf;
	int rc = 0;
	if (eh->e_shnum != 6) {
		fprintf(stderr, "nodataw: e_shnum=%u, want 6 (legacy)\n",
		    eh->e_shnum);
		rc = -1;
	}
	if (find_section(buf, ".data") >= 0) {
		fprintf(stderr, "nodataw: .data section present, want absent\n");
		rc = -1;
	}

	free(buf);
	unlink(src);
	unlink(obj);
	return rc;
}

int
main(void)
{
	const char *bin = getenv("BIN");
	if (!bin) bin = "out/bin";

	int fail = 0;
	if (test_dataw_emits_writable_section(bin) != 0) fail++;
	if (test_no_dataw_keeps_legacy_layout(bin)   != 0) fail++;

	if (fail) {
		fprintf(stderr, "dataw: %d/2 subtests failed\n", fail);
		return 1;
	}
	printf("dataw: 2/2 ok\n");
	return 0;
}