#include <efi.h>
#include <efigpt.h>
#include <efilib.h>
#include <stddef.h>

#include "bootloader.h"
#include "ukiboot-efi.h"

/* This adds minimal support for the boot loader interface defined in
 * https://systemd.io/BOOT_LOADER_INTERFACE/.
 *
 * At the moment we only set Loasystemd-gpt-auto-generatorderDevicePartUUID, which is enough to
 * allow systemd-gpt-auto-generator to auto-find all mountpoints.
 */

/* Defined in  */
#define BOOT_LOADER_INTERFACE_VENDOR_UUID                                                          \
    {0x4a67b082, 0x0a4c, 0x41cf, {0xb6, 0xc7, 0x44, 0x0b, 0x29, 0xbb, 0x8c, 0x4f}}

static EFI_STATUS
set_efivar(const EFI_GUID *vendor, const CHAR16 *name, const VOID *buf, size_t size, UINT32 flags)
{
    flags |= EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS;
    return uefi_call_wrapper(RT->SetVariable,
			     5,
			     (CHAR16 *)name,
			     (EFI_GUID *)vendor,
			     flags,
			     size,
			     (VOID *)buf);
}

static EFI_STATUS
set_efivar_str(const EFI_GUID *vendor, const CHAR16 *name, const CHAR16 *value, UINT32 flags)
{
    return set_efivar(vendor, name, value, value ? (StrLen(value) + 1) * sizeof(CHAR16) : 0, flags);
}

static EFI_GUID *
device_path_find_partition_guid(EFI_DEVICE_PATH_PROTOCOL *dp)
{
    while (!IsDevicePathEnd(dp)) {
	if (DevicePathType(dp) == MEDIA_DEVICE_PATH &&
	    DevicePathSubType(dp) == MEDIA_HARDDRIVE_DP) {
	    HARDDRIVE_DEVICE_PATH *Hd = (HARDDRIVE_DEVICE_PATH *)dp;
	    if (Hd->SignatureType == SIGNATURE_TYPE_GUID) {
		return (EFI_GUID *)Hd->Signature;
	    }
	}
	dp = NextDevicePathNode(dp);
    }

    return NULL;
}

EFI_STATUS
bootloader_interface_init(EFI_HANDLE loaded_image_handle)
{
    EFI_STATUS rc;
    _cleanup_protocol Protocol proto = PROTOCOL_INIT;
    EFI_GUID lidpp = EFI_LOADED_IMAGE_DEVICE_PATH_PROTOCOL_GUID;
    EFI_GUID boot_loader_iface = BOOT_LOADER_INTERFACE_VENDOR_UUID;

    rc = efi_open_protocol_exclusive(loaded_image_handle, &lidpp, loaded_image_handle, &proto);
    if (EFI_ERROR(rc)) {
	log(L"Could not open loaded_image devicepath protocol");
	return rc;
    }
    EFI_DEVICE_PATH_PROTOCOL *dp = proto.iface;

    EFI_GUID *guid = device_path_find_partition_guid(dp);
    if (guid != NULL) {
	CHAR16 guid_str[64];
	GuidToString(guid_str, guid);
	set_efivar_str(&boot_loader_iface, L"LoaderDevicePartUUID", guid_str, 0);
    }

    return EFI_SUCCESS;
}
