/*
 * obj.c — load an ELF64 relocatable object emitted by w6a, append its
 * .text bytes to the combined image, and pull its symbols and
 * relocations into the global tables (with offsets adjusted to the
 * combined section).
 */
#include "l.h"
#include <stdlib.h>
#include <string.h>

#define ET_REL		1
#define EM_X86_64	62

#define SHT_PROGBITS	1
#define SHT_SYMTAB	2
#define SHT_STRTAB	3
#define SHT_RELA	4

#pragma pack(push, 1)
typedef struct {
	u8 e_ident[16];
	u16 e_type, e_machine;
	u32 e_version;
	u64 e_entry, e_phoff, e_shoff;
	u32 e_flags;
	u16 e_ehsize, e_phentsize, e_phnum, e_shentsize, e_shnum, e_shstrndx;
} Ehdr;

typedef struct {
	u32 sh_name, sh_type;
	u64 sh_flags, sh_addr, sh_offset, sh_size;
	u32 sh_link, sh_info;
	u64 sh_addralign, sh_entsize;
} Shdr;

typedef struct {
	u32 st_name;
	u8 st_info, st_other;
	u16 st_shndx;
	u64 st_value, st_size;
} Sym64;

typedef struct {
	u64 r_offset;
	u64 r_info;
	i64 r_addend;
} Rela64;
#pragma pack(pop)

#define ELF64_R_SYM(i)	((u32)((i) >> 32))
#define ELF64_R_TYPE(i)	((u32)((i) & 0xffffffff))
#define ELF64_ST_TYPE(i)	((i) & 0xf)
#define ELF64_ST_BIND(i)	((i) >> 4)

int
l_read_all(const char *path, u8 **out, u64 *len)
{
	FILE *f = fopen(path, "rb");
	if (f == NULL) return -1;
	fseek(f, 0, SEEK_END);
	long n = ftell(f);
	fseek(f, 0, SEEK_SET);
	if (n < 0) { fclose(f); return -1; }
	u8 *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;
	*len = (u64)n;
	return 0;
}

#define read_all l_read_all

static void
emit_text(Lnk *l, const u8 *src, u64 n)
{
	if (l->textlen + n > l->textcap) {
		u64 nc = l->textcap ? l->textcap * 2 : 4096;
		while (nc < l->textlen + n) nc *= 2;
		l->text = realloc(l->text, nc);
		l->textcap = nc;
	}
	memcpy(l->text + l->textlen, src, n);
	l->textlen += n;
}

/* Internal: load a single ELF .o image already in memory. The caller
 * gives us the bytes (we own them) and a path tag for diagnostics.
 * If the bytes look like an archive (magic "!<arch>\n") we recurse
 * over each member instead.
 */
static int load_image(Lnk *l, const char *path, u8 *buf, u64 len);

static u64
ar_field(const u8 *p, int n)
{
	/* decimal field, space-padded */
	u64 v = 0;
	for (int i = 0; i < n; i++) {
		if (p[i] >= '0' && p[i] <= '9') v = v * 10 + (p[i] - '0');
		else if (p[i] == ' ') break;
		else if (p[i] == 0) break;
	}
	return v;
}

/* Read an ELF .o image's globally-defined symbol names without
 * actually appending it to the link. Returns a heap-allocated
 * NULL-terminated array; caller frees the array (not the strings,
 * which point into the .o image and must remain alive).
 */
