#define MODULE "elf"

#include "sys/elf/ELFLoader.h"
#include "init/InitRD.h"
#include "log/Log.h"
#include "memory/MemoryManager.h"
#include "misc/utils.h"
#include "std/string.h"
#include "sys/elf/ELF.h"

void* ELFLoader::load_elf_from_initrd(const char* filename)
{
    InitRD::File elf_file = InitRD::open(filename);
    if (!elf_file.addr)
    {
        kwarnln("failed to open file %s for loading", filename);
        return 0;
    }

    return load_elf_from_address((uintptr_t)elf_file.addr);
}

void* ELFLoader::load_elf_from_address(uintptr_t addr)
{
    Elf64_Ehdr* elf_ehdr = (Elf64_Ehdr*)addr;
    if (strncmp((const char*)elf_ehdr->e_ident, ELFMAG, SELFMAG) != 0)
    {
        kwarnln("ELF file has invalid magic, skipping");
        return 0;
    }
    if (elf_ehdr->e_ident[EI_CLASS] != ELFCLASS64)
    {
        kwarnln("ELF file is not ELF64, skipping");
        return 0;
    }
    if (elf_ehdr->e_ident[EI_DATA] != ELFDATA2LSB)
    {
        kwarnln("ELF file is not little-endian, skipping");
        return 0;
    }
    if (elf_ehdr->e_type != ET_EXEC)
    {
        kwarnln("not supported: ELF file is not an executable");
        return 0;
    }
    if (elf_ehdr->e_machine != EM_MACH)
    {
        kwarnln("unsupported target machine");
        return 0;
    }
    if (elf_ehdr->e_phnum == 0)
    {
        kwarnln("ELF file has no PHDRS");
        return 0;
    }
    int i;
    Elf64_Phdr* phdr;
    for (phdr = (Elf64_Phdr*)((uint64_t)addr + elf_ehdr->e_phoff), i = 0; i < elf_ehdr->e_phnum;
         i++, phdr = (Elf64_Phdr*)((uint8_t*)phdr + elf_ehdr->e_phentsize))
    {
        if (phdr->p_type == PT_LOAD)
        {
            kdbgln("loading loadable segment at vaddr %lx, filesz %ld, memsz %ld", phdr->p_vaddr, phdr->p_filesz,
                   phdr->p_memsz);
            if (!phdr->p_vaddr)
            {
                kerrorln("vaddr is NULL, this is invalid :(");
                return 0;
            }
            int pages = Utilities::get_blocks_from_size(0x1000, phdr->p_memsz);
            void* buffer = MemoryManager::get_pages_at(phdr->p_vaddr, pages,
                                                       phdr->p_flags & 2 ? MAP_READ_WRITE | MAP_USER : MAP_USER);
            memcpy(buffer, (void*)(addr + phdr->p_offset), phdr->p_filesz);
            memset((void*)((uint64_t)buffer + phdr->p_filesz), 0, phdr->p_memsz - phdr->p_filesz);
        }
        else { kdbgln("skipping non-loadable segment"); }
    }
    return (void*)elf_ehdr->e_entry;
}