/*-
 * Copyright (c) 1998 Robert Nordier
 * All rights reserved.
 * Copyright (c) 2001 Robert Drehmel
 * All rights reserved.
 * Copyright (c) 2014 Nathan Whitehorn
 * All rights reserved.
 * Copyright (c) 2015 Eric McCorkle
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms are freely
 * permitted provided that the above copyright notice and this
 * paragraph and the following disclaimer are duplicated in all
 * such forms.
 *
 * This software is provided "AS IS" and without any express or
 * implied warranties, including, without limitation, the implied
 * warranties of merchantability and fitness for a particular
 * purpose.
 */

#include <sys/param.h>
#include <sys/stdarg.h>

#include <machine/elf.h>

#include <stand.h>

#include <efi.h>
#include <eficonsctl.h>
#include <efichar.h>

#include "boot_module.h"
#include "paths.h"
#include "proto.h"

static void efi_panic(EFI_STATUS s, const char *fmt, ...) __dead2 __printflike(2, 3);

const boot_module_t *boot_modules[] =
{
#ifdef EFI_ZFS_BOOT
        &zfs_module,
#endif
#ifdef EFI_UFS_BOOT
        &ufs_module
#endif
};
const UINTN num_boot_modules = nitems(boot_modules);

static EFI_GUID BlockIoProtocolGUID = BLOCK_IO_PROTOCOL;
static EFI_GUID DevicePathGUID = DEVICE_PATH_PROTOCOL;
static EFI_GUID LoadedImageGUID = LOADED_IMAGE_PROTOCOL;
static EFI_GUID ConsoleControlGUID = EFI_CONSOLE_CONTROL_PROTOCOL_GUID;

static EFI_PHYSICAL_ADDRESS heap;
static UINTN heapsize;

/*
 * try_boot only returns if it fails to load the loader. If it succeeds
 * it simply boots, otherwise it returns the status of last EFI call.
 */
EFI_STATUS
try_boot(const boot_module_t *mod, dev_info_t *dev, void *loaderbuf, size_t loadersize)
{
        size_t bufsize, cmdsize;
        void *buf;
        char *cmd;
        EFI_HANDLE loaderhandle;
        EFI_LOADED_IMAGE *loaded_image;
        EFI_STATUS status;

        /*
         * Read in and parse the command line from /boot.config or /boot/config,
         * if present. We'll pass it the next stage via a simple ASCII
         * string. loader.efi has a hack for ASCII strings, so we'll use that to
         * keep the size down here. We only try to read the alternate file if
         * we get EFI_NOT_FOUND because all other errors mean that the boot_module
         * had troubles with the filesystem. We could return early, but we'll let
         * loading the actual kernel sort all that out. Since these files are
         * optional, we don't report errors in trying to read them.
         */
        cmd = NULL;
        cmdsize = 0;
        status = mod->load(PATH_DOTCONFIG, dev, &buf, &bufsize);
        if (status == EFI_NOT_FOUND)
                status = mod->load(PATH_CONFIG, dev, &buf, &bufsize);
        if (status == EFI_SUCCESS) {
                cmdsize = bufsize + 1;
                cmd = malloc(cmdsize);
                if (cmd == NULL)
                        goto errout;
                memcpy(cmd, buf, bufsize);
                cmd[bufsize] = '\0';
                free(buf);
                buf = NULL;
        }

        /*
         * See if there's any env variables the module wants to set. If so,
         * append it to any config present.
         */
        if (mod->extra_env != NULL) {
                const char *env = mod->extra_env();
                if (env != NULL) {
                        size_t newlen = cmdsize + strlen(env) + 1;

                        cmd = realloc(cmd, newlen);
                        if (cmd == NULL)
                                goto errout;
                        if (cmdsize > 0)
                                strlcat(cmd, " ", newlen);
                        strlcat(cmd, env, newlen);
                        cmdsize = strlen(cmd);
                        free(__DECONST(char *, env));
                }
        }

        if ((status = BS->LoadImage(TRUE, IH, efi_devpath_last_node(dev->devpath),
            loaderbuf, loadersize, &loaderhandle)) != EFI_SUCCESS) {
                printf("Failed to load image provided by %s, size: %zu, (%lu)\n",
                     mod->name, loadersize, DECODE_ERROR(status));
                goto errout;
        }

        status = OpenProtocolByHandle(loaderhandle, &LoadedImageGUID,
            (void **)&loaded_image);
        if (status != EFI_SUCCESS) {
                printf("Failed to query LoadedImage provided by %s (%lu)\n",
                    mod->name, DECODE_ERROR(status));
                goto errout;
        }

        if (cmd != NULL)
                printf("    command args: %s\n", cmd);

        loaded_image->DeviceHandle = dev->devhandle;
        loaded_image->LoadOptionsSize = cmdsize;
        loaded_image->LoadOptions = cmd;

        DPRINTF("Starting '%s' in 5 seconds...", PATH_LOADER_EFI);
        DSTALL(1000000);
        DPRINTF(".");
        DSTALL(1000000);
        DPRINTF(".");
        DSTALL(1000000);
        DPRINTF(".");
        DSTALL(1000000);
        DPRINTF(".");
        DSTALL(1000000);
        DPRINTF(".\n");

        if ((status = BS->StartImage(loaderhandle, NULL, NULL)) !=
            EFI_SUCCESS) {
                printf("Failed to start image provided by %s (%lu)\n",
                    mod->name, DECODE_ERROR(status));
                loaded_image->LoadOptionsSize = 0;
                loaded_image->LoadOptions = NULL;
        }

errout:
        if (cmd != NULL)
                free(cmd);
        if (buf != NULL)
                free(buf);
        if (loaderbuf != NULL)
                free(loaderbuf);

        return (status);
}