static char **
elf_globals(const u8 *buf, u64 len)
{
	if (len < sizeof(Ehdr)) return NULL;
	Ehdr *eh = (Ehdr *)buf;
	if (memcmp(eh->e_ident, "\x7f""ELF", 4) != 0) return NULL;
	Shdr *sh = (Shdr *)(buf + eh->e_shoff);
	int idx_text = -1, idx_symtab = -1;
	const char *shstr = (const char *)(buf + sh[eh->e_shstrndx].sh_offset);
	for (u16 i = 0; i < eh->e_shnum; i++) {
		if (sh[i].sh_type == SHT_PROGBITS &&
		    strcmp(shstr + sh[i].sh_name, ".text") == 0)
			idx_text = i;
		else if (sh[i].sh_type == SHT_SYMTAB)
			idx_symtab = i;
	}
	if (idx_text < 0 || idx_symtab < 0) return NULL;
	int idx_strtab = sh[idx_symtab].sh_link;
	Sym64 *symtab = (Sym64 *)(buf + sh[idx_symtab].sh_offset);
	u64 nsyms = sh[idx_symtab].sh_size / sizeof(Sym64);
	const char *str = (const char *)(buf + sh[idx_strtab].sh_offset);

	char **out = calloc(nsyms + 1, sizeof *out);
	int n = 0;
	for (u64 i = 1; i < nsyms; i++) {
		if (symtab[i].st_shndx == 0) continue;
		if ((symtab[i].st_info >> 4) != 1) continue; /* STB_GLOBAL */
		if ((int)symtab[i].st_shndx != idx_text) continue;
		out[n++] = strdup(str + symtab[i].st_name);
	}
	out[n] = NULL;
	return out;
}

typedef struct ArMember ArMember;
struct ArMember {
	u8	*data;	/* heap copy; freed if never loaded */
	u64	size;
	char	**defs;	/* NULL-terminated list of defined globals */
	int	loaded;
	ArMember *next;
};

static int
member_defines_undef(Lnk *l, ArMember *m)
{
	if (m->defs == NULL) return 0;
	for (int i = 0; m->defs[i]; i++) {
		Lsym *s = l_lookup(l, m->defs[i]);
		if (s != NULL && !s->defined) return 1;
	}
	return 0;
}

static int
load_archive(Lnk *l, const char *path, u8 *buf, u64 len)
{
	/* Pass 1: index members. We copy each member's bytes (cheap; few
	 * tens of KB per stdlib module) so the archive buffer can be
	 * freed once we're done indexing. */
	ArMember *head = NULL, *tail = NULL;
	u64 pos = 8;	/* past "!<arch>\n" */
	while (pos + 60 <= len) {
		const u8 *hdr = buf + pos;
		u64 size = ar_field(hdr + 48, 10);
		u64 hdr_end = pos + 60;
		if (hdr_end + size > len) break;
		if (hdr[0] != '/' && hdr[0] != 0 && hdr[0] != ' ') {
			ArMember *m = calloc(1, sizeof *m);
			m->size = size;
			m->data = malloc((size_t)size);
			memcpy(m->data, buf + hdr_end, (size_t)size);
			m->defs = elf_globals(m->data, size);
			if (head == NULL) head = m;
			else tail->next = m;
			tail = m;
		}
		pos = hdr_end + size;
		if (size & 1) pos++;
	}
	free(buf);

	/* Pass 2: iteratively pull members that define a currently-
	 * undefined symbol. Each pull may introduce new undefs, so loop. */
	int changed = 1;
	while (changed) {
		changed = 0;
		for (ArMember *m = head; m; m = m->next) {
			if (m->loaded) continue;
			if (!member_defines_undef(l, m)) continue;
			u8 *copy = malloc((size_t)m->size);
			memcpy(copy, m->data, m->size);
			if (load_image(l, path, copy, m->size) == 0) {
				m->loaded = 1;
				changed = 1;
			}
		}
	}

	/* Free unloaded members; loaded ones had their bytes consumed
	 * by load_image (which took the copy). */
	while (head) {
		ArMember *next = head->next;
		free(head->data);
		if (head->defs) {
			for (int i = 0; head->defs[i]; i++) free(head->defs[i]);
			free(head->defs);
		}
		free(head);
		head = next;
	}
	return 0;
}

int
l_load(Lnk *l, const char *path)
{
	u8 *buf;
	u64 len;
	if (read_all(path, &buf, &len) < 0) return -1;
	if (len >= 8 && memcmp(buf, "!<arch>\n", 8) == 0)
		return load_archive(l, path, buf, len);
	/* Shared object: dispatch to dyn.c, which re-reads (small loss
	 * for a much cleaner separation of static vs dynamic loaders). */
	if (len >= sizeof(Ehdr)) {
		Ehdr *eh = (Ehdr *)buf;
		if (memcmp(eh->e_ident, "\x7f""ELF", 4) == 0
		    && eh->e_type == 3 /* ET_DYN */) {
			free(buf);
			return l_load_so(l, path);
		}
	}
	return load_image(l, path, buf, len);
}

