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

#include <elf.h>

#include "registers.h"

CHAR16 *DEFAULT_KERNEL_NAME = L"kernel";

EFI_STATUS
check_header(SIMPLE_READ_FILE file)
{
    return EFI_SUCCESS;    
}

EFI_STATUS
load_image(SIMPLE_READ_FILE file)
{
    Print(L"Start load image.\n");

    UINTN size = sizeof(Elf64_Ehdr);
    Print(L"File size is %d\n", size);
    CHAR16 *buffer = L"NULL";

    Print(L"ReadSimpleLoadKernel. size %d\n", size);
    EFI_STATUS status = ReadSimpleReadFile(file ,0 , &size, (VOID *)buffer);

    Print(L"ReadSimpleReadFile :%r\n", status);
    if (EFI_ERROR(status)) {
        Print(L"ReadSimpleReadFile :%r\n", status);
        return status;
    }

    Print(L"buffer %s\n", buffer);

    return EFI_SUCCESS;
}

UINTN
is_space(CHAR16 c)
{
    return  c == L' ' || c == L'\t' || c == L'\r' || c == L'\n';
}

CHAR16*
get_kernel_name(CHAR16 *options, UINT32 options_size)
{
    /* 
     * skip first option,
     * first option is file name of this efi executable.
     */ 

    UINT32 count = options_size;

    while(count && !is_space(*options)) {
        options++;
        count--;
    }

    while (is_space(*options)) {
        options++;
    }

    CHAR16 *start_ptr = options;

    while(!is_space(*options)) {
        options++;
    }

    CHAR16 *kernel_name = (CHAR16 *)AllocatePool((options - start_ptr + 1) * sizeof(CHAR16));

    RtCopyMem(kernel_name, start_ptr, (options - start_ptr) * sizeof(CHAR16));
    
    // delete last 0
    kernel_name[options - start_ptr] = 0;

    Print(L"kernel name: -%s- name size: %d\n", kernel_name, (options - start_ptr));

    return kernel_name;
    
}

EFI_STATUS
load_kernel(CHAR16 *kernel_name)
{
    Print(L"Start load kernel: %s\n", kernel_name);

    EFI_STATUS status = EFI_SUCCESS;
    UINTN handle_count;
    EFI_HANDLE *handle_buffer;

    status = uefi_call_wrapper(BS->LocateHandleBuffer,
                5,
                ByProtocol,
                &FileSystemProtocol,
                NULL,
                &handle_count,
                &handle_buffer);

    if (EFI_ERROR(status)) {
        Print(L"LocateHandleBuffer is %r\n", status);
        return status;
    }

    EFI_DEVICE_PATH *path = NULL;
    UINTN handle_idx = 0;
    SIMPLE_READ_FILE read_handle;

    for (handle_idx = 0; handle_idx < handle_count; handle_idx++) {
        EFI_HANDLE device_handle;

        path = FileDevicePath(handle_buffer[handle_idx], kernel_name);

        if (!path) {
            status = EFI_NOT_FOUND;
            break;
        }

        status = OpenSimpleReadFile(TRUE, NULL, 0, &path, &device_handle, &read_handle);

        if (!EFI_ERROR(status)) {
            break;
        }

        FreePool(path);
        path = NULL;
    }

    if (!EFI_ERROR(status)) {
        status = load_image(read_handle);
    }

    if (read_handle) {
        CloseSimpleReadFile(read_handle);
    }

    if (path) {
        FreePool(path);
    }

    FreePool(handle_buffer);

    return status;
}

EFI_STATUS
efi_main(EFI_HANDLE image, EFI_SYSTEM_TABLE *systab)
{
    EFI_LOADED_IMAGE *loaded_image;

    InitializeLib(image, systab);

    uefi_call_wrapper(systab->ConOut->Reset, 2, systab->ConOut, FALSE);

    Print(L"---- start ----\n");

    EFI_STATUS status;
    
    Print(L"Open LoadedImage Protocol\n");
    status = uefi_call_wrapper(
        BS->OpenProtocol,
        6,
        image,
        &LoadedImageProtocol,
        (void **)&loaded_image,
        image,
        NULL,
        EFI_OPEN_PROTOCOL_BY_HANDLE_PROTOCOL);

    if (EFI_ERROR(status)) {
        Print(L"OpenProtocol is %r\n", status);
    }

    Print(L"%s\n", loaded_image->LoadOptions);

    CHAR16 *kernel_name = get_kernel_name(loaded_image->LoadOptions, loaded_image->LoadOptionsSize);

    Print(L"%s\n", kernel_name);

    load_kernel(kernel_name);

    Print(L"Close LoadedImage Protocol\n");
    status = uefi_call_wrapper(
        BS->CloseProtocol,
        4,
        image,
        &LoadedImageProtocol,
        image,
        NULL);

    if (EFI_ERROR(status)) {
        Print(L"CloseProtocol is %r\n", status);
    }

    Print(L"---- Faild start kernel ----\n");

    return EFI_SUCCESS;
}