#define MODULE "elf"

#include "sys/elf/ELFLoader.h"
#include "fs/VFS.h"
#include "init/InitRD.h"
#include "log/Log.h"
#include "memory/Memory.h"
#include "memory/MemoryManager.h"
#include "memory/VMM.h"
#include "misc/utils.h"
#include "std/assert.h"
#include "std/errno.h"
#include "std/stdlib.h"
#include "std/string.h"
#include "sys/elf/ELF.h"
#include "utils/Addresses.h"

static const char* format_permissions(uint32_t flags)
{
    static char perms[4];
    perms[0] = (flags & 4) > 0 ? 'r' : '-';
    perms[1] = (flags & 2) > 0 ? 'w' : '-';
    perms[2] = (flags & 1) > 0 ? 'x' : '-';
    perms[3] = 0;
    return perms;
}

ELFImage* ELFLoader::load_elf_from_filesystem(const char* filename)
{
    VFS::Node* node = VFS::resolve_path(filename);

    if (!node)
    {
        kwarnln("Failed to open file %s for loading", filename);
        return 0;
    }

    if (node->type == VFS_DIRECTORY)
    {
        kwarnln("Failed to load %s: is a directory", filename);
        return 0;
    }

    ELFImage* result = load_elf_from_vfs(node);
    return result;
}

long ELFLoader::check_elf_image_from_filesystem(const char* filename)
{
    VFS::Node* node = VFS::resolve_path(filename);

    if (!node)
    {
        kwarnln("Failed to open file %s for checking", filename);
        return -ENOENT;
    }

    if (node->type == VFS_DIRECTORY)
    {
        kwarnln("Failed to check %s: is a directory", filename);
        return -EISDIR;
    }

    return check_elf_image(node);
}

ELFImage* ELFLoader::load_elf_from_vfs(VFS::Node* node)
{
    Elf64_Ehdr elf_ehdr;
    ASSERT(VFS::read(node, 0, sizeof(elf_ehdr), (char*)&elf_ehdr) >= 0);
    ASSERT(strncmp((const char*)elf_ehdr.e_ident, ELFMAG, SELFMAG) ==
           0); // If you haven't checked the ELF executable with check_elf_image() first, then an assertion fail is your
               // fault =D
    ASSERT(elf_ehdr.e_ident[EI_CLASS] == ELFCLASS64);
    ASSERT(elf_ehdr.e_ident[EI_DATA] == ELFDATA2LSB);
    ASSERT(elf_ehdr.e_type == ET_EXEC);
    ASSERT(elf_ehdr.e_machine == EM_MACH);
    ASSERT(elf_ehdr.e_phnum != 0);
    ELFImage* image = (ELFImage*)kmalloc(sizeof(ELFImage) - sizeof(ELFSection));
    memset(image, 0, sizeof(ELFImage) - sizeof(ELFSection));
    image->entry = elf_ehdr.e_entry;
    int i;
    Elf64_Phdr phdr;
    for (VFS::read(node, elf_ehdr.e_phoff, sizeof(Elf64_Phdr), (char*)&phdr), i = 0; i < elf_ehdr.e_phnum;
         i++, VFS::read(node, elf_ehdr.e_phoff + (i * elf_ehdr.e_phentsize), sizeof(Elf64_Phdr), (char*)&phdr))
    {
        if (phdr.p_type == PT_LOAD)
        {
            kdbgln("Loading loadable segment at address %lx, file size %ld, mem size %ld, permissions %s", phdr.p_vaddr,
                   phdr.p_filesz, phdr.p_memsz, format_permissions(phdr.p_flags));
            ASSERT(phdr.p_vaddr);

            uint64_t pages = Utilities::get_blocks_from_size(PAGE_SIZE, (phdr.p_vaddr % PAGE_SIZE) + phdr.p_memsz);
            void* buffer = (void*)((uint64_t)MemoryManager::get_pages_at(round_down_to_nearest_page(phdr.p_vaddr),
                                                                         pages, MAP_READ_WRITE) +
                                   (phdr.p_vaddr % PAGE_SIZE));

            if (VMM::is_using_kernel_address_space()) { VMM::switch_to_previous_user_address_space(); }
            VMM::apply_address_space();

            VFS::read(node, phdr.p_offset, phdr.p_filesz, (char*)buffer);
            memset((void*)((uint64_t)buffer + phdr.p_filesz), 0, phdr.p_memsz - phdr.p_filesz);

            VMM::switch_back_to_kernel_address_space();
            VMM::apply_address_space();
            VMM::switch_to_previous_user_address_space();

            MemoryManager::protect(buffer, pages, phdr.p_flags & 2 ? MAP_READ_WRITE | MAP_USER : MAP_USER);

            image = (ELFImage*)krealloc(image, (sizeof(ELFImage) - sizeof(ELFSection)) +
                                                   (image->section_count + 1) * sizeof(ELFSection));
            ELFSection& section = image->sections[image->section_count];
            section.base = (uintptr_t)buffer;
            section.pages = pages;
            image->section_count++;
        }
        else { kdbgln("skipping non-loadable segment"); }
    }
    ASSERT(image->section_count);
    return image;
}

