/*
 * Copyright (c) 2026, Chloe M.
 * Provided under the BSD-3 clause.
 *
 * Description: ELF64 loader
 * Author:      Chloe M.
 */

#include <ex/elf64.h>
#include <ex/pool.h>
#include <hal/page.h>
#include <ptapi/ptmm.h>
#include <elf.h>
#include <mm/vm.h>
#include <string.h>

/* When allocating ELF related pools */
#define ELF_POOL_TAG 'ELF'

/* Maximum loadable program headers */
#define MAX_PHDRS 64

/* Prototypes */
static PT_STATUS Elf64Load(MMU_VAS *Vas, VOID *Data, LOADER_PROGRAM *Result);

/* ELF64 operations */
LOADER_OPS gElf64Ops = {
    .LoadProgram = Elf64Load
};

/*
 * Program data associated with a loader
 *
 * @LoadedSegments: Segments that have been loaded
 * @SegmentCount:   Number of loaded segments
 */
typedef struct {
    MM_RANGE LoadedSegments[MAX_PHDRS];
    USIZE SegmentCount;
} ELF64_PROGDATA;

/*
 * Verify the header of an ELF64 binary to see if it is
 * compatible with PTOS.
 */
static BOOLEAN
Elf64Verify(const Elf64_Ehdr *Eh)
{
    if (Eh == NULL) {
        return false;
    }

    if (RtlMemCmp(&Eh->e_ident[EI_MAG0], ELFMAG, SELFMAG) != 0)
        return false;
    if (Eh->e_ident[EI_CLASS] != ELFCLASS64)
        return false;
    if (Eh->e_ident[EI_DATA] != ELFDATA2LSB)
        return false;
    if (Eh->e_version != EV_CURRENT)
        return false;

    return true;
}

/*
 * Load an ELF64 program header
 *
 * @Vas:    Virtual address space to map within
 * @Phdr:   Program header to load
 * @Image:  Image result is written here
 */
static PT_STATUS
ElfLoadPhdr(MMU_VAS *Vas, Elf64_Ehdr *Eh, Elf64_Phdr *Phdr, ELF64_PROGDATA *Image)
{
    USIZE MapLength, MisAlign;
    USIZE PageCount;
    PT_STATUS Status;
    MM_PROT Prot = PAGE_READ | PAGE_WRITE | PAGE_USER;
    MM_RANGE AllocatedPages, MapRange;
    VOID *Src, *Dest;

    if (Vas == NULL || Eh == NULL) {
        return STATUS_INVALID_PARAM;
    }

    if (Phdr == NULL || Image == NULL) {
        return STATUS_INVALID_PARAM;
    }

    /* Set the permissions (TODO: Make unwritable by default) */
    if (ISSET(Phdr->p_flags, PF_X))
        Prot |= PAGE_EXEC;

    /* Compute the number of pages to map */
    MisAlign = Phdr->p_vaddr & (PAGESIZE - 1);
    MapLength = ALIGN_UP(Phdr->p_memsz + MisAlign, PAGESIZE);
    PageCount = MapLength / PAGESIZE;

    /* Grab some pages for our segment data */
    Status = MmAllocatePages(PageCount, &AllocatedPages);
    if (PT_ERROR(Status)) {
        return STATUS_NO_MEMORY;
    }

    MapRange.VirtualBase = Phdr->p_vaddr;
    MapRange.PhysicalBase = AllocatedPages.PhysicalBase;
    MapRange.Length = MapLength;
    Status = MmMapRange(
        Vas,
        &MapRange,
        Prot,
        PAGESIZE_4K
    );

    /* TODO: Free previously allocated pages */
    if (PT_ERROR(Status)) {
        return Status;
    }

    Src = PTR_OFFSET(Eh, Phdr->p_offset);
    Dest = PMA_TO_VMA(AllocatedPages.PhysicalBase);
    RtlMemCpy(Dest, Src, Phdr->p_filesz);
    Image->LoadedSegments[Image->SegmentCount++] = MapRange;
    return STATUS_SUCCESS;
}

/*
 * Load every program header of an ELF binary
 *
 * @Vas:    Virtual address space to map within
 * @Eh:     ELF header
 * @Image:  Loaded program image
 */
static PT_STATUS
ElfLoad(MMU_VAS *Vas, Elf64_Ehdr *Eh, LOADER_PROGRAM *Image)
{
    Elf64_Phdr *PhdrBase, *Phdr;
    USIZE PhdrIndex;
    PT_STATUS Status;

    if (Vas == NULL || Eh == NULL) {
        return STATUS_INVALID_PARAM;
    }

    if (Image == NULL) {
        return STATUS_INVALID_PARAM;
    }

    /* Allocate our own program data */
    Image->ProgramData = ExAllocatePoolWithTag(
        POOL_NON_PAGED,
        sizeof(ELF64_PROGDATA),
        ELF_POOL_TAG
    );

    if (Image->ProgramData == NULL) {
        return STATUS_NO_MEMORY;
    }

#define PHDR_INDEX(INDEX)   \
    PTR_OFFSET(PhdrBase, (INDEX) * Eh->e_phentsize)
    PhdrBase = PTR_OFFSET(Eh, Eh->e_phoff);
    for (PhdrIndex = 0; PhdrIndex < Eh->e_phnum; ++PhdrIndex) {
        Phdr = PHDR_INDEX(PhdrIndex);

        switch (Phdr->p_type) {
        case PT_LOAD:
            /* TODO: Clean up on failure */
            Status = ElfLoadPhdr(Vas, Eh, Phdr, Image->ProgramData);
            if (PT_ERROR(Status)) {
                ExFreePoolWithTag(Image->ProgramData, ELF_POOL_TAG);
                Image->ProgramData = NULL;
                return Status;
            }

            break;
        default:
            continue;
        }
    }
#undef PHDR_INDEX

    return STATUS_SUCCESS;
}

/*
 * Load an ELF64 program
 */
static PT_STATUS
Elf64Load(MMU_VAS *Vas, VOID *Data, LOADER_PROGRAM *Result)
{
    Elf64_Ehdr *Eh;
    PT_STATUS Status;

    if (Data == NULL || Result == NULL) {
        return STATUS_INVALID_PARAM;
    }

    /* Verify that it is executable */
    Eh = (Elf64_Ehdr *)Data;
    if (!Elf64Verify(Eh)) {
        return STATUS_NOT_EXEC;
    }

    Status = ElfLoad(
        Vas,
        Eh,
        Result
    );

    if (PT_ERROR(Status)) {
        return Status;
    }

    return STATUS_SUCCESS;
}