#include <string.h>
#include <stdlib.h>
#include <stdio.h>
#include <mm/pmm.h>
#include <mm/paging.h>

#include <drivers/elf.h>

#define PAGE_SIZE 4096
#define ALIGN_DOWN(x) ((x) & ~(PAGE_SIZE - 1))
#define ALIGN_UP(x)   (((x) + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1))


elf_executable_t* load_elf32(void* elf_data)
{
  Elf32_Ehdr* ehdr = (Elf32_Ehdr*)elf_data;

  /* Check ELF magic */
  if (memcmp(ehdr->e_ident, (uint8_t*)("\x7F""ELF"), 4) != 0)
  {
    printf("Invalid ELF file\n");
    return NULL;
  }

  if (ehdr->e_machine != 3 || ehdr->e_type != 2)
  {
    printf("Unsupported ELF type or architecture\n");
    return NULL;
  }

  Elf32_Phdr* phdrs = (Elf32_Phdr*)((uint8_t*)elf_data + ehdr->e_phoff);

  for (int i = 0; i < ehdr->e_phnum; i++) {
    Elf32_Phdr* phdr = &phdrs[i];
    if (phdr->p_type != PT_LOAD)
    {
      continue;
    }

    uint32_t vaddr_start = ALIGN_DOWN(phdr->p_vaddr);
    uint32_t vaddr_end   = ALIGN_UP(phdr->p_vaddr + phdr->p_memsz);
    uint32_t segment_pages = (vaddr_end - vaddr_start) / PAGE_SIZE;

    for (uint32_t page = 0; page < segment_pages; page++)
    {
      void* phys = alloc_page();
      void* virt = (void*)(vaddr_start + page * PAGE_SIZE);
      map_page(phys, virt);
      memset(virt, 0, PAGE_SIZE);
    }

    void* dest = (void*)(uintptr_t)phdr->p_vaddr;
    void* src  = (uint8_t*)elf_data + phdr->p_offset;
    memcpy(dest, src, phdr->p_filesz);

    printf("Mapped segment %d: 0x%08x–0x%08x (memsz: %u, filesz: %u)\n", i, phdr->p_vaddr, phdr->p_vaddr + phdr->p_memsz, phdr->p_memsz, phdr->p_filesz);
  }

  elf_executable_t* result = malloc(sizeof(elf_executable_t));
  result->entry_point = (void*)(uintptr_t)ehdr->e_entry;
  return result;
}