// Bootloader
#include <efi.h>
#include <efilib.h>
#include <utils.h>

#include <pe_types.h>
#include "bootinfo.h"

#define bfunc(func, msg)  s = func;               \
                          if (s != EFI_SUCCESS) { \
                            berror(msg, s);       \
                          }

#define putline Print(L"%d\r\n", __LINE__)

#define berror(msg, status)  Print(msg);                               \
                             Print(L"EFI error code: %r\r\n", status); \
                             return status;

EFI_STATUS EFIAPI efi_main(EFI_HANDLE ImageHandle,
    EFI_SYSTEM_TABLE* SystemTable) {
    InitializeLib(ImageHandle, SystemTable);
    Print(L"Welcome to The Vegan OS\r\n");
    Print(L"Starting boot sequence\r\n");
    Print(L"Finding graphics modes\r\n");
    EFI_STATUS s;

//==========Get graphics modes===========
    EFI_GRAPHICS_OUTPUT_PROTOCOL* graphP;
    EFI_GRAPHICS_OUTPUT_MODE_INFORMATION* graphInfo;
    UINTN handleCount = 0;
    EFI_HANDLE* graphBuf;
    EFI_GUID graphGUID = EFI_GRAPHICS_OUTPUT_PROTOCOL_GUID;
    bfunc(SystemTable->BootServices->LocateHandleBuffer(ByProtocol, &graphGUID, NULL, &handleCount, &graphBuf), L"Failed to locate handle buffer\r\n");
    bfunc(SystemTable->BootServices->HandleProtocol(graphBuf[0], &graphGUID, (VOID **)&graphP), L"Failed to handle protocol\r\n");
    UINTN modeNum;
    UINTN sizeOfInfo;
    UINT32 bestMode = 0, hHres = 0, hVres = 0;

    for (modeNum = 0; (s = graphP->QueryMode(graphP, modeNum, &sizeOfInfo, &graphInfo)) == EFI_SUCCESS; modeNum++) {
        if (graphInfo->HorizontalResolution > hHres && graphInfo->VerticalResolution > hVres && graphInfo->PixelFormat == PixelBlueGreenRedReserved8BitPerColor){
            bestMode = modeNum;
            hHres = graphInfo->HorizontalResolution;
            hVres = graphInfo->VerticalResolution;
            break;
        }
    }
    Print(L"Found best graphics mode: %d with resolution %d x %d. Pixels per scan line: %d\r\n", bestMode, hHres, hVres, graphInfo->PixelsPerScanLine);

    EFI_FILE* kernel;

    //==========Open kernel binary==========
    {
        EFI_GUID gLoadedImageProtocol = EFI_LOADED_IMAGE_PROTOCOL_GUID;
        EFI_GUID gSimpleFileSystemProtocol = EFI_SIMPLE_FILE_SYSTEM_PROTOCOL_GUID;

        EFI_HANDLE_PROTOCOL HandleProtocol =
            SystemTable->BootServices->HandleProtocol;
        EFI_LOADED_IMAGE_PROTOCOL* LoadedImage;
        bfunc(HandleProtocol(ImageHandle, &gLoadedImageProtocol,
            (void**)&LoadedImage),
            L"Failed to load image\r\n");

        EFI_SIMPLE_FILE_SYSTEM_PROTOCOL* FileSystem;
        bfunc(HandleProtocol(LoadedImage->DeviceHandle, &gSimpleFileSystemProtocol,
            (void**)&FileSystem),
            L"Failed to load filesystem\r\n");

        EFI_FILE_PROTOCOL* Root;
        bfunc(FileSystem->OpenVolume(FileSystem, &Root),
            L"Failed to load root directory\r\n");

        bfunc(Root->Open(Root, &kernel, L"kernel.exe", EFI_FILE_MODE_READ,
            EFI_FILE_READ_ONLY),
            L"Failed to load kernel binary\r\n");

        Print(L"Kernel binary successfully found\r\n");
    }
    //==========Parse kernel binary==========. Do not do anything after so not to mess up memory

    IMAGE_DOS_HEADER dosHdr;

    UINTN size = sizeof(IMAGE_DOS_HEADER);
    kernel->Read(kernel, &size, &dosHdr);

    if (dosHdr.e_magic != 23117) {
        Print(L"Incorrect magic number: %d\r\n", dosHdr.e_magic);
    }

    kernel->SetPosition(kernel, dosHdr.e_lfanew);

    IMAGE_FILE_HEADER fileHeader;
    size = sizeof(IMAGE_FILE_HEADER);
    kernel->Read(kernel, &size, &fileHeader);
    if (fileHeader.Signature != 17744) {
        Print(L"Incorrect signature: %d\r\n", fileHeader.Signature);
        berror(L"Bad signature\r\n", -1);
    }
    if (fileHeader.Machine != IMAGE_FILE_MACHINE_X64) {
        Print(L"Kernel binary is incorrect machine type: %d", fileHeader.Machine);
        berror(L"Bad machine type\r\n", -1);
    }

    UINTN sizeOfOptHdr = (UINTN)fileHeader.SizeOfOptionalHeader;

    IMAGE_OPTIONAL_HEADER optHdr;
    kernel->Read(kernel, &sizeOfOptHdr, &optHdr);

    if (optHdr.Magic != 0x20b) {
        if (optHdr.Magic == 0x10b)
            Print(L"Kernel binary is PE32 and will not be read properly.\r\n");
        else
            Print(L"Kernel binary has corrupted or incorrect magic number: %d\r\n",
                optHdr.Magic);
        berror(L"Incorrect magic number\r\n", -1);
    }

    Print(L"Kernel executable is not faulty and is of the correct types\r\n");
    kernel->SetPosition(kernel, 0);
    kernel->Read(kernel, &optHdr.SizeOfImage, (UINT64*) optHdr.ImageBase);
    UINT64 entry = optHdr.ImageBase + optHdr.AddressOfEntryPoint;
    Print(L"Kernel loaded at %x\r\n", entry);
    Print(L"Jumping to kernel\r\n");
    SystemTable->ConOut->ClearScreen(SystemTable->ConOut);
    bfunc(graphP->SetMode(graphP, bestMode), L"Failed to set graphics mode\r\n");
    //Exit boot services and jump to kernel
    // Get memory map
    EFI_MEMORY_DESCRIPTOR map;
    UINTN mapSize, mapKey;
    UINTN descriptorSize;
    UINT32 descriptorVersion;

    EFI_GET_MEMORY_MAP GetMemoryMap = SystemTable->BootServices->GetMemoryMap;
    GetMemoryMap(&mapSize, &map, &mapKey, &descriptorSize, &descriptorVersion);

    s = SystemTable->BootServices->ExitBootServices(ImageHandle, mapKey);
    if (s != EFI_SUCCESS) {
        //TODO: Jump back to text mode and say failed
        Print(L"Failed to exit boot services. EFI: %r\r\n", s);
        Pause();
    }
    else {
        UINT32* at = (UINT32*)graphP->Mode->FrameBufferBase;
        UINT32 col = 0x00ff99ff;
        for (unsigned i = 0; i < 500000; i++)
            *at++ = col;
        BOOTINFO bootinfo = { SystemTable->RuntimeServices, &map, graphInfo, graphP->Mode->FrameBufferBase };
        typedef VOID(__stdcall* kEntryPoint) (BOOTINFO* b);
        kEntryPoint kernel_main = (kEntryPoint)entry;
        kernel_main(&bootinfo);
    }

    return EFI_SUCCESS;
}