#include <stdint.h>
#include <stdbool.h>
#include <stdio.h>
#include <string.h>
#include <stdbool.h>
#include <assert.h>

#define ET_NONE 0x00
#define ET_REL  0x01
#define ET_EXEC 0x02
#define ET_DYN  0x03
#define ET_CORE 0x04

struct Elf32Header
{
    uint8_t  ident[16];
    uint16_t type;
    uint16_t machine;
    uint32_t version;
    uint32_t entry;
    uint32_t phoff;
    uint32_t shoff;
    uint32_t flags;
    uint16_t ehsize;
    uint16_t phentsize;
    uint16_t phnum;
    uint16_t shentsize;
    uint16_t shnum;
    uint16_t shstrndx;
};

#define PT_NULL    0x00000000
#define PT_LOAD    0x00000001
#define PT_DYNAMIC 0x00000002
#define PT_INTERP  0x00000003
#define PT_NOTE    0x00000004
#define PT_SHLIB   0x00000005
#define PT_PHDR    0x00000006
#define PT_TLS     0x00000007

#define PF_X 0x1
#define PF_W 0x2
#define PF_R 0x4

struct Elf32ProgramHeader
{
    uint32_t type;
    uint32_t offset;
    uint32_t vaddr;
    uint32_t paddr;
    uint32_t filesz;
    uint32_t memsz;
    uint32_t flags;
    uint32_t align;
};

bool load_elf(const char* file, char* mem, uint32_t mem_size, uint32_t* start_address)
{
    FILE* fp = fopen(file, "rb");
    if (fp == NULL)
    {
        printf("Couldn't open input file\n");
        return false;
    }

    struct Elf32Header header;
    size_t read = fread(&header, sizeof(struct Elf32Header), 1, fp);
    if (read != 1)
    {
        printf("ELF header too small\n");
        return false;
    }

    static uint8_t kValidIdent[7] = {
        0x7f, 0x45, 0x4c, 0x46,
        1, 1, 1
    };
    if (memcmp(kValidIdent, header.ident, 7) != 0)
    {
        printf("Invalid ELF header\n");
        return false;
    }

    if (header.type != ET_EXEC)
    {
        printf("Can only link EXEC ELF files\n");
        return false;
    }

    memset(mem, 0, mem_size);
    
    assert(header.phnum == 0 || header.phentsize == sizeof(struct Elf32ProgramHeader));
    for (uint32_t i = 0; i < header.phnum; ++i)
    {
        struct Elf32ProgramHeader pheader;

        fseek(fp, header.phoff + i * header.phentsize, SEEK_SET);
        read = fread(&pheader, sizeof(struct Elf32ProgramHeader), 1, fp);
        if (read != 1)
        {
            printf("Error reading program header at index %u\n", i);
            return false;
        }

        if (pheader.type == PT_LOAD)
        {
            if (pheader.paddr + pheader.memsz > mem_size)
            {
                printf("Memory not large enough for ELF file\n");
                return false;
            }

            fseek(fp, pheader.offset, SEEK_SET);
            read = fread(mem + pheader.paddr, pheader.filesz, 1, fp);
            if (read != 1)
            {
                printf("Error when copying ELF segment\n");
                return false;
            }
        }
    }

    if (header.entry > mem_size)
    {
        printf("Entry address out of bounds\n");
        return false;
    }

    *start_address = header.entry;
    fclose(fp);

    return true;
}