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

typedef struct {
    UINT8  e_ident[16];
    UINT16 e_type;
    UINT16 e_machine;
    UINT32 e_version;
    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;

#define PT_LOAD 1
#define ELFMAG  0x464C457F

#define PAGE_SIZE_4K 4096

static void load_kernel(EFI_HANDLE ImageHandle, EFI_SYSTEM_TABLE *SystemTable) {
    EFI_STATUS status;
    EFI_GUID loadedImageGuid = EFI_LOADED_IMAGE_PROTOCOL_GUID;
    EFI_GUID fsGuid = EFI_SIMPLE_FILE_SYSTEM_PROTOCOL_GUID;
    EFI_LOADED_IMAGE_PROTOCOL *loadedImage = NULL;
    EFI_SIMPLE_FILE_SYSTEM_PROTOCOL *fs = NULL;
    EFI_FILE_PROTOCOL *root = NULL;
    EFI_FILE_PROTOCOL *kernelFile = NULL;
    VOID *phdrBuf = NULL;
    Elf64_Phdr *phdrs = NULL;
    Elf64_Ehdr ehdr;
    UINTN readSize;
    UINT16 i;

    status = uefi_call_wrapper(BS->OpenProtocol, 6,
        ImageHandle, &loadedImageGuid, (void**)&loadedImage,
        ImageHandle, NULL, EFI_OPEN_PROTOCOL_GET_PROTOCOL);
    if (EFI_ERROR(status)) {
        uefi_call_wrapper(SystemTable->ConOut->OutputString, 2, SystemTable->ConOut,
            L"ERR: loaded_image proto\n");
        goto halt;
    }

    status = uefi_call_wrapper(BS->OpenProtocol, 6,
        loadedImage->DeviceHandle, &fsGuid, (void**)&fs,
        ImageHandle, NULL, EFI_OPEN_PROTOCOL_GET_PROTOCOL);
    if (EFI_ERROR(status)) {
        uefi_call_wrapper(SystemTable->ConOut->OutputString, 2, SystemTable->ConOut,
            L"ERR: file system proto\n");
        goto halt;
    }

    status = uefi_call_wrapper(fs->OpenVolume, 1, fs, &root);
    if (EFI_ERROR(status)) {
        uefi_call_wrapper(SystemTable->ConOut->OutputString, 2, SystemTable->ConOut,
            L"ERR: open volume\n");
        goto halt;
    }

    status = uefi_call_wrapper(root->Open, 5, root, &kernelFile, L"Kernel.elf",
        EFI_FILE_MODE_READ, 0);
    if (EFI_ERROR(status)) {
        uefi_call_wrapper(SystemTable->ConOut->OutputString, 2, SystemTable->ConOut,
            L"ERR: open Kernel.elf\n");
        goto close_root;
    }

    readSize = sizeof(ehdr);
    status = uefi_call_wrapper(kernelFile->Read, 3, kernelFile, &readSize, &ehdr);
    if (EFI_ERROR(status) || readSize != sizeof(ehdr)) {
        uefi_call_wrapper(SystemTable->ConOut->OutputString, 2, SystemTable->ConOut,
            L"ERR: read ELF hdr\n");
        goto close_all;
    }

    if (*(UINT32*)ehdr.e_ident != ELFMAG) {
        uefi_call_wrapper(SystemTable->ConOut->OutputString, 2, SystemTable->ConOut,
            L"ERR: bad ELF magic\n");
        goto close_all;
    }

    {
        UINTN phdrSize = ehdr.e_phnum * ehdr.e_phentsize;
        phdrBuf = AllocatePool(phdrSize);
        if (!phdrBuf) {
            uefi_call_wrapper(SystemTable->ConOut->OutputString, 2, SystemTable->ConOut,
                L"ERR: alloc PHDRs\n");
            goto close_all;
        }

        uefi_call_wrapper(kernelFile->SetPosition, 2, kernelFile, ehdr.e_phoff);
        readSize = phdrSize;
        status = uefi_call_wrapper(kernelFile->Read, 3, kernelFile, &readSize, phdrBuf);
        if (EFI_ERROR(status) || readSize != phdrSize) {
            uefi_call_wrapper(SystemTable->ConOut->OutputString, 2, SystemTable->ConOut,
                L"ERR: read PHDRs\n");
            goto free_phdr;
        }
        phdrs = (Elf64_Phdr *)phdrBuf;
    }

    {
        UINT64 minAddr = ~0ULL;
        UINT64 maxAddr = 0;

        for (i = 0; i < ehdr.e_phnum; i++) {
            if (phdrs[i].p_type == PT_LOAD) {
                UINT64 end = phdrs[i].p_vaddr + phdrs[i].p_memsz;
                if (phdrs[i].p_vaddr < minAddr) minAddr = phdrs[i].p_vaddr;
                if (end > maxAddr) maxAddr = end;
            }
        }

        UINT64 base = minAddr & ~(PAGE_SIZE_4K - 1);
        UINT64 top = (maxAddr + PAGE_SIZE_4K - 1) & ~(PAGE_SIZE_4K - 1);
        UINTN numPages = (UINTN)((top - base) / PAGE_SIZE_4K);
        EFI_PHYSICAL_ADDRESS allocAddr = (EFI_PHYSICAL_ADDRESS)base;

        status = uefi_call_wrapper(BS->AllocatePages, 4,
            AllocateAddress, EfiLoaderData, numPages, &allocAddr);
        if (EFI_ERROR(status)) {
            uefi_call_wrapper(SystemTable->ConOut->OutputString, 2, SystemTable->ConOut,
                L"ERR: alloc kernel mem\n");
            goto free_phdr;
        }
    }

    for (i = 0; i < ehdr.e_phnum; i++) {
        if (phdrs[i].p_type != PT_LOAD) continue;

        uefi_call_wrapper(kernelFile->SetPosition, 2, kernelFile, phdrs[i].p_offset);
        readSize = (UINTN)phdrs[i].p_filesz;
        uefi_call_wrapper(kernelFile->Read, 3, kernelFile, &readSize,
            (VOID*)(UINTN)phdrs[i].p_vaddr);

        if (phdrs[i].p_memsz > phdrs[i].p_filesz) {
            UINT8 *bss = (UINT8*)(UINTN)(phdrs[i].p_vaddr + phdrs[i].p_filesz);
            UINTN bssSize = (UINTN)(phdrs[i].p_memsz - phdrs[i].p_filesz);
            for (UINTN j = 0; j < bssSize; j++) bss[j] = 0;
        }
    }

    FreePool(phdrBuf);
    phdrBuf = NULL;
    uefi_call_wrapper(kernelFile->Close, 1, kernelFile);
    kernelFile = NULL;
    uefi_call_wrapper(root->Close, 1, root);
    root = NULL;

    uefi_call_wrapper(SystemTable->ConOut->OutputString, 2, SystemTable->ConOut,
        L"* Jumping to kernel...\n");

    {
        typedef void (*KernelEntry)(EFI_HANDLE, EFI_SYSTEM_TABLE *);
        KernelEntry entry = (KernelEntry)(UINTN)ehdr.e_entry;
        entry(ImageHandle, SystemTable);
    }

    uefi_call_wrapper(SystemTable->ConOut->OutputString, 2, SystemTable->ConOut,
        L"ERR: kernel returned\n");
    goto halt;

free_phdr:
    if (phdrBuf) FreePool(phdrBuf);
close_all:
    if (kernelFile) uefi_call_wrapper(kernelFile->Close, 1, kernelFile);
close_root:
    if (root) uefi_call_wrapper(root->Close, 1, root);
halt:
    while (1) __asm__ volatile("hlt");
}

EFI_STATUS EFIAPI efi_main(EFI_HANDLE ImageHandle, EFI_SYSTEM_TABLE *SystemTable) {
    InitializeLib(ImageHandle, SystemTable);

    uefi_call_wrapper(SystemTable->ConOut->ClearScreen, 1, SystemTable->ConOut);
    uefi_call_wrapper(SystemTable->ConOut->SetCursorPosition, 3, SystemTable->ConOut, 0, 1);
    uefi_call_wrapper(SystemTable->ConOut->OutputString, 2, SystemTable->ConOut, L"Welcome to SYLVA OS!\n");
    uefi_call_wrapper(SystemTable->ConOut->SetCursorPosition, 3, SystemTable->ConOut, 0, 2);
    uefi_call_wrapper(SystemTable->ConOut->OutputString, 2, SystemTable->ConOut, L"* The system has booted successfully.\n");

    uefi_call_wrapper(SystemTable->ConOut->SetCursorPosition, 3, SystemTable->ConOut, 0, 3);
    uefi_call_wrapper(SystemTable->ConOut->OutputString, 2, SystemTable->ConOut, L"* Loading kernel...\n");

    load_kernel(ImageHandle, SystemTable);

    return EFI_SUCCESS;
}