#include "binfmt/ELF.h"
#include "Log.h"
#include "arch/CPU.h"
#include "arch/MMU.h"
#include "memory/MemoryManager.h"
#include <bits/mmap-flags.h>
#include <luna/Alignment.h>
#include <luna/Alloc.h>
#include <luna/CString.h>
#include <luna/Common.h>
#include <luna/ScopeGuard.h>

static bool can_execute_segment(u32 flags)
{
    return flags & 1;
}

static bool can_write_segment(u32 flags)
{
    return flags & 2;
}

/*static bool can_write_and_execute_segment(u32 flags)
{
    return can_write_segment(flags) && can_execute_segment(flags);
}*/

Result<bool> ELFLoader::sniff()
{
    u8 buf[SELFMAG];
    usize nread = TRY(m_inode->read(buf, 0, sizeof buf));
    if (nread < SELFMAG) return false;

    return !memcmp(buf, ELFMAG, SELFMAG);
}

Result<u64> ELFLoader::load(AddressSpace* space)
{
    Elf64_Ehdr elf_header;
    usize nread = TRY(m_inode->read((u8*)&elf_header, 0, sizeof elf_header));

    if (nread < sizeof elf_header)
    {
        kerrorln("Error while loading ELF: ELF header does not fit in file");
        return err(ENOEXEC);
    }

    if (memcmp(elf_header.e_ident, ELFMAG, SELFMAG) != 0)
    {
        kerrorln("Error while loading ELF: ELF header has no valid magic");
        return err(ENOEXEC);
    }

    if (elf_header.e_ident[EI_CLASS] != ELFCLASS64)
    {
        kerrorln("Error while loading ELF: ELF object is not 64-bit");
        return err(ENOEXEC);
    }

    if (elf_header.e_ident[EI_DATA] != ELFDATA2LSB)
    {
        kerrorln("Error while loading ELF: ELF object is not 2's complement little-endian");
        return err(ENOEXEC);
    }

    if (elf_header.e_type != ET_EXEC)
    {
        kerrorln("Error while loading ELF: ELF object is not an executable");
        return err(ENOEXEC);
    }

    if (elf_header.e_machine != EM_MACH)
    {
        kerrorln("Error while loading ELF: ELF object's target architecture does not match the current one (%s)",
                 CPU::platform_string().chars());
        return err(ENOEXEC);
    }

    if (elf_header.e_phnum == 0)
    {
        kerrorln("Error while loading ELF: ELF object has no program headers");
        return err(ENOEXEC);
    }

#ifdef ELF_DEBUG
    kdbgln("ELF: Loading ELF with entry=%#.16lx", elf_header.e_entry);
#endif

    usize i;
    Elf64_Phdr program_header;

    for (TRY(m_inode->read((u8*)&program_header, elf_header.e_phoff, sizeof program_header)), i = 0;
         i < elf_header.e_phnum;
         i++, TRY(m_inode->read((u8*)&program_header, elf_header.e_phoff + (i * elf_header.e_phentsize),
                                sizeof program_header)))
    {
        if (program_header.p_type == PT_LOAD)
        {
#ifdef ELF_DEBUG
            kdbgln("ELF: Loading segment (offset=%zu, base=%#.16lx, filesize=%zu, memsize=%zu)",
                   program_header.p_offset, program_header.p_vaddr, program_header.p_filesz, program_header.p_memsz);
#endif

            u64 base_vaddr = align_down<ARCH_PAGE_SIZE>(program_header.p_vaddr);
            u64 vaddr_diff = program_header.p_vaddr - base_vaddr;
            /*expect(!can_write_and_execute_segment(program_header.p_flags),
                   "Segment is both writable and executable");*/

            int flags = MMU::User | MMU::NoExecute;
            if (can_write_segment(program_header.p_flags)) flags |= MMU::ReadWrite;
            if (can_execute_segment(program_header.p_flags)) flags &= ~MMU::NoExecute;

            int prot = PROT_READ;
            if (can_write_segment(program_header.p_flags)) prot |= PROT_WRITE;
            if (can_execute_segment(program_header.p_flags)) prot |= PROT_EXEC;

            if (!TRY(space->test_and_alloc_region(base_vaddr,
                                                  ceil_div(program_header.p_memsz + vaddr_diff, ARCH_PAGE_SIZE), prot,
                                                  MAP_ANONYMOUS | MAP_PRIVATE, 0, true)))
                return err(ENOMEM);

            // Allocate physical memory for the segment
            TRY(MemoryManager::alloc_at(base_vaddr, ceil_div(program_header.p_memsz + vaddr_diff, ARCH_PAGE_SIZE),
                                        flags));

            // Zero out unused memory (before the start of the segment)
            memset((void*)base_vaddr, 0, vaddr_diff);

            // Load the file section of the segment
            m_inode->read((u8*)program_header.p_vaddr, program_header.p_offset, program_header.p_filesz);

            // Fill out the rest of the segment with 0s
            memset((void*)(program_header.p_vaddr + program_header.p_filesz), 0,
                   program_header.p_memsz - program_header.p_filesz);
        }
        else { kwarnln("ELF: Encountered non-loadable program header, skipping"); }
    }

    return elf_header.e_entry;
}

Result<Vector<String>> ELFLoader::cmdline(const String&, Vector<String> args)
{
    return args;
}

ELFLoader::ELFLoader(SharedPtr<VFS::Inode> inode, int recursion_level) : BinaryFormatLoader(inode, recursion_level)
{
}

Result<SharedPtr<BinaryFormatLoader>> ELFLoader::create(SharedPtr<VFS::Inode> inode, void*, int recursion_level)
{
    return (SharedPtr<BinaryFormatLoader>)TRY(make_shared<ELFLoader>(inode, recursion_level));
}