#include "elf.h"
#include "vfs.h"
#include "memory/mm.h"

#include <lib/memory.h>

uint32_t load_elf_program(const char* prgm) {
    int fd = vfs_open(prgm, 0);
    if (fd == -1) {
        return -1;
    }

    elf_header_t* e_hdr = kalloc(sizeof(elf_header_t));
    memset((void*)e_hdr, 0, sizeof(elf_header_t));

    uint32_t bytes_rd = vfs_read(fd, e_hdr, sizeof(elf_header_t));
    if (bytes_rd <= 0) {
        return -1;
    }

    uint32_t ph_table_ptr = e_hdr->e_phoff;
    elf_program_header_t* ph = kalloc(sizeof(elf_program_header_t));
    uint32_t heap_start = 0;

    for (uint32_t i = 0; i < e_hdr->e_phnum; i++) {
        vfs_seek(fd, ph_table_ptr);
        memset((void*)ph, 0, sizeof(elf_program_header_t));
        vfs_read(fd, ph, sizeof(elf_program_header_t));

        if (ph->p_type == PT_LOAD) {
            uint32_t vstart = ph->p_vaddr;
            uint32_t vend = vstart + ph->p_memsz;

            if (vend > heap_start) {
                heap_start = vend;
            }

            uint32_t page_start = vstart & ~(PAGE_SIZE - 1);
            uint32_t page_end = (vend + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1);

            for (uint32_t v = page_start; v < page_end; v += PAGE_SIZE) {
                void* phys_frame = p_alloc_frame();
                map_page(v, (uint32_t)phys_frame, PAGE_PRESENT | PAGE_WRITABLE | PAGE_USER);
            }

            uint32_t data_ptr = ph->p_offset;
            vfs_seek(fd, data_ptr);
            vfs_read(fd, (void*)vstart, ph->p_filesz);

            // 0 out bss
            if (ph->p_memsz > ph->p_filesz) {
                size_t bss_size = ph->p_memsz - ph->p_filesz;
                memset((void*)(vstart + ph->p_filesz), 0, bss_size);
            }
        }

        ph_table_ptr += sizeof(elf_program_header_t);
    }

    uint32_t eip = e_hdr->e_entry;
    kfree(ph);
    kfree(e_hdr);

    return eip;
}