#include "ELF.h"
#include "Log.h"
#include "arch/CPU.h"
#include "arch/MMU.h"
#include "memory/MemoryManager.h"
#include <luna/Alignment.h>
#include <luna/Alloc.h>
#include <luna/CString.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);
}*/

namespace ELFLoader
{
    // FIXME: Check that all calls to read_contents() read the proper amount of bytes.
    Result<ELFData> load(const TarStream::Entry& elf_entry, const TarStream& stream)
    {
        Elf64_Ehdr elf_header;
        usize nread = stream.read_contents(elf_entry, &elf_header, 0, sizeof elf_header);
        if (nread < sizeof elf_header)
        {
            kdbgln("Error while loading ELF: ELF header does not fit in entry");
            return err(ENOEXEC);
        }

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

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

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

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

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

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

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

        usize i;
        Elf64_Phdr program_header;

        for (stream.read_contents(elf_entry, &program_header, elf_header.e_phoff, sizeof program_header), i = 0;
             i < elf_header.e_phnum;
             i++, stream.read_contents(elf_entry, &program_header, elf_header.e_phoff + (i * elf_header.e_phentsize),
                                       sizeof program_header))
        {
            if (program_header.p_type == PT_LOAD)
            {
                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);

                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;

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

                // Load the file section of the segment
                stream.read_contents(elf_entry, (void*)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 { kdbgln("ELF: Encountered non-loadable program header, skipping"); }
        }

        return ELFData { elf_header.e_entry };
    }
}