/**
 * PROJECT:         ExectOS
 * COPYRIGHT:       See COPYING.md in the top level directory
 * FILE:            xtldr/xtldr.c
 * DESCRIPTION:     XTOS UEFI Boot Loader
 * DEVELOPERS:      Rafal Kupiec <belliash@codingworkshop.eu.org>
 */

#include <xtldr.h>


XTCDECL
EFI_STATUS
BlLoadModule(IN PWCHAR Module);

XTCDECL
EFI_STATUS
BlLoadModules(IN PWCHAR ModulesList);

/**
 * Initializes a list of operating systems for XTLDR boot menu.
 *
 * @param MenuEntries
 *        Supplies a pointer to memory area where operating systems list will be stored.
 *
 * @param EntriesCount
 *        Supplies a pointer to memory area where number of menu entries will be stored.
 *
 * @param DefaultId
 *        Supplies a pointer to memory area where ID of default menu entry will be stored.
 *
 * @return This routine does not return any value.
 *
 * @since XT 1.0
 */
XTCDECL
VOID
BlInitializeBootMenuList(OUT PXTBL_BOOTMENU_ITEM MenuEntries,
                         OUT PULONG EntriesCount,
                         OUT PULONG DefaultId)
{
    PWCHAR DefaultMenuEntry, MenuEntryName;
    PLIST_ENTRY MenuEntrySectionList, MenuEntryList;
    PXTBL_CONFIG_SECTION MenuEntrySection;
    PXTBL_CONFIG_ENTRY MenuEntryOption;
    PXTBL_BOOTMENU_ITEM OsList;
    ULONG DefaultOS, NumberOfEntries;

    /* Set default values */
    DefaultOS = 0;
    NumberOfEntries = 0;
    OsList = NULL;

    /* Get default menu entry from configuration */
    DefaultMenuEntry = BlGetConfigValue(L"DEFAULT");

    /* Iterate through all menu sections */
    MenuEntrySectionList = BlpMenuList->Flink;
    while(MenuEntrySectionList != BlpMenuList)
    {
        /* NULLify menu entry name */
        MenuEntryName = NULL;

        /* Get menu section */
        MenuEntrySection = CONTAIN_RECORD(MenuEntrySectionList, XTBL_CONFIG_SECTION, Flink);

        /* Check if this is the default menu entry */
        if(RtlCompareWideStringInsensitive(MenuEntrySection->SectionName, DefaultMenuEntry, 0) == 0)
        {
            /* Set default OS ID */
            DefaultOS = NumberOfEntries;
        }

        /* Iterate through all entry parameters */
        MenuEntryList = MenuEntrySection->Options.Flink;
        while(MenuEntryList != &MenuEntrySection->Options)
        {
            /* Get menu entry parameter */
            MenuEntryOption = CONTAIN_RECORD(MenuEntryList, XTBL_CONFIG_ENTRY, Flink);

            /* Check if this is the menu entry display name */
            if(RtlCompareWideStringInsensitive(MenuEntryOption->Name, L"SYSTEMNAME", 0) == 0)
            {
                /* Set menu entry display name */
                MenuEntryName = MenuEntryOption->Value;
            }

            /* Get next parameter for this menu entry */
            MenuEntryList = MenuEntryList->Flink;
        }

        /* Add OS to the boot menu list */
        OsList[NumberOfEntries].EntryName = MenuEntryName;
        OsList[NumberOfEntries].Options = &MenuEntrySection->Options;

        /* Get next menu entry */
        MenuEntrySectionList = MenuEntrySectionList->Flink;
        NumberOfEntries++;
    }

    /* Set return values */
    *DefaultId = DefaultOS;
    *EntriesCount = NumberOfEntries;
    MenuEntries = OsList;
}

/**
 * Loads all necessary modules and invokes boot protocol.
 *
 * @param OptionsList
 *        Supplies a pointer to list of options associated with chosen boot menu entry.
 *
 * @return This routine returns a status code.
 *
 * @since XT 1.0
 */
