// ============================================================================
// efi_elf_parser.cpp - Parse and load an ELF64 kernel binary
//
// Why we don't just execute the file in-place:
//  1) Layout mismatch — the on-disk ELF layout differs from the in-memory
//     layout (segments are not necessarily contiguous or page-aligned).
//  2) Fixed link addresses — the kernel is ET_EXEC (non-relocatable) linked
//     at a specific virtual address (e.g. 0x100000 for x86_64).
//  3) .bss - the .bss segment has p_filesz == 0 but p_memsz > 0; those
//     zero-initialised globals occupy no space in the file but need real
//     zeroed memory at runtime.
//
// So we load the file into a temporary buffer, then parse + copy segments
// into a freshly allocated region at the correct address.
// ============================================================================

#include <elf/elf.h>
#include <boot/efi.h>

namespace {

// Local memcpy/memset — we have no libc in the EFI environment.
static void* memcpy_local(void* dst, const void* src, uint64_t n) noexcept {
    auto* d = static_cast<uint8_t*>(dst);
    const auto* s = static_cast<const uint8_t*>(src);
    for (uint64_t i = 0; i < n; ++i) d[i] = s[i];
    return dst;
}

static void* memset_local(void* dst, int val, uint64_t n) noexcept {
    auto* d = static_cast<uint8_t*>(dst);
    for (uint64_t i = 0; i < n; ++i) d[i] = static_cast<uint8_t>(val);
    return dst;
}

} // anonymous namespace

// Target machine type for ELF validation — set by the build system.
// TODO: somehow remove it to arch-specific code
#if defined(EFI_ARCH_X86_64)
    static constexpr uint16_t EXPECTED_ARCH = EM_X86_64;
#elif defined(EFI_ARCH_AARCH64)
    static constexpr uint16_t EXPECTED_ARCH = EM_AARCH64;
#else
    #error "Unknown EFI_ARCH"
#endif

ElfLoadResult elf_load(const uint8_t* file_data, uint64_t file_size, EFI_BOOT_SERVICES* bs) {
    ElfLoadResult result{};
    result.success = false;
    result.error = ELF_ERR_NONE;

    // The file must be at least as large as the ELF header.
    if (file_size < sizeof(Elf64_Ehdr)) {
        result.error = ELF_ERR_FILE_TOO_SMALL;
        return result;
    }

    // Got ELF header
    const auto* ehdr = reinterpret_cast<const Elf64_Ehdr*>(file_data);

    // Validate magic, class, endianness, machine type, and object type.
    if (*reinterpret_cast<const uint32_t*>(ehdr->e_ident) != ELF_MAGIC) {
        result.error = ELF_ERR_INVALID_MAGIC;
        return result;
    }
    if (ehdr->e_ident[EI_CLASS] != ELFCLASS64) {
        result.error = ELF_ERR_INVALID_CLASS;
        return result;
    }
    if (ehdr->e_ident[EI_DATA] != ELFDATA2LSB) {
        result.error = ELF_ERR_INVALID_IDENT;
        return result;
    }
    if (ehdr->e_machine != EXPECTED_ARCH) {
        result.error = ELF_ERR_INVALID_ARCH;
        return result;
    }
    if (ehdr->e_type != ET_EXEC) {
        result.error = ELF_ERR_INVALID_TYPE;
        return result;
    }

    // Set entry point
    result.entry_point = ehdr->e_entry;

    // First pass: find the virtual address span of all PT_LOAD segments.
    // This tells us how many pages to allocate and where.
    uint64_t vaddr_min = ~0ULL;
    uint64_t vaddr_max = 0;

    /* Get Program Header */
    const auto* phdrs = reinterpret_cast<const Elf64_Phdr*>(file_data + ehdr->e_phoff);

    /* Iterates via headers and get only loadable segments and computes min and max vaddr
        to determine the total memory footprint the kernel needs.
     */
    for (uint16_t i = 0; i < ehdr->e_phnum; i++) {
        const auto& ph = phdrs[i];
        if (ph.p_type != PT_LOAD || ph.p_memsz == 0) continue;
        if (ph.p_vaddr < vaddr_min) vaddr_min = ph.p_vaddr;
        if (ph.p_vaddr + ph.p_memsz > vaddr_max) vaddr_max = ph.p_vaddr + ph.p_memsz;
    }

    if (vaddr_min >= vaddr_max) {
        result.error = ELF_ERR_NO_LOAD_SEGS;
        return result;
    }

    result.virt_base  = vaddr_min;
    result.total_size = vaddr_max - vaddr_min;

    // Allocate pages at the kernel's linked virtual address so that
    // absolute references in the non-relocatable ELF resolve correctly.
    const uint64_t num_pages  = (result.total_size + ELF_PAGE_MASK) / ELF_PAGE_SIZE;
    EFI_PHYSICAL_ADDRESS phys_addr = vaddr_min;

    EFI_STATUS status = bs->AllocatePages(AllocateAddress, EfiLoaderData, num_pages, &phys_addr);
    if (EFI_ERROR(status)) {
        result.error = ELF_ERR_ALLOC_FAILED;
        result.efi_alloc_status = status;
        return result;
    }

    result.phys_base = phys_addr;

    // Zero the entire region — this covers .bss and any gaps between segments.
    memset_local(reinterpret_cast<void*>(phys_addr), 0, num_pages * ELF_PAGE_SIZE);

    // Second pass: load PT_LOAD segments
    // Iterates via PT_LOAD segments and copying each segment's file data (p_filesz) to the
    // correct offset within the allocted region.
    // Segments where p_memsz > p_filesz (e.g. .bss) are already zeroed above.
    for (uint16_t i = 0; i < ehdr->e_phnum; i++) {
        const auto& ph = phdrs[i];
        if (ph.p_type != PT_LOAD || ph.p_memsz == 0) continue;

        const uint64_t offset_in_image = ph.p_vaddr - vaddr_min;
        auto* dest = reinterpret_cast<uint8_t*>(phys_addr + offset_in_image);

        if (ph.p_filesz > 0)
            memcpy_local(dest, file_data + ph.p_offset, ph.p_filesz);
    }

    result.success = true;
    return result;
}