#include "elf.h"
#include "kernel.h"
#include "klibc.h"
#include "mem.h"
#include "paging.h"
#include "thread.h"
#include "types.h"

/**
 * Make sure the program is in a valid ELF format, map it into memory,
 * and return the address of its entry point (ie _start function)
 *
 * @return 0 when the program is not a valid ELF
 */

uaddr_t loadElfProg(const char *prog, struct process *proc)
{
    int i;
    uaddr_t lastUserAddr = 0;

/* e_ident value */
#define ELFMAG0 0x7f
#define ELFMAG1 'E'
#define ELFMAG2 'L'
#define ELFMAG3 'F'

/* e_ident offsets */
#define EI_MAG0 0
#define EI_MAG1 1
#define EI_MAG2 2
#define EI_MAG3 3
#define EI_CLASS 4
#define EI_DATA 5
#define EI_VERSION 6
#define EI_PAD 7

/* e_ident[EI_CLASS] */
#define ELFCLASSNONE 0
#define ELFCLASS32 1
#define ELFCLASS64 2

/* e_ident[EI_DATA] */
#define ELFDATANONE 0
#define ELFDATA2LSB 1
#define ELFDATA2MSB 2

/* e_type */
#define ET_NONE 0        /* No file type       */
#define ET_REL 1         /* Relocatable file   */
#define ET_EXEC 2        /* Executable file    */
#define ET_DYN 3         /* Shared object file */
#define ET_CORE 4        /* Core file          */
#define ET_LOPROC 0xff00 /* Processor-specific */
#define ET_HIPROC 0xffff /* Processor-specific */

/* e_machine */
#define EM_NONE 0  /* No machine     */
#define EM_M32 1   /* AT&T WE 32100  */
#define EM_SPARC 2 /* SPARC          */
#define EM_386 3   /* Intel 80386    */
#define EM_68K 4   /* Motorola 68000 */
#define EM_88K 5   /* Motorola 88000 */
#define EM_860 7   /* Intel 80860    */
#define EM_MIPS 8  /* MIPS RS3000    */

/* e_version */
#define EV_NONE 0    /* invalid version */
#define EV_CURRENT 1 /* current version */


/* Reserved segment types p_type */
#define PT_NULL 0
#define PT_LOAD 1
#define PT_DYNAMIC 2
#define PT_INTERP 3
#define PT_NOTE 4
#define PT_SHLIB 5
#define PT_PHDR 6
#define PT_LOPROC 0x70000000
#define PT_HIPROC 0x7fffffff

/* p_flags */
#define PF_X 1
#define PF_W 2
#define PF_R 4

    Elf32_Ehdr_t *elf_hdr = (Elf32_Ehdr_t *)prog;
    Elf32_Phdr_t *elf_phdrs;

    /* Macro to check expected values for some fields in the ELF header */
#define ELF_CHECK(hdr, field, expected_value)                                                 \
    ({                                                                                        \
        if ((hdr)->field != (expected_value)) {                                               \
            printf("ELF prog : for %s, expected %x, got %x\n", #field,                        \
                   (unsigned)(expected_value), (unsigned)(hdr)->field);                       \
            return 0;                                                                         \
        }                                                                                     \
    })

    ELF_CHECK(elf_hdr, e_ident[EI_MAG0], ELFMAG0);
    ELF_CHECK(elf_hdr, e_ident[EI_MAG1], ELFMAG1);
    ELF_CHECK(elf_hdr, e_ident[EI_MAG2], ELFMAG2);
    ELF_CHECK(elf_hdr, e_ident[EI_MAG3], ELFMAG3);
    ELF_CHECK(elf_hdr, e_ident[EI_CLASS], ELFCLASS32);
    ELF_CHECK(elf_hdr, e_ident[EI_DATA], ELFDATA2LSB);
    ELF_CHECK(elf_hdr, e_type, ET_EXEC);
    ELF_CHECK(elf_hdr, e_version, EV_CURRENT);

    /* Get the begining of the program header table */
    elf_phdrs = (Elf32_Phdr_t *)(prog + elf_hdr->e_phoff);

    /* Map the program segment in R/W mode. To make things clean, we
       should iterate over the sections, not the program header */
    for (i = 0; i < elf_hdr->e_phnum; i++) {
        uaddr_t uaddr;

        /* Ignore the empty program headers that are not marked "LOAD" */
        if (elf_phdrs[i].p_type != PT_LOAD) {
            continue;
        }

        if (elf_phdrs[i].p_vaddr < PAGING_BASE_USER_ADDRESS) {
            printf("User program has an incorrect address 0x%p\n", (void *)elf_phdrs[i].p_vaddr);
            return (uaddr_t)NULL;
        }

        /* Map pages of physical memory into user space */
        for (uaddr = ALIGN_DOWN(elf_phdrs[i].p_vaddr, PAGE_SIZE);
             uaddr < elf_phdrs[i].p_vaddr + elf_phdrs[i].p_memsz; uaddr += PAGE_SIZE) {
            paddr_t ppage;
            ppage = allocPhyPage(1);

            assert(pageMap(uaddr, ppage,
                           PAGING_MEM_USER | PAGING_MEM_WRITE | PAGING_MEM_READ) == 0);

            unrefPhyPage(ppage);
        }

        /* Copy segment into memory */
        memcpy((void *)elf_phdrs[i].p_vaddr, (void *)(prog + elf_phdrs[i].p_offset),
               elf_phdrs[i].p_filesz);
        if (lastUserAddr < uaddr) {
            lastUserAddr = uaddr;
        }
    }

    processInitHeap(proc, lastUserAddr);

    return elf_hdr->e_entry;
}