XTCDECL
EFI_STATUS
BlInvokeBootProtocol(IN PLIST_ENTRY OptionsList)
{
    PWCHAR ModulesList, ProtocolName;
    PLIST_ENTRY OptionsListEntry;
    PXTBL_CONFIG_ENTRY Option;
    EFI_STATUS Status;

    /* Set default values */
    ModulesList = NULL;
    ProtocolName = NULL;

    /* Iterate through all options provided by boot menu entry */
    OptionsListEntry = OptionsList->Flink;
    while(OptionsListEntry != OptionsList)
    {
        /* Get option */
        Option = CONTAIN_RECORD(OptionsListEntry, XTBL_CONFIG_ENTRY, Flink);

        /* Look for boot protocol and modules list */
        if(RtlCompareWideStringInsensitive(Option->Name, L"SYSTEMTYPE", 0) == 0)
        {
            /* Boot protocol found */
            ProtocolName = Option->Value;
        }
        else if(RtlCompareWideStringInsensitive(Option->Name, L"BOOTMODULES", 0) == 0)
        {
            /* Set protocol name */
            ModulesList = Option->Value;
        }

        /* Move to the next option entry */
        OptionsListEntry = OptionsListEntry->Flink;
    }

    /* Load all necessary modules */
    Status = BlLoadModules(ModulesList);
    if(Status != STATUS_EFI_SUCCESS)
    {
        /* Failed to load modules, print error message and return status code */
        BlDebugPrint(L"Failed to load XTLDR modules\n");
        return Status;
    }
    for(;;);

    /* This point should never be reached */
    return STATUS_EFI_SUCCESS;
}

XTCDECL
EFI_STATUS
BlLoadModules(IN PWCHAR ModulesList)
{
    PWCHAR LastModule, Module;
    EFI_STATUS Status;

    if(ModulesList != NULL)
    {
        /* Tokenize provided list of modules */
        Module = RtlTokenizeWideString(ModulesList, L" ", &LastModule);

        /* Iterate over all arguments passed to boot loader */
        while(Module != NULL)
        {
            Status = BlLoadModule(Module);
            if(Status != STATUS_EFI_SUCCESS)
            {
                /* Failed to load module, print error message and return status code */
                BlDebugPrint(L"Failed to load module '%S', status = 0x%lx\n", Module, Status);
                // return Status;
            }

            /* Take next module from the list */
            Module = RtlTokenizeWideString(NULL, L" ", &LastModule);
        }
    }

    /* Return success */
    return STATUS_EFI_SUCCESS;
}