EFI_STATUS
efi_main(EFI_HANDLE Ximage, EFI_SYSTEM_TABLE *Xsystab)
{
        EFI_HANDLE *handles;
        EFI_LOADED_IMAGE *img;
        EFI_DEVICE_PATH *imgpath;
        EFI_STATUS status;
        EFI_CONSOLE_CONTROL_PROTOCOL *ConsoleControl = NULL;
        SIMPLE_TEXT_OUTPUT_INTERFACE *conout = NULL;
        UINTN i, hsize, nhandles;
        CHAR16 *text;

        /* Basic initialization*/
        ST = Xsystab;
        IH = Ximage;
        BS = ST->BootServices;
        RS = ST->RuntimeServices;

        heapsize = 64 * 1024 * 1024;
        status = BS->AllocatePages(AllocateAnyPages, EfiLoaderData,
            EFI_SIZE_TO_PAGES(heapsize), &heap);
        if (status != EFI_SUCCESS) {
                ST->ConOut->OutputString(ST->ConOut,
                    __DECONST(CHAR16 *,
                    L"Failed to allocate memory for heap.\r\n"));
                BS->Exit(IH, status, 0, NULL);
        }

        setheap((void *)(uintptr_t)heap, (void *)(uintptr_t)(heap + heapsize));

        /* Set up the console, so printf works. */
        status = BS->LocateProtocol(&ConsoleControlGUID, NULL,
            (VOID **)&ConsoleControl);
        if (status == EFI_SUCCESS)
                (void)ConsoleControl->SetMode(ConsoleControl,
                    EfiConsoleControlScreenText);
        /*
         * Reset the console enable the cursor. Later we'll choose a better
         * console size through GOP/UGA.
         */
        conout = ST->ConOut;
        conout->Reset(conout, TRUE);
        /* Explicitly set conout to mode 0, 80x25 */
        conout->SetMode(conout, 0);
        conout->EnableCursor(conout, TRUE);
        conout->ClearScreen(conout);

        printf("\n>> FreeBSD EFI boot block\n");
        printf("   Loader path: %s\n\n", PATH_LOADER_EFI);
        printf("   Initializing modules:");
        for (i = 0; i < num_boot_modules; i++) {
                printf(" %s", boot_modules[i]->name);
                if (boot_modules[i]->init != NULL)
                        boot_modules[i]->init();
        }
        putchar('\n');

        /* Fetch all the block I/O handles, we have to search through them later */
        hsize = 0;
        BS->LocateHandle(ByProtocol, &BlockIoProtocolGUID, NULL,
            &hsize, NULL);
        handles = malloc(hsize);
        if (handles == NULL)
                efi_panic(EFI_OUT_OF_RESOURCES, "Failed to allocate %d handles\n",
                    hsize);
        status = BS->LocateHandle(ByProtocol, &BlockIoProtocolGUID,
            NULL, &hsize, handles);
        if (status != EFI_SUCCESS)
                efi_panic(status, "Failed to get device handles\n");
        nhandles = hsize / sizeof(*handles);

        /* Determine the devpath of our image so we can prefer it. */
        status = OpenProtocolByHandle(IH, &LoadedImageGUID, (void **)&img);
        imgpath = NULL;
        if (status == EFI_SUCCESS) {
                text = efi_devpath_name(img->FilePath);
                if (text != NULL) {
                        printf("   Load Path: %S\n", text);
                        efi_setenv_freebsd_wcs("Boot1Path", text);
                        efi_free_devpath_name(text);
                }

                status = OpenProtocolByHandle(img->DeviceHandle,
                    &DevicePathGUID, (void **)&imgpath);
                if (status != EFI_SUCCESS) {
                        DPRINTF("Failed to get image DevicePath (%lu)\n",
                            DECODE_ERROR(status));
                } else {
                        text = efi_devpath_name(imgpath);
                        if (text != NULL) {
                                printf("   Load Device: %S\n", text);
                                efi_setenv_freebsd_wcs("Boot1Dev", text);
                                efi_free_devpath_name(text);
                        }
                }
        }

        choice_protocol(handles, nhandles, imgpath);

        /* If we get here, we're out of luck... */
        efi_panic(EFI_LOAD_ERROR, "No bootable partitions found!");
}

/*
 * add_device adds a device to the passed devinfo list.
 */
void
add_device(dev_info_t **devinfop, dev_info_t *devinfo)
{
        dev_info_t *dev;

        if (*devinfop == NULL) {
                *devinfop = devinfo;
                return;
        }

        for (dev = *devinfop; dev->next != NULL; dev = dev->next)
                ;

        dev->next = devinfo;
}

void
efi_exit(EFI_STATUS s)
{

        BS->FreePages(heap, EFI_SIZE_TO_PAGES(heapsize));
        BS->Exit(IH, s, 0, NULL);
        __unreachable();
}

void
exit(int error)
{
        efi_exit(errno_to_efi_status(error));
        __unreachable();
}

/*
 * OK. We totally give up. Exit back to EFI with a sensible status so
 * it can try the next option on the list.
 */
static void
efi_panic(EFI_STATUS s, const char *fmt, ...)
{
        va_list ap;

        printf("panic: ");
        va_start(ap, fmt);
        vprintf(fmt, ap);
        va_end(ap);
        printf("\n");

        efi_exit(s);
        __unreachable();
}

int getchar(void)
{
        return (-1);
}

void
putchar(int c)
{
        CHAR16 buf[2];

        if (c == '\n') {
                buf[0] = '\r';
                buf[1] = 0;
                ST->ConOut->OutputString(ST->ConOut, buf);
        }
        buf[0] = c;
        buf[1] = 0;
        ST->ConOut->OutputString(ST->ConOut, buf);
}