#include "ELF.h"
#include "memorytypes.h"

#include <efi.h>
#include <efilib.h>
#include <efidef.h>

UINTN EFIAPI elf_ceil2(float x)
{
  UINTN ret = x;
  if (x - (int)x > 0) {
      x += 1;
      ret = (int)x;
  }

  return ret;
}

void EFIAPI PrintELFInfo(ELF * file)
{
  Print(L"LoadAddress: 0x%x\n", file);

  Print(L"Kernel Entry Point: 0x%llx\n", file->e_entry);

  Print(L"program Header offset: %d\n", file->e_phoff);
  Print(L"Program Header Size: %d\n", file->e_phentsize);
  Print(L"Program Header number: %d\n", file->e_phnum);
  Print(L"== FILE PH START: 0x%llx ==\n", (uint64_t)file + file->e_phoff);

  for(uint64_t i = 0; i < file->e_phnum; i++)
  {
    Print(L"====== PH %d\n", (uint64_t)(i+1));
    PH* ph = (PH*)(((uint64_t)file) + file->e_phoff+i*file->e_phentsize);
    Print(L"TOTAL LOAD offset: 0x%llx\n", (uint64_t)(file) + ph->p_offset);
    Print(L"UNCAST LOAD offset: 0x%llx\n", (uint64_t)(file + ph->p_offset));
    Print(L"LOAD offset: 0x%llx\n", ph->p_offset);
    Print(L"LOAD vaddr: 0x%llx\n", ph->p_vaddr);
    Print(L"LOAD paddr: 0x%llx\n", ph->p_paddr);
    Print(L"LOAD size in file: 0x%llx\n", ph->p_filesz);
    Print(L"LOAD size in memory: 0x%llx\n", ph->p_memsz);
  }
}

Loaded_ELF* EFIAPI LoadELF(ELF* file)
{
  Print(L"Loading ELF file into memory\n");
  Loaded_ELF* loaded_elf;

  uefi_call_wrapper(BS->AllocatePool, 3, EfiLoaderData, sizeof(Loaded_ELF), (void**)&loaded_elf);

  loaded_elf->ph_num = file->e_phnum;
  loaded_elf->entry = file->e_entry;
  uefi_call_wrapper(BS->AllocatePool, 3, EfiLoaderData, file->e_phnum*sizeof(Loaded_PH), (void**)&(loaded_elf->ph));

  Print(L"Loading Program Segments...\n");
  for(uint32_t i = 0; i < file->e_phnum; i++)
  {
    PH* ph = (PH*)(((uint64_t)file) + file->e_phoff+i*file->e_phentsize);
    loaded_elf->ph[i].mem_size = ph->p_memsz;
    loaded_elf->ph[i].vaddr = ph->p_vaddr;

    uint64_t pages = (ph->p_memsz/1024)/4;

    // uefi_call_wrapper(BS->AllocatePool, 3, EfiLoaderData, ph->p_memsz, (void**)&(loaded_elf->ph[i].data)); // <- wrong, need to allocate aligned pages
    UINTN ipages = elf_ceil2(pages);
    EFI_PHYSICAL_ADDRESS tmp;

    /*EFI_STATUS allocstatus = */uefi_call_wrapper(BS->AllocatePages, 4, AllocateAnyPages, MEM_KERNEL, ipages, &tmp);
    loaded_elf->ph[i].data = (void*)tmp;

    // Print(L"PH: %d Addr: 0x%llX\n", i, file + ph->p_offset);
    Print(L"LOADED PH: %d Addr: 0x%llX\n", i, loaded_elf->ph[i].data);
    uefi_call_wrapper(BS->CopyMem, 3, loaded_elf->ph[i].data, (void*)((uint64_t)(file) + ph->p_offset), ph->p_filesz);
  }
  Print(L"\nFinished Loading ELF file into memory\n\n");

  return loaded_elf;
}