XTCDECL
EFI_STATUS
BlLoadModule(IN PWCHAR ModuleName)
{
    EFI_MEMMAP_DEVICE_PATH ModuleDevicePath[2];
    PEFI_FILE_HANDLE DirHandle, FsHandle;
    EFI_HANDLE DiskHandle, ModuleHandle;
    PPECOFF_IMAGE_SECTION_HEADER SectionHeader;
    PPECOFF_IMAGE_DOS_HEADER DosHeader;
    PPECOFF_IMAGE_PE_HEADER PeHeader;
    PLIST_ENTRY ModuleListEntry;
    WCHAR ModuleFileName[1024];
    USHORT SectionIndex;
    SIZE_T ModuleSize;
    EFI_STATUS Status;
    PVOID ModuleData;
    PWCHAR DepsData;
    PXTBL_MODULE_INFO ModuleInfo;

    ModuleListEntry = BlpLoadedModules.Flink;
    while(ModuleListEntry != &BlpLoadedModules)
    {
        /* Get module information */
        ModuleInfo = CONTAIN_RECORD(ModuleListEntry, XTBL_MODULE_INFO, Flink);

        if(RtlCompareWideStringInsensitive(ModuleInfo->ModuleName, ModuleName, 0) == 0)
        {
            /* Module already loaded */
            BlDebugPrint(L"Module '%S' already loaded\n", ModuleName);
            return STATUS_EFI_SUCCESS;
        }

        /* Move to the module */
        ModuleListEntry = ModuleListEntry->Flink;
    }

    /* Print debug message */
    BlDebugPrint(L"Loading module '%S' ...\n", ModuleName);

    /* Set module path */
    RtlCopyMemory(ModuleFileName, ModuleName, sizeof(ModuleFileName) / sizeof(WCHAR));
    RtlConcatenateWideString(ModuleFileName, L".EFI", 0);

    /* Open EFI volume */
    Status = BlOpenVolume(NULL, &DiskHandle, &FsHandle);
    if(Status != STATUS_EFI_SUCCESS)
    {
        /* Failed to open a volume */
        return Status;
    }

    /* Open XTLDR modules directory and close the FS immediately */
    Status = FsHandle->Open(FsHandle, &DirHandle, XTBL_MODULES_DIRECTORY_PATH, EFI_FILE_MODE_READ, 0);
    FsHandle->Close(FsHandle);

    /* Check if directory opened successfully */
    if(Status != STATUS_EFI_SUCCESS)
    {
        /* Failed to open directory */
        BlCloseVolume(DiskHandle);
        return Status;
    }

    /* Read module file from disk and close EFI volume */
    Status = BlReadFile(DirHandle, ModuleFileName, &ModuleData, &ModuleSize);
    BlCloseVolume(DiskHandle);

    /* Make sure module file was read successfully */
    if(Status != STATUS_EFI_SUCCESS)
    {
        /* Failed to read file */
        return Status;
    }

    /* Allocate memory for new option */
    Status = BlMemoryAllocatePool(sizeof(XTBL_MODULE_INFO), (PVOID*)&ModuleInfo);
    if(Status != STATUS_EFI_SUCCESS)
    {
        /* Failed to allocate memory */
        return Status;
    }

    RtlZeroMemory(ModuleInfo, sizeof(XTBL_MODULE_INFO));

    /* Setup PE/COFF EFI image headers */
    DosHeader = (PPECOFF_IMAGE_DOS_HEADER)ModuleData;
    PeHeader = (PPECOFF_IMAGE_PE_HEADER)(ModuleData + DosHeader->e_lfanew);
    SectionHeader = (PPECOFF_IMAGE_SECTION_HEADER)((PUCHAR)&PeHeader->OptionalHeader +
                                                   PeHeader->FileHeader.SizeOfOptionalHeader);

    /* Look for .moddeps section */
    for(SectionIndex = 0; SectionIndex < PeHeader->FileHeader.NumberOfSections; SectionIndex++)
    {
        /* Check section name */
        if(RtlCompareString((PCHAR)SectionHeader[SectionIndex].Name, ".moddeps", 8) == 0)
        {
            /* Store address of .moddeps data segment */
            DepsData = ModuleData + SectionHeader[SectionIndex].PointerToRawData;

            /* Iterate over all dependencies stored */
            while(*DepsData != 0)
            {
                /* Load dependency module */
                BlDebugPrint(L"Module '%S' requires '%S' ...\n", ModuleName, DepsData);
                Status = BlLoadModule(DepsData);

                if(Status != STATUS_EFI_SUCCESS)
                {
                    /* Failed to load module, print error message and return status code */
                    BlDebugPrint(L"Failed to load dependency module '%S', status = 0x%lx\n", DepsData, Status);
                    return STATUS_EFI_UNSUPPORTED;
                }

                /* Add dependency module name to the list */
                RtlConcatenateWideString(ModuleInfo->Dependencies, DepsData, RtlWideStringLength(DepsData, 0));
                RtlConcatenateWideString(ModuleInfo->Dependencies, L" ", 1);

                /* Get next dependency module name */
                DepsData += 8;
            }
        }
        else if(RtlCompareString((PCHAR)SectionHeader[SectionIndex].Name, ".modinfo", 8) == 0)
        {
            /* Store module description */
            ModuleInfo->ModuleDescription = ModuleData + SectionHeader[SectionIndex].PointerToRawData;
        }
    }

    /* Store module name */
    ModuleInfo->ModuleName = ModuleName;

    /* Setup module device path */
    ModuleDevicePath[0].Header.Length[0] = sizeof(EFI_MEMMAP_DEVICE_PATH);
    ModuleDevicePath[0].Header.Length[1] = sizeof(EFI_MEMMAP_DEVICE_PATH) >> 8;
    ModuleDevicePath[0].Header.Type = EFI_HARDWARE_DEVICE_PATH;
    ModuleDevicePath[0].Header.SubType = EFI_HARDWARE_MEMMAP_DP;
    ModuleDevicePath[0].MemoryType = EfiLoaderData;
    ModuleDevicePath[0].StartingAddress = (UINT_PTR)ModuleData;
    ModuleDevicePath[0].EndingAddress = (UINT_PTR)ModuleData + ModuleSize;
    ModuleDevicePath[1].Header.Length[0] = sizeof(EFI_DEVICE_PATH_PROTOCOL);
    ModuleDevicePath[1].Header.Length[1] = sizeof(EFI_DEVICE_PATH_PROTOCOL) >> 8;
    ModuleDevicePath[1].Header.Type = EFI_END_DEVICE_PATH;
    ModuleDevicePath[1].Header.SubType = EFI_END_ENTIRE_DP;

    /* Load EFI image */
    Status = EfiSystemTable->BootServices->LoadImage(FALSE, EfiImageHandle, (PEFI_DEVICE_PATH_PROTOCOL)ModuleDevicePath,
                                                     ModuleData, ModuleSize, &ModuleHandle);
    if(Status != STATUS_EFI_SUCCESS)
    {
        /* Failed to load module image */
        return Status;
    }

    /* Start EFI image */
    Status = EfiSystemTable->BootServices->StartImage(ModuleHandle, NULL, NULL);
    if(Status != STATUS_EFI_SUCCESS)
    {
        /* Failed to start module image */
        return Status;
    }

    /* Add module to the list of loaded modules */
    RtlInsertTailList(&BlpLoadedModules, &ModuleInfo->Flink);

    /* Return success */
    return STATUS_EFI_SUCCESS;
}

