/* Copyright (c) 2024 Freya Murphy */

#ifndef __LINK_H__
#define __LINK_H__

#include <linux/limits.h>
#include <mlimits.h>
#include <merror.h>
#include <stdint.h>
#include <elf.h>


// when mapping porinters, we need to bounds check to
// make sure its in the mapped object file
//
// this checks that
// 1. the end is in the file
// 2. the off and len doesnt integer overflow
#define BOUND_CHK(obj, len, off) \
	(off > UINT32_MAX - len || off + len > obj->mapped_size)

// when relocating segments, we need to bounds check to
// make sure it wont overflow the addresses past the 32bit
// ELF file
#define ADDR_CHK(lnk_f, seg_f, max) \
	((lnk_f) > max - (seg_f) || (lnk_f) + (seg_f) > max)

// checks if a phdr and shdr are matches
#define PHDR_SHDR_MATCH(phdr, shdr) (              \
	((phdr)->p_offset == (shdr)->sh_offset) && \
	((phdr)->p_filesz == (shdr)->sh_size))     \

// start addresses for each tyoe of segment
#define TEXT_VADDR_MIN 0x00400000
#define DATA_VADDR_MIN 0x10000000

// alignment of a section
#define SEC_ALIGN 0x1000

// pre define
struct linker;
struct object;
struct segment;
struct string_table;
struct symbol_table;
struct symbol_table_mapping;
struct glboff_table;

///
/// relocation table
///

struct relocation_table {
	uint32_t type;
	union {
		void *raw;
		Elf32_Rel  *rel;
		Elf32_Rela *rela;
	};
	size_t len;
	struct symbol_table *symtab;
};

///
/// string table
///

struct string_table {
	char *data;
	size_t len;
};

int strtab_init(struct string_table *strtab);
void strtab_free(struct string_table *strtab);

int strtab_push(struct string_table *strtab, const char *str, size_t *res);
int strtab_get(struct string_table *strtab, const char *str, size_t *res);

///
/// symbol table
///

struct symbol_table {
	struct string_table *strtab;
	Elf32_Sym *syms;
	size_t len;
	size_t size;

	// mapping for global symbol table
	// can be null
	// (metadata for each symbol)
	struct symbol_table_mapping *map;
};

int symtab_init(struct symbol_table *symtab);
void symtab_free(struct symbol_table *symtab);

int symtab_push(struct symbol_table *symtab, const Elf32_Sym *sym);
int symtab_get(struct symbol_table *symtab, Elf32_Sym **sym, const char *name,
	       int32_t obj_idx);

///
/// symbol table map
/// (metadata for each symbol)
/// stores each object that a symbol
/// is in
///

struct symbol_table_mapping {
	uint32_t len;
	uint32_t size;
	struct object **meta;
};

int symtab_map_push(struct symbol_table_mapping *symtabm, struct object *seg);

int symtab_map_init(struct symbol_table_mapping *symtabm);
void symtab_map_free(struct symbol_table_mapping *symtabm);

///
/// segment
///

/* a loadable program segment */
struct segment {
	// segment data
	char *name;
	unsigned char *bytes;

	// current loc
	uint32_t off;
	uint32_t vaddr;

	// new loc
	uint32_t new_off;
	uint32_t new_vaddr;

	// meta
	bool read;
	bool write;
	bool execute;
	uint32_t align;

	// size
	uint32_t size;

	// phdr
	Elf32_Phdr *phdr;
	uint32_t phdr_idx;

	// shdr
	Elf32_Shdr *shdr;
	uint32_t shdr_idx;

	// object im related to
	struct object *obj;
	// segment table entry im related to
	struct segment_table_entry *ent;

	// relocation table
	struct relocation_table reltab;
};

int segment_load(struct object *object, struct segment *segment, size_t index);

///
/// segment table
/// holds each segment that
/// share the same name
/// (i.e. all .text segments)
///

struct segment_table_entry {
	char *name;
	uint32_t len;
	uint32_t size;
	uint32_t off;
	uint32_t vaddr;
	// weak segment pointers. we do not own these!!!
	// (but we own the array)
	struct segment **parts;
};

int segtab_ent_init(struct segment_table_entry *ent);
void segtab_ent_free(struct segment_table_entry *ent);

int segtab_ent_push(struct segment_table_entry *ent, struct segment *seg);
uint32_t segtab_ent_size(struct segment_table_entry *ent);

///
/// segment table
/// holds each segment table entry by name
///
struct segment_table {
	uint32_t len;
	uint32_t size;
	struct segment_table_entry *entries;
};

int segtab_init(struct segment_table *segtab);
void segtab_free(struct segment_table *segtab);

/* create a new entry with <seg> as its first segment part */
int segtab_push(struct segment_table *segtab, struct segment_table_entry **ent,
		struct segment *seg);

/* find a segment table entry with a given name */
int segtab_get(struct segment_table *segtab, struct segment_table_entry **ent,
	       const char *name);

///
/// object file
///

struct object {
	// file
	int fd;
	char *mapped;
	size_t mapped_size;

	// ehdr
	Elf32_Ehdr *ehdr;

	// section header table
	Elf32_Shdr *shdr; // weak
	size_t shdr_len;

	// program table
	Elf32_Phdr *phdr; // owned
	size_t phdr_len;

	// phdr <=> shdr mappings
	// need a way to find any assoiciated
	// phdr with a shdr, and also reverse
	uint32_t *phdr_to_shdr_mapping;

	// object meta
	const char *name;
	size_t index;

	// segments
	size_t segment_len;
	struct segment *segments;

	// section header strtab
	struct string_table *shstrtab;

	// strtabs
	struct string_table *strtabs;
	// symtabs
	struct symbol_table *symtabs;
};

int object_load(struct object *object, char *path, uint32_t index);

void object_free(struct object *object);

///
/// linker
///

struct linker {
	size_t obj_len;
	struct object *objects;

	struct linker_arguments *args;

	// current pointers to relocate
	// sections
	uint32_t off;
	uint32_t text_vaddr;
	uint32_t data_vaddr;

	// elf tables
	struct string_table shstrtab;
	struct string_table strtab;
	struct symbol_table symtab;
	struct symbol_table_mapping symtab_map;

	// output elf
	Elf32_Ehdr ehdr;
	Elf32_Phdr *phdr;
	uint32_t phdr_len;
	Elf32_Shdr *shdr;
	uint32_t shdr_len;

	uint32_t symtab_shidx;
	uint32_t strtab_shidx;
	uint32_t shstrtab_shidx;

	// all segments
	struct segment_table segments;
};

/* defines arguments to the linker */
struct linker_arguments {
	char **in_files;
	int in_count;
	char *out_file;
	bool freestanding;
};

/* link object files */
int link_files(struct linker_arguments args);

#endif /* __LINK_H__ */