long ELFLoader::check_elf_image(VFS::Node* node)
{
    Elf64_Ehdr elf_ehdr;
    if (VFS::read(node, 0, sizeof(elf_ehdr), (char*)&elf_ehdr) < (long)sizeof(elf_ehdr))
    {
        kwarnln("Unable to read ELF header");
        return -ENOEXEC;
    }
    if (strncmp((const char*)elf_ehdr.e_ident, ELFMAG, SELFMAG) != 0)
    {
        kwarnln("ELF file has invalid magic, skipping");
        return -ENOEXEC;
    }
    if (elf_ehdr.e_ident[EI_CLASS] != ELFCLASS64)
    {
        kwarnln("ELF file is not ELF64, skipping");
        return -ENOEXEC;
    }
    if (elf_ehdr.e_ident[EI_DATA] != ELFDATA2LSB)
    {
        kwarnln("ELF file is not little-endian, skipping");
        return -ENOEXEC;
    }
    if (elf_ehdr.e_type != ET_EXEC)
    {
        kwarnln("not supported: ELF file is not an executable");
        return -ENOEXEC;
    }
    if (elf_ehdr.e_machine != EM_MACH)
    {
        kwarnln("Unsupported target machine");
        return -ENOEXEC;
    }
    if (elf_ehdr.e_phnum == 0)
    {
        kwarnln("ELF file has no PHDRS");
        return -ENOEXEC;
    }
    int i;
    int loadable_sections = 0;
    long memusage = 0;
    Elf64_Phdr phdr;
    for (VFS::read(node, elf_ehdr.e_phoff, sizeof(Elf64_Phdr), (char*)&phdr), i = 0; i < elf_ehdr.e_phnum;
         i++, VFS::read(node, elf_ehdr.e_phoff + (i * elf_ehdr.e_phentsize), sizeof(Elf64_Phdr), (char*)&phdr))
    {
        if (phdr.p_type == PT_LOAD)
        {
            if (!phdr.p_vaddr)
            {
                kerrorln("segment address is NULL, this is invalid :(");
                return -ENOEXEC;
            }
            if (Memory::is_kernel_address(phdr.p_vaddr) || Memory::is_kernel_address(phdr.p_vaddr + phdr.p_memsz))
            {
                kerrorln("trying to load ELF into kernel memory");
                return -ENOEXEC;
            }
            loadable_sections++;
            memusage += Utilities::get_blocks_from_size(PAGE_SIZE, phdr.p_memsz) * PAGE_SIZE;
        }
    }
    if (!loadable_sections)
    {
        kwarnln("No loadable sections");
        return -ENOEXEC;
    }
    return memusage;
}

void ELFLoader::release_elf_image(ELFImage* image)
{
    for (uint64_t i = 0; i < image->section_count; i++)
    {
        ELFSection& section = image->sections[i];
        kdbgln("Freeing up section %lx, was using %ld pages", section.base, section.pages);
        MemoryManager::release_pages((void*)round_down_to_nearest_page(section.base), section.pages);
    }
    kfree(image);
}