/*
 * 520_datar — DATAR reloc smoke. Assemble + link a tiny .s that
 * uses DATAR to patch a .data slot with a symbol's runtime VA,
 * then run the binary and confirm the linker filled the slot
 * correctly: the program prints "hello" via the relocated pointer
 * and exits 0.
 *
 * Also confirms the .o has a `.rela.data` section with one
 * R_X86_64_64 entry against the expected symbol — catches
 * regressions where the asm emits DATAR but obj.c doesn't reach
 * the writable-data reloc path.
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <unistd.h>
#include <sys/wait.h>

#define SHT_RELA	4
#define R_X86_64_64	1

#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 {
	uint64_t r_offset;
	uint64_t r_info;
	int64_t  r_addend;
} Rela64;
#pragma pack(pop)

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

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 const char *prog =
    "TEXT _start,$0\n"
    "\tMOVQ\t$1, AX\n"			/* write syscall */
    "\tMOVQ\t$1, DI\n"			/* fd = 1 (stdout) */
    "\tMOVQ\tpair(SB), SI\n"		/* ptr (relocated) */
    "\tMOVQ\t$5, DX\n"			/* count = 5 */
    "\tSYSCALL\n"
    "\tMOVQ\t$60, AX\n"			/* exit */
    "\tMOVQ\t$0, DI\n"
    "\tSYSCALL\n"
    "\n"
    "DATAW pair(SB),\"\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00"
    "\\x05\\x00\\x00\\x00\\x00\\x00\\x00\\x00\"\n"
    "DATAR pair+0(SB),greeting(SB)\n"
    "DATA greeting(SB),\"hello\"\n";

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

	char src[64], obj[64], exe[64], cmd[1024];
	snprintf(src, sizeof src, "/tmp/wwt_datar_%d.s", getpid());
	snprintf(obj, sizeof obj, "/tmp/wwt_datar_%d.o", getpid());
	snprintf(exe, sizeof exe, "/tmp/wwt_datar_%d.x", getpid());

	FILE *f = fopen(src, "wb");
	if (!f) return 1;
	fputs(prog, f);
	fclose(f);

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

	int fail = 0;

	/* Inspect the .o: there must be a .rela.data section with one
	 * R_X86_64_64 entry against `greeting`. */
	uint8_t *buf = NULL; size_t n = 0;
	if (slurp(obj, &buf, &n) < 0) {
		fprintf(stderr, "datar: cannot read .o\n");
		fail++;
		goto run_step;
	}
	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);
	int idx_relad = -1, idx_strtab = -1, idx_symtab = -1;
	for (uint16_t i = 0; i < eh->e_shnum; i++) {
		const char *nm = shstr + sh[i].sh_name;
		if (sh[i].sh_type == SHT_RELA && strcmp(nm, ".rela.data") == 0)
			idx_relad = i;
		if (sh[i].sh_type == 2 /* SHT_SYMTAB */)
			idx_symtab = i;
	}
	if (idx_symtab >= 0)
		idx_strtab = (int)sh[idx_symtab].sh_link;
	if (idx_relad < 0) {
		fprintf(stderr, "datar: .rela.data section missing\n");
		fail++;
	} else {
		const Rela64 *rd = (const Rela64 *)(buf + sh[idx_relad].sh_offset);
		uint64_t nrel = sh[idx_relad].sh_size / sizeof(Rela64);
		if (nrel != 1) {
			fprintf(stderr, "datar: .rela.data has %lu entries, want 1\n",
			    (unsigned long)nrel);
			fail++;
		} else {
			uint32_t kind = (uint32_t)(rd[0].r_info & 0xffffffff);
			if (kind != R_X86_64_64) {
				fprintf(stderr, "datar: reloc kind %u, want %d\n",
				    kind, R_X86_64_64);
				fail++;
			}
			if (rd[0].r_offset != 0) {
				fprintf(stderr, "datar: reloc offset %lu, want 0\n",
				    (unsigned long)rd[0].r_offset);
				fail++;
			}
		}
	}
	free(buf);

run_step:
	snprintf(cmd, sizeof cmd, "%s/w6l -o %s %s", bin, exe, obj);
	if (system(cmd) != 0) {
		fprintf(stderr, "datar: w6l failed\n");
		unlink(src); unlink(obj);
		return 1;
	}

	/* Run: should print "hello" and exit 0. We pipe stdout to a
	 * tmp file and check both the output and the exit code. */
	char outf[64];
	snprintf(outf, sizeof outf, "/tmp/wwt_datar_%d.out", getpid());
	snprintf(cmd, sizeof cmd, "%s > %s", exe, outf);
	int rc = run_exit(cmd);
	if (rc != 0) {
		fprintf(stderr, "datar: exit=%d, want 0\n", rc);
		fail++;
	}
	uint8_t *out = NULL; size_t on = 0;
	if (slurp(outf, &out, &on) < 0 || on != 5
	    || memcmp(out, "hello", 5) != 0) {
		fprintf(stderr, "datar: stdout != 'hello'\n");
		fail++;
	}
	free(out);

	unlink(src); unlink(obj); unlink(exe); unlink(outf);
	if (fail) {
		fprintf(stderr, "datar: %d check(s) failed\n", fail);
		return 1;
	}
	printf("datar: ok\n");
	(void)idx_strtab;
	return 0;
}