#include "elf.h"
#include "stdio.h"
#include "string.h"
#include "mm/pmm.h"
#include "mm/vmm.h"
#include "mm/memory.h"
#include "fs/ext2.h"

extern uintptr_t g_hhdm_offset;
extern struct pagemap *kernel_pagemap;

#define ELF_BUFFER_SIZE (1024 * 1024)


#define PTE_PRESENT  (1ULL << 0)
#define PTE_WRITABLE (1ULL << 1)
#define PTE_USER     (1ULL << 2)

bool ELF_Read(const char* path, void** entryPoint)
{
    uint32_t size;

    uint8_t* elf_buffer = kmalloc(ELF_BUFFER_SIZE);
    if (!elf_buffer) {
        printf("ELF: kmalloc failed\n");
        return false;
    }

    // ── load file ─────────────────────────────────────
    if (!ext2_read_file_from_root(path, elf_buffer, &size)) {
        printf("ELF: failed to read file\n");
        kfree(elf_buffer);
        return false;
    }

    if (size < sizeof(ELFHeader)) {
        printf("ELF: file too small\n");
        kfree(elf_buffer);
        return false;
    }

    ELFHeader* header = (ELFHeader*)elf_buffer;

    // ── validate ELF ──────────────────────────────────
    if (memcmp(header->Magic, ELF_MAGIC, 4) != 0) {
        printf("ELF: bad magic\n");
        kfree(elf_buffer);
        return false;
    }

    if (header->Bitness != ELF_BITNESS_64BIT) {
        printf("ELF: not 64-bit\n");
        kfree(elf_buffer);
        return false;
    }

    if (header->Endianness != ELF_ENDIANNESS_LITTLE) {
        printf("ELF: wrong endianness\n");
        kfree(elf_buffer);
        return false;
    }

    if (header->Type != ELF_TYPE_EXECUTABLE) {
        printf("ELF: not executable\n");
        kfree(elf_buffer);
        return false;
    }

    if (header->InstructionSet != ELF_INSTRUCTION_SET_X64) {
        printf("ELF: wrong arch\n");
        kfree(elf_buffer);
        return false;
    }

    *entryPoint = (void*)header->ProgramEntryPosition;

    // ── program headers ───────────────────────────────
    uint8_t* ph_table = elf_buffer + header->ProgramHeaderTablePosition;

    for (uint32_t i = 0; i < header->ProgramHeaderTableEntryCount; i++)
    {
        ELFProgramHeader* ph = (ELFProgramHeader*)(ph_table +
            i * header->ProgramHeaderTableEntrySize);

        if (ph->Type != ELF_PROGRAM_TYPE_LOAD)
            continue;

        uint64_t virt   = ph->VirtualAddress;
        uint64_t offset = ph->Offset;
        uint64_t memsz  = ph->MemorySize;
        uint64_t filesz = ph->FileSize;

        if (memsz == 0)
            continue;

        // ── align to page boundary ─────────────────────
        uint64_t aligned_virt   = ALIGN_DOWN(virt, PAGE_SIZE);
        uint64_t page_offset    = virt & 0xFFF;
        uint64_t aligned_memsz  = ALIGN_UP(memsz + page_offset, PAGE_SIZE);

        uint64_t pages = aligned_memsz / PAGE_SIZE;

        // Allocate physical pages
        uint64_t phys_base = (uint64_t)pmm_alloc(pages);
        if (!phys_base) {
            printf("ELF: pmm_alloc failed for %lu pages\n", pages);
            kfree(elf_buffer);
            return false;
        }

        // ── map each page individually using new vmm_map_page ─────
        for (uint64_t p = 0; p < pages; p++) {
            uint64_t virt_addr = aligned_virt + p * PAGE_SIZE;
            uint64_t phys_addr = phys_base + p * PAGE_SIZE;

            bool success = vmm_map_page(
                kernel_pagemap,
                virt_addr,
                phys_addr,
                PAGE_READ | PAGE_WRITE | PAGE_USER,   // RW + User mode
                Size4KiB
            );

            if (!success) {
                printf("ELF: failed to map page at 0x%lx\n", virt_addr);
                // TODO: cleanup previously mapped pages + free phys
                kfree(elf_buffer);
                return false;
            }
        }

        // ── copy segment data ───────────────────────────────
        uint8_t* dst = (uint8_t*)(phys_base + MEM_PHYS_OFFSET);  // via HHDM
        uint8_t* src = elf_buffer + offset;

        memcpy(dst + page_offset, src, filesz);

        // ── zero BSS section ────────────────────────────────
        if (memsz > filesz) {
            memset(dst + page_offset + filesz, 0, memsz - filesz);
        }
    }

    kfree(elf_buffer);
    return true;
}