static int
load_image(Lnk *l, const char *path, u8 *buf, u64 len)
{
	if (len < sizeof(Ehdr)) { free(buf); return -1; }
	Ehdr *eh = (Ehdr *)buf;
	if (memcmp(eh->e_ident, "\x7f""ELF", 4) != 0 || eh->e_ident[4] != 2
	    || eh->e_machine != EM_X86_64 || eh->e_type != ET_REL) {
		fprintf(stderr, "w6l: %s: not an amd64 ELF64 relocatable\n", path);
		free(buf);
		return -1;
	}
	Shdr *sh = (Shdr *)(buf + eh->e_shoff);
	if (eh->e_shstrndx >= eh->e_shnum) { free(buf); return -1; }
	const char *shstr = (const char *)(buf + sh[eh->e_shstrndx].sh_offset);

	/* find .text, .symtab, .strtab, .rela.text */
	int idx_text = -1, idx_symtab = -1, idx_strtab = -1, idx_rela = -1;
	for (u16 i = 0; i < eh->e_shnum; i++) {
		const char *nm = shstr + sh[i].sh_name;
		if (sh[i].sh_type == SHT_PROGBITS && strcmp(nm, ".text") == 0)
			idx_text = i;
		else if (sh[i].sh_type == SHT_SYMTAB)
			idx_symtab = i;
		else if (sh[i].sh_type == SHT_RELA && strcmp(nm, ".rela.text") == 0)
			idx_rela = i;
	}
	if (idx_text < 0 || idx_symtab < 0) {
		fprintf(stderr, "w6l: %s: missing .text or .symtab\n", path);
		free(buf);
		return -1;
	}
	idx_strtab = sh[idx_symtab].sh_link;

	Lobj *ob = calloc(1, sizeof *ob);
	ob->path = strdup(path);
	ob->buf = buf;
	ob->len = len;
	ob->text_off = l->textlen;
	ob->text_size = sh[idx_text].sh_size;
	ob->next = l->objs;
	l->objs = ob;

	/* append .text */
	emit_text(l, buf + sh[idx_text].sh_offset, sh[idx_text].sh_size);

	/* per-object: load symbols */
	Sym64 *symtab = (Sym64 *)(buf + sh[idx_symtab].sh_offset);
	u64 nsyms = sh[idx_symtab].sh_size / sizeof(Sym64);
	const char *str = (const char *)(buf + sh[idx_strtab].sh_offset);

	/* map per-object sym index → global Lsym */
	Lsym **map = calloc(nsyms, sizeof *map);
	for (u64 i = 1; i < nsyms; i++) {
		const char *nm = str + symtab[i].st_name;
		if (nm[0] == '\0') continue;
		Lsym *gs = l_intern(l, nm);
		if (symtab[i].st_shndx != 0 /* SHN_UNDEF */
		    && symtab[i].st_shndx == idx_text) {
			if (gs->defined) {
				fprintf(stderr, "w6l: %s: duplicate symbol %s\n",
				    path, nm);
				l->errs++;
			} else {
				gs->defined = 1;
				gs->owner = ob;
				gs->idx_in_owner = (int)i;
				gs->val = ob->text_off + symtab[i].st_value;
			}
		}
		map[i] = gs;
	}

	/* per-object: collect relocations */
	if (idx_rela >= 0) {
		Rela64 *rt = (Rela64 *)(buf + sh[idx_rela].sh_offset);
		u64 nrel = sh[idx_rela].sh_size / sizeof(Rela64);
		for (u64 i = 0; i < nrel; i++) {
			Lrel *r = calloc(1, sizeof *r);
			r->off = ob->text_off + rt[i].r_offset;
			r->kind = (int)ELF64_R_TYPE(rt[i].r_info);
			u32 sidx = ELF64_R_SYM(rt[i].r_info);
			r->sym = (sidx < nsyms) ? map[sidx] : NULL;
			r->addend = rt[i].r_addend;
			r->next = l->rels;
			l->rels = r;
		}
	}
	free(map);
	return 0;
}