/**
 * This routine is the entry point of the XT EFI boot loader.
 *
 * @param ImageHandle
 *        Firmware-allocated handle that identifies the image.
 *
 * @param SystemTable
 *        Provides the EFI system table.
 *
 * @return This routine returns status code.
 *
 * @since XT 1.0
 */
XTCDECL
EFI_STATUS
BlStartXtLoader(IN EFI_HANDLE ImageHandle,
                IN PEFI_SYSTEM_TABLE SystemTable)
{
    EFI_STATUS Status;

    /* Set the system table and image handle */
    EfiImageHandle = ImageHandle;
    EfiSystemTable = SystemTable;

    /* Initialize XTLDR and */
    BlpInitializeEfiBootLoader();

    /* Parse configuration options passed from UEFI shell */
    Status = BlpParseCommandLine();
    if(Status != STATUS_EFI_SUCCESS)
    {
        /* Failed to parse command line options */
        BlDisplayErrorDialog(L"XTLDR", L"Failed to parse command line parameters.");
    }

    /* Attempt to early initialize debug console */
    if(DEBUG)
    {
        Status = BlpInitializeDebugConsole();
        if(Status != STATUS_EFI_SUCCESS)
        {
            /* Initialization failed, notify user on stdout */
            BlDisplayErrorDialog(L"XTLDR", L"Failed to initialize debug console.");
        }
    }

    /* Load XTLDR configuration file */
    Status = BlpLoadConfiguration();
    if(Status != STATUS_EFI_SUCCESS)
    {
        /* Failed to load/parse config file */
        BlDisplayErrorDialog(L"XTLDR", L"Failed to load and parse configuration file.");
    }

    /* Reinitialize debug console if it was not initialized earlier */
    if(DEBUG)
    {
        Status = BlpInitializeDebugConsole();
        if(Status != STATUS_EFI_SUCCESS)
        {
            /* Initialization failed, notify user on stdout */
            BlDisplayErrorDialog(L"XTLDR", L"Failed to initialize debug console.");
        }
    }

    /* Disable watchdog timer */
    Status = EfiSystemTable->BootServices->SetWatchdogTimer(0, 0x10000, 0, NULL);
    if(Status != STATUS_EFI_SUCCESS)
    {
        /* Failed to disable the timer, print message */
        BlDebugPrint(L"WARNING: Failed to disable watchdog timer\n");
    }

    /* Register loader protocol */
    Status = BlpRegisterXtLoaderProtocol();
    if(Status != STATUS_EFI_SUCCESS)
    {
        /* Failed to register loader protocol */
        BlDebugPrint(L"ERROR: Failed to register XTLDR boot protocol\n");
        return Status;
    }

    /* Discover and enumerate EFI block devices */
    Status = BlEnumerateBlockDevices();
    if(Status != STATUS_EFI_SUCCESS)
    {
        /* Failed to enumerate block devices */
        BlDebugPrint(L"ERROR: Failed to discover and enumerate block devices\n");
        return Status;
    }

    BlLoadModules(BlGetConfigValue(L"MODULES"));

    BlLoadModule(L"DUMMY");


    BlConsolePrint(L"\n\n\n\n\n\n\n\nList of loaded modules:");
    PLIST_ENTRY ModuleListEntry;
    PXTBL_MODULE_INFO ModuleInfo;

    ModuleListEntry = BlpLoadedModules.Flink;
    while(ModuleListEntry != &BlpLoadedModules)
    {
        /* Get module information */
        ModuleInfo = CONTAIN_RECORD(ModuleListEntry, XTBL_MODULE_INFO, Flink);

        /* Module already loaded */
        BlConsolePrint(L"\n%.8S", ModuleInfo->ModuleName);
        if(ModuleInfo->ModuleDescription != 0)
        {
            BlConsolePrint(L" (%S)", ModuleInfo->ModuleDescription);
        }

        if(ModuleInfo->Dependencies[0] != 0)
        {
            BlConsolePrint(L"\n  - Uses: %S", ModuleInfo->Dependencies);
        }

        /* Move to the module */
        ModuleListEntry = ModuleListEntry->Flink;
    }

    BlConsolePrint(L"\n\n END OF LIST\n");

    for(;;);

    for(;;)
    {
        if(BlpStatus.BootMenu != NULL)
        {
            /* Display alternative boot menu */
            BlpStatus.BootMenu();
        }
        else
        {
            /* Display default boot menu */
            BlDisplayBootMenu();
        }
    }

    /* This point should be never reached, if this happen return error code */
    return STATUS_EFI_LOAD_ERROR;
}