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

#include "boot_info.h"

#define ELF_MAGIC 0x464c457f
#define PT_LOAD 1

typedef struct {
    UINT32 e_magic;
    UINT8 e_class;
    UINT8 e_data;
    UINT8 e_version;
    UINT8 e_osabi;
    UINT8 e_abiversion;
    UINT8 e_pad[7];
    UINT16 e_type;
    UINT16 e_machine;
    UINT32 e_version2;
    UINT64 e_entry;
    UINT64 e_phoff;
    UINT64 e_shoff;
    UINT32 e_flags;
    UINT16 e_ehsize;
    UINT16 e_phentsize;
    UINT16 e_phnum;
    UINT16 e_shentsize;
    UINT16 e_shnum;
    UINT16 e_shstrndx;
} Elf64_Ehdr;

typedef struct {
    UINT32 p_type;
    UINT32 p_flags;
    UINT64 p_offset;
    UINT64 p_vaddr;
    UINT64 p_paddr;
    UINT64 p_filesz;
    UINT64 p_memsz;
    UINT64 p_align;
} Elf64_Phdr;

static EFI_STATUS open_root_volume(EFI_HANDLE ImageHandle, EFI_FILE_PROTOCOL **Root)
{
    EFI_STATUS Status;
    EFI_LOADED_IMAGE *LoadedImage = NULL;
    EFI_SIMPLE_FILE_SYSTEM_PROTOCOL *FileSystem = NULL;

    Status = uefi_call_wrapper(BS->HandleProtocol, 3, ImageHandle,
                               &gEfiLoadedImageProtocolGuid, (VOID **)&LoadedImage);
    if (EFI_ERROR(Status)) {
        return Status;
    }

    Status = uefi_call_wrapper(BS->HandleProtocol, 3, LoadedImage->DeviceHandle,
                               &gEfiSimpleFileSystemProtocolGuid, (VOID **)&FileSystem);
    if (EFI_ERROR(Status)) {
        return Status;
    }

    return uefi_call_wrapper(FileSystem->OpenVolume, 2, FileSystem, Root);
}

static EFI_STATUS read_file_to_buffer(EFI_HANDLE ImageHandle, CHAR16 *Path,
                                      VOID **Buffer, UINTN *Size)
{
    EFI_STATUS Status;
    EFI_FILE_PROTOCOL *Root = NULL;
    EFI_FILE_PROTOCOL *File = NULL;
    EFI_FILE_INFO *Info = NULL;
    UINTN InfoSize = 0;
    UINTN FileSize = 0;

    Status = open_root_volume(ImageHandle, &Root);
    if (EFI_ERROR(Status)) {
        return Status;
    }

    Status = uefi_call_wrapper(Root->Open, 5, Root, &File, Path,
                               EFI_FILE_MODE_READ, 0);
    if (EFI_ERROR(Status)) {
        return Status;
    }

    Status = uefi_call_wrapper(File->GetInfo, 4, File, &gEfiFileInfoGuid,
                               &InfoSize, NULL);
    if (Status != EFI_BUFFER_TOO_SMALL) {
        uefi_call_wrapper(File->Close, 1, File);
        return Status;
    }

    Status = uefi_call_wrapper(BS->AllocatePool, 3, EfiLoaderData,
                               InfoSize, (VOID **)&Info);
    if (EFI_ERROR(Status)) {
        uefi_call_wrapper(File->Close, 1, File);
        return Status;
    }

    Status = uefi_call_wrapper(File->GetInfo, 4, File, &gEfiFileInfoGuid,
                               &InfoSize, Info);
    if (EFI_ERROR(Status)) {
        uefi_call_wrapper(BS->FreePool, 1, Info);
        uefi_call_wrapper(File->Close, 1, File);
        return Status;
    }

    FileSize = (UINTN)Info->FileSize;
    Status = uefi_call_wrapper(BS->AllocatePool, 3, EfiLoaderData,
                               FileSize, Buffer);
    if (EFI_ERROR(Status)) {
        uefi_call_wrapper(BS->FreePool, 1, Info);
        uefi_call_wrapper(File->Close, 1, File);
        return Status;
    }

    Status = uefi_call_wrapper(File->Read, 3, File, &FileSize, *Buffer);
    if (EFI_ERROR(Status)) {
        uefi_call_wrapper(BS->FreePool, 1, Info);
        uefi_call_wrapper(File->Close, 1, File);
        return Status;
    }

    uefi_call_wrapper(BS->FreePool, 1, Info);
    uefi_call_wrapper(File->Close, 1, File);

    *Size = FileSize;
    return EFI_SUCCESS;
}

static EFI_STATUS load_elf_kernel(VOID *Image, UINTN Size, UINT64 *EntryOut)
{
    Elf64_Ehdr *Ehdr = (Elf64_Ehdr *)Image;
    Elf64_Phdr *Phdr = NULL;
    UINT16 Index = 0;

    if (Size < sizeof(Elf64_Ehdr)) {
        return EFI_LOAD_ERROR;
    }

    if (Ehdr->e_magic != ELF_MAGIC || Ehdr->e_class != 2 || Ehdr->e_data != 1) {
        return EFI_LOAD_ERROR;
    }

    if ((Ehdr->e_phoff + (UINT64)Ehdr->e_phnum * Ehdr->e_phentsize) > Size) {
        return EFI_LOAD_ERROR;
    }

    Phdr = (Elf64_Phdr *)((UINT8 *)Image + Ehdr->e_phoff);

    for (Index = 0; Index < Ehdr->e_phnum; Index++) {
        Elf64_Phdr *Segment = (Elf64_Phdr *)((UINT8 *)Phdr + (Index * Ehdr->e_phentsize));
        UINT64 SegmentStart = 0;
        UINT64 SegmentEnd = 0;
        UINT64 SegmentPages = 0;
        EFI_PHYSICAL_ADDRESS Address = 0;

        if (Segment->p_type != PT_LOAD || Segment->p_memsz == 0) {
            continue;
        }

        if ((Segment->p_offset + Segment->p_filesz) > Size) {
            return EFI_LOAD_ERROR;
        }

        SegmentStart = Segment->p_vaddr & ~0xFFFULL;
        SegmentEnd = (Segment->p_vaddr + Segment->p_memsz + 0xFFFULL) & ~0xFFFULL;
        SegmentPages = (SegmentEnd - SegmentStart) >> 12;

        Address = (EFI_PHYSICAL_ADDRESS)SegmentStart;
        if (EFI_ERROR(uefi_call_wrapper(BS->AllocatePages, 4, AllocateAddress,
                                        EfiLoaderData, SegmentPages, &Address))) {
            return EFI_OUT_OF_RESOURCES;
        }

        CopyMem((VOID *)(UINTN)Segment->p_vaddr,
                (UINT8 *)Image + Segment->p_offset,
                (UINTN)Segment->p_filesz);
        if (Segment->p_memsz > Segment->p_filesz) {
            SetMem((VOID *)(UINTN)(Segment->p_vaddr + Segment->p_filesz),
                   (UINTN)(Segment->p_memsz - Segment->p_filesz), 0);
        }
    }

    *EntryOut = Ehdr->e_entry;
    return EFI_SUCCESS;
}

static EFI_STATUS loader_clear_screen(void)
{
    return uefi_call_wrapper(ST->ConOut->ClearScreen, 1, ST->ConOut);
}

static EFI_STATUS loader_set_attribute(UINTN Attribute)
{
    return uefi_call_wrapper(ST->ConOut->SetAttribute, 2, ST->ConOut, Attribute);
}

static EFI_STATUS loader_read_key(EFI_INPUT_KEY *Key)
{
    UINTN Index = 0;
    uefi_call_wrapper(BS->WaitForEvent, 3, 1, &ST->ConIn->WaitForKey, &Index);
    return uefi_call_wrapper(ST->ConIn->ReadKeyStroke, 2, ST->ConIn, Key);
}

static EFI_STATUS loader_try_read_key(EFI_INPUT_KEY *Key)
{
    return uefi_call_wrapper(ST->ConIn->ReadKeyStroke, 2, ST->ConIn, Key);
}

static void loader_shutdown(void)
{
    uefi_call_wrapper(RT->ResetSystem, 4, EfiResetShutdown, EFI_SUCCESS, 0, NULL);
}

static EFI_STATUS loader_alloc_pages(UINTN pages, EFI_PHYSICAL_ADDRESS *addr)
{
    return uefi_call_wrapper(BS->AllocatePages, 4,
                             AllocateAnyPages, EfiLoaderData, pages, addr);
}

static EFI_STATUS loader_free_pages(EFI_PHYSICAL_ADDRESS addr, UINTN pages)
{
    return uefi_call_wrapper(BS->FreePages, 2, addr, pages);
}

EFI_STATUS
EFIAPI
efi_main(EFI_HANDLE ImageHandle, EFI_SYSTEM_TABLE *SystemTable)
{
    EFI_STATUS Status;
    VOID *KernelImage = NULL;
    UINTN KernelSize = 0;
    UINT64 KernelEntry = 0;
    BootInfo Boot;
    KernelEntryFn EntryFn = NULL;
    
    // Initialize the GNU-EFI library
    InitializeLib(ImageHandle, SystemTable);
    
    Print(L"Loading kernel...\n\r");

    Status = read_file_to_buffer(ImageHandle, L"\\kernel.elf", &KernelImage, &KernelSize);
    if (EFI_ERROR(Status)) {
        Print(L"Failed to read kernel.elf: %r\n\r", Status);
        return Status;
    }

    Status = load_elf_kernel(KernelImage, KernelSize, &KernelEntry);
    if (EFI_ERROR(Status)) {
        Print(L"Failed to load kernel.elf: %r\n\r", Status);
        return Status;
    }

    Boot.SystemTable = ST;
    Boot.print = Print;
    Boot.clear_screen = loader_clear_screen;
    Boot.set_attribute = loader_set_attribute;
    Boot.read_key = loader_read_key;
    Boot.try_read_key = loader_try_read_key;
    Boot.shutdown = loader_shutdown;
    Boot.alloc_pages = loader_alloc_pages;
    Boot.free_pages = loader_free_pages;

    EntryFn = (KernelEntryFn)(UINTN)KernelEntry;
    EntryFn(&Boot);

    Print(L"Kernel returned. Halting.\n\r");
    return EFI_SUCCESS;
}