/*
 * obj.c — emit a tiny ELF64 relocatable object.
 *
 * Layout (in file order):
 *    [0] ELF header
 *    [1] Section .text (program bytes)
 *    [2] Section .rela.text (relocations)
 *    [3] Section .symtab
 *    [4] Section .strtab
 *    [5] Section .shstrtab
 *    [6] Section header table
 *
 * Symtab indices: 0 = STN_UNDEF, 1 = file (skipped), 2.. = our syms.
 * For simplicity we emit GLOBAL symbols only (no LOCAL ordering rules
 * to worry about).
 */
#include "a.h"
#include <stdlib.h>
#include <string.h>
#include <stdio.h>

/* ELF constants */
#define ELFMAG		"\x7f""ELF"
#define ELFCLASS64	2
#define ELFDATA2LSB	1
#define EV_CURRENT	1
#define ET_REL		1
#define EM_X86_64	62

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

#define SHF_ALLOC	0x2
#define SHF_EXECINSTR	0x4
#define SHF_INFO_LINK	0x40

#define STB_LOCAL	0
#define STB_GLOBAL	1
#define STT_NOTYPE	0
#define STT_FUNC	2
#define ELF64_ST_INFO(b,t)	(((b) << 4) + ((t) & 0xf))

#define R_X86_64_PC32	2
#define R_X86_64_PLT32	4
#define ELF64_R_INFO(s,t)	(((u64)(s) << 32) | ((u64)(t) & 0xffffffff))

/* growable byte buffer */
typedef struct Buf Buf;
struct Buf { u8 *p; size_t n, cap; };

static void
bput(Buf *b, const void *src, size_t n)
{
	if (b->n + n > b->cap) {
		size_t nc = b->cap ? b->cap * 2 : 256;
		while (nc < b->n + n) nc *= 2;
		b->p = realloc(b->p, nc);
		b->cap = nc;
	}
	memcpy(b->p + b->n, src, n);
	b->n += n;
}

static u32 stput(Buf *st, const char *s) {
	u32 off = (u32)st->n;
	bput(st, s, strlen(s) + 1);
	return off;
}

#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)

int
a_emit_elf(Asm *a, FILE *f)
{
	Buf shstr = {0}, str = {0}, sym = {0}, rela = {0};
	stput(&shstr, "");	/* idx 0 = empty */
	stput(&str, "");

	/* Section name offsets */
	u32 shn_text	  = stput(&shstr, ".text");
	u32 shn_rela	  = stput(&shstr, ".rela.text");
	u32 shn_symtab	  = stput(&shstr, ".symtab");
	u32 shn_strtab	  = stput(&shstr, ".strtab");
	u32 shn_shstrtab  = stput(&shstr, ".shstrtab");

	/* Symbol 0 — STN_UNDEF */
	{
		Sym64 z = {0};
		bput(&sym, &z, sizeof z);
	}

	/* Section indices: 1=.text, 2=.rela.text, 3=.symtab, 4=.strtab, 5=.shstrtab */
	const u16 SH_TEXT = 1;

	/* Build symbols (defined = global; undefined = global UND) */
	int idx = 1;
	for (Asym *s = a->syms; s; s = s->next) {
		Sym64 e = {0};
		e.st_name = stput(&str, s->name);
		if (s->defined) {
			e.st_info = ELF64_ST_INFO(STB_GLOBAL, STT_FUNC);
			e.st_shndx = SH_TEXT;
			e.st_value = s->addr;
			e.st_size = 0;
		} else {
			e.st_info = ELF64_ST_INFO(STB_GLOBAL, STT_NOTYPE);
			e.st_shndx = 0;
		}
		bput(&sym, &e, sizeof e);
		s->idx = idx++;
	}

	/* Build relocations */
	for (Areloc *r = a->relocs; r; r = r->next) {
		Rela64 re;
		re.r_offset = r->off;
		re.r_info = ELF64_R_INFO((u64)r->sym->idx, (u64)r->kind);
		re.r_addend = r->addend;
		bput(&rela, &re, sizeof re);
	}

	/* Layout offsets in the file */
	u64 off = sizeof(Ehdr);
	u64 off_text = off; off += a->textlen;
	u64 off_rela = off; off += rela.n;
	u64 off_sym  = off; off += sym.n;
	u64 off_str  = off; off += str.n;
	u64 off_shstr= off; off += shstr.n;
	/* align to 8 */
	while (off % 8) off++;
	u64 off_shdr = off;
	const int NSECT = 6;	/* null + 5 real */

	Ehdr eh = {0};
	memcpy(eh.e_ident, "\x7f""ELF", 4);
	eh.e_ident[4] = ELFCLASS64;
	eh.e_ident[5] = ELFDATA2LSB;
	eh.e_ident[6] = EV_CURRENT;
	eh.e_type = ET_REL;
	eh.e_machine = EM_X86_64;
	eh.e_version = EV_CURRENT;
	eh.e_shoff = off_shdr;
	eh.e_ehsize = sizeof(Ehdr);
	eh.e_shentsize = sizeof(Shdr);
	eh.e_shnum = NSECT;
	eh.e_shstrndx = 5;

	fwrite(&eh, 1, sizeof eh, f);
	if (a->textlen) fwrite(a->text, 1, a->textlen, f);
	fwrite(rela.p, 1, rela.n, f);
	fwrite(sym.p, 1, sym.n, f);
	fwrite(str.p, 1, str.n, f);
	fwrite(shstr.p, 1, shstr.n, f);
	while ((u64)ftell(f) % 8) fputc(0, f);

	/* section header table */
	Shdr sh = {0};
	fwrite(&sh, 1, sizeof sh, f);	/* SHT_NULL */

	memset(&sh, 0, sizeof sh);
	sh.sh_name = shn_text;
	sh.sh_type = SHT_PROGBITS;
	sh.sh_flags = SHF_ALLOC | SHF_EXECINSTR;
	sh.sh_offset = off_text;
	sh.sh_size = a->textlen;
	sh.sh_addralign = 1;
	fwrite(&sh, 1, sizeof sh, f);

	memset(&sh, 0, sizeof sh);
	sh.sh_name = shn_rela;
	sh.sh_type = SHT_RELA;
	sh.sh_flags = SHF_INFO_LINK;
	sh.sh_offset = off_rela;
	sh.sh_size = rela.n;
	sh.sh_link = 3;	/* symtab */
	sh.sh_info = 1;	/* applies to .text */
	sh.sh_addralign = 8;
	sh.sh_entsize = sizeof(Rela64);
	fwrite(&sh, 1, sizeof sh, f);

	memset(&sh, 0, sizeof sh);
	sh.sh_name = shn_symtab;
	sh.sh_type = SHT_SYMTAB;
	sh.sh_offset = off_sym;
	sh.sh_size = sym.n;
	sh.sh_link = 4;	/* strtab */
	sh.sh_info = 1;	/* one local: STN_UNDEF */
	sh.sh_addralign = 8;
	sh.sh_entsize = sizeof(Sym64);
	fwrite(&sh, 1, sizeof sh, f);

	memset(&sh, 0, sizeof sh);
	sh.sh_name = shn_strtab;
	sh.sh_type = SHT_STRTAB;
	sh.sh_offset = off_str;
	sh.sh_size = str.n;
	sh.sh_addralign = 1;
	fwrite(&sh, 1, sizeof sh, f);

	memset(&sh, 0, sizeof sh);
	sh.sh_name = shn_shstrtab;
	sh.sh_type = SHT_STRTAB;
	sh.sh_offset = off_shstr;
	sh.sh_size = shstr.n;
	sh.sh_addralign = 1;
	fwrite(&sh, 1, sizeof sh, f);

	free(shstr.p); free(str.p); free(sym.p); free(rela.p);
	return 0;
}