// SPDX-License-Identifier: GPL-3.0-or-later
#include "elilong.h"

// Helper to get the file system volume elilong was started from
EFI_FILE_HANDLE get_esp_root(EFI_HANDLE image_handle) {
    EFI_LOADED_IMAGE *loaded_image = NULL;
    EFI_GUID loaded_image_protocol = LOADED_IMAGE_PROTOCOL;
    EFI_GUID fs_protocol = SIMPLE_FILE_SYSTEM_PROTOCOL;

    EFI_FILE_IO_INTERFACE *io_device = NULL;
    EFI_FILE_HANDLE root = NULL;
    // Retrieve the Loaded Image protocol to find the device handle
    uefi_call_wrapper(ST->BootServices->HandleProtocol, 3, image_handle,
                               &loaded_image_protocol, (void**)&loaded_image);
    // Use the device handle to get the file system interface
    uefi_call_wrapper(ST->BootServices->HandleProtocol, 3, loaded_image->DeviceHandle,
                               &fs_protocol, (void**)&io_device);
    // Open the volume's root directory
    if (io_device) {
        uefi_call_wrapper(io_device->OpenVolume, 2, io_device, &root);
    }
    return root;
}

// Helper to load a file into memory
EFI_STATUS FileRead(EFI_FILE_HANDLE root, CHAR16 *path, VOID **buffer, UINTN *size) {
  EFI_FILE_HANDLE file = NULL;
  EFI_FILE_INFO *info = NULL;
  UINTN info_size = 0;
  EFI_STATUS status;
  EFI_GUID info_guid = EFI_FILE_INFO_ID;
  // Open the target file for reading.
  status = uefi_call_wrapper(root->Open, 5, root, &file, path, EFI_FILE_MODE_READ, 0);
  if (EFI_ERROR(status)) return status;

  // Get file size
  uefi_call_wrapper(file->GetInfo, 4, file, &info_guid, &info_size, NULL);
  uefi_call_wrapper(ST->BootServices->AllocatePool, 3, EfiLoaderData, info_size, (void**)&info);
  uefi_call_wrapper(file->GetInfo, 4, file, &info_guid, &info_size, info);
  // Allocate a buffer for the file data
  *size = info->FileSize;
  uefi_call_wrapper(ST->BootServices->AllocatePool, 3, EfiLoaderData, *size, buffer);
  // Read the actual content into the buffer
  status = uefi_call_wrapper(file->Read, 3, file, size, *buffer);
  // Cleanup temporary info and close the handle
  uefi_call_wrapper(ST->BootServices->FreePool, 1, info);
  uefi_call_wrapper(file->Close, 1, file);
  return status;
}
// manual_pe_load_and_jump: Generic loader for EFI PE/COFF images (Linux or Windows)
EFI_STATUS manual_pe_load_and_jump(EFI_HANDLE image_handle, EFI_FILE_HANDLE root, BOOT_ENTRY *entry) {
    EFI_STATUS status;
    EFI_HANDLE target_handle = NULL;
    EFI_LOADED_IMAGE *loaded_img = NULL;
    EFI_LOADED_IMAGE *parent_img = NULL;
    EFI_DEVICE_PATH *full_path = NULL;
    CHAR16 final_kernel_path[256];
    CHAR16 final_initrd_path[256];
    // add efi path if it doesn't begin with \ .
    if (entry->kernel_path[0] != L'\\') {
        StrCpy(final_kernel_path, L"\\EFI\\" SUBDIR L"\\");
        StrCat(final_kernel_path, entry->kernel_path);
    } else {
        StrCpy(final_kernel_path, entry->kernel_path);
    }

    if (entry->initrd_path[0] != L'\0' && entry->initrd_path[0] != L'\\') {
        StrCpy(final_initrd_path, L"\\EFI\\" SUBDIR L"\\");
        StrCat(final_initrd_path, entry->initrd_path);
    } else {
        StrCpy(final_initrd_path, entry->initrd_path);
    }
    Print(L"Attempting to boot: %s at %s\n", entry->label, entry->kernel_path);
    uefi_call_wrapper(ST->BootServices->Stall, 1, 5000000);

    // path = FileDevicePath(NULL, entry->kernel_path);

    /* status = uefi_call_wrapper(ST->BootServices->LoadImage, 6,
                               FALSE, image_handle, path,
                               NULL, 0, &target_handle);

    if (EFI_ERROR(status)) {
        Print(L"Error: Could not load %s (%r)\n", entry->kernel_path,status);
        return status;
    }  */
    // Get the Loaded Image protocol for elilong itself to find our partition
    status = uefi_call_wrapper(ST->BootServices->HandleProtocol, 3,
                               image_handle, &LoadedImageProtocol, (VOID**)&parent_img);
    if (EFI_ERROR(status)) {
        Print(L"Error: Could not get parent image protocol (%r)\n", status);
        return status;
    }

    // Create a FULL Device Path (Partition Handle + File Path)
    full_path = FileDevicePath(parent_img->DeviceHandle, entry->kernel_path);
    if (full_path == NULL) {
        Print(L"Error: Could not create device path for %s\n", entry->kernel_path);
        return EFI_OUT_OF_RESOURCES;
    }

    // Load the image using the full device path
    status = uefi_call_wrapper(ST->BootServices->LoadImage, 6,
                               FALSE,           // BootPolicy
                               image_handle,    // Parent image handle
                               full_path,       // The full path we just built
                               NULL,            // SourceBuffer
                               0,               // SourceSize
                               &target_handle);

    // We can free the path memory now that the image is loaded
    FreePool(full_path);

    if (EFI_ERROR(status)) {
        Print(L"Error: Could not load %s (%r)\n", entry->kernel_path, status);
        // Pause so the user can see the error before falling back to ELILO
        uefi_call_wrapper(ST->BootServices->Stall, 1, 3000000);
        return status;
    }

    if (entry->initrd_path[0] != L'\0' || entry->append[0] != L'\0') {
        status = uefi_call_wrapper(ST->BootServices->HandleProtocol, 3,
                               target_handle, &gEfiLoadedImageProtocolGuid, (VOID**)&loaded_img);

        if (!EFI_ERROR(status)) {
            static CHAR16 options[1024];
            ZeroMem(options, sizeof(options));
            options[0] = L'\0';

            // Construct the combined command: "initrd=\path\to\initrd append_string"
            if (entry->initrd_path[0] != L'\0') {
                StrCat(options, L"initrd=");
                StrCat(options, entry->initrd_path);
                StrCat(options, L" ");
            }
            if (entry->root_device[0] != L'\0') {
                StrCat(options, L"root=");
                StrCat(options, entry->root_device);
                StrCat(options, L" ");
            }
            if (entry->mount_mode == MOUNT_RO) {
                StrCat(options, L"ro ");
            } else if (entry->mount_mode == MOUNT_RW) {
                StrCat(options, L"rw ");
            }
            // Add the rest of the append options
            if (entry->append[0] != L'\0') {
                StrCat(options, entry->append);
            }

            // Assign the options to the loaded image.
            if (StrLen(options) > 0) {
                loaded_img->LoadOptions = options;
                loaded_img->LoadOptionsSize = (StrLen(options) + 1) * sizeof(CHAR16);
                Print(L"Kernel Command Line: %s\n", options);
            }
        }
    }

    // Start the kernel or other efi pe/coff binary.
    // If successful, control is handed off and this function never returns.
    Print(L"Starting %s...\n",entry->label);
    status = uefi_call_wrapper(ST->BootServices->StartImage, 3, target_handle, NULL, NULL);

    if (EFI_ERROR(status)) {
        Print(L"Error: Failed to start image (%r)\n", status);
        uefi_call_wrapper(ST->BootServices->UnloadImage, 1, target_handle);
        uefi_call_wrapper(ST->BootServices->Stall, 1, 3000000);
    }

    return status;
}
EFI_STATUS load_config(EFI_FILE_HANDLE root, CHAR16 *path, BOOT_CONFIG *config) {
    VOID *buffer = NULL;
    UINTN size = 0;
    EFI_STATUS status = FileRead(root, path, &buffer, &size);

    if (status == EFI_SUCCESS) {
        parse_config((char*)buffer, size, config);
        if (config->entries == NULL) return EFI_NOT_FOUND;
        return EFI_SUCCESS;
    }
    return status;
}
BOOLEAN wait_for_specific_char(CHAR16 target) {
    EFI_INPUT_KEY key;
    UINTN index;
    // Wait for a keystroke
    uefi_call_wrapper(ST->BootServices->WaitForEvent, 3, 1, &ST->ConIn->WaitForKey, &index);
    uefi_call_wrapper(ST->ConIn->ReadKeyStroke, 2, ST->ConIn, &key);

    // Check if it matches (case-insensitive for convenience)
    if (key.UnicodeChar == target || (target >= L'a' && target <= L'z' && key.UnicodeChar == (target - 32))) {
        return TRUE;
    }
    return FALSE;
}
