#include "config.h"

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

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

#define UKIBOOTCTL_PARTITION_TYPE_UUID                                                             \
    {0xfefd9070, 0x346f, 0x4c9a, {0x85, 0xe6, 0x17, 0xf0, 0x7f, 0x92, 0x27, 0x73}}

#define UKIBOOT_PARTITION_TYPE_UUID                                                                \
    {0xdf331e4d, 0xbe00, 0x463f, {0xb4, 0xa7, 0x8b, 0x43, 0xe1, 0x8f, 0xb5, 0x3a}}

#define UKIBOOT_PARTITION_TYPE_MBR_IDS	 {0x46, 0x47}
#define UKIBOOTCTL_PARTITION_TYPE_MBR_ID 0x48

typedef struct PACKED {
    UINT8 priority;
    UINT8 tries_remaining;
    UINT8 successful_boot;
    UINT8 reserved;
} UkiBootSlotInfo;

#define UKI_BOOT_CTRL_MAGIC 1420550408

/*
 * The size (in bytes) of the EFI_DEVICE_PATH_PROTOCOL header, which consists of three
 * fields: Type (1 byte), SubType (1 byte), and Length (2 bytes). So, the header i
 * always 4 bytes.
 */
#define DEVICE_PATH_HEADER_SIZE 4

typedef struct PACKED {
    UINT32 magic;
    UkiBootSlotInfo slots[2];
    UINT8 reserved[8];
    UINT32 CRC32;
} UkiBootCtrl;

CHAR16 *slot_names[] = {
    L"A",
    L"B",
};

CHAR16 *slot_partition_names[] = {
    L"ukiboot_a",
    L"ukiboot_b",
};

CHAR16 *slot_efi_name[] = {
    L"ukiboot_a.efi",
    L"ukiboot_b.efi",
};

UINT32
bootctl_slot_other(UINT32 slot)
{
    if (slot == 0) {
	return 1;
    }

    return 0;
}

BOOLEAN
bootctl_slot_valid(UkiBootSlotInfo *slot)
{
    return slot->successful_boot != 0 || slot->tries_remaining > 0;
}

EFI_STATUS
bootctl_update_crc32(UkiBootCtrl *bootctl)
{
    UINT32 crc;
    EFI_STATUS rc;

    rc = efi_calculate_crc32(bootctl, sizeof(UkiBootCtrl) - sizeof(UINT32), &crc);
    if (EFI_ERROR(rc)) {
	log(L"Can't compute CRC32");
	return rc;
    }

    bootctl->CRC32 = crc;
    return EFI_SUCCESS;
}

EFI_STATUS
bootctl_init(UkiBootCtrl *bootctl)
{
    SetMem(bootctl, sizeof(UkiBootCtrl), 0);
    bootctl->magic = UKI_BOOT_CTRL_MAGIC;
    bootctl->slots[0].priority = 15;
    bootctl->slots[0].tries_remaining = 7;
    bootctl->slots[0].successful_boot = 0;
    bootctl->slots[1].priority = 14;
    bootctl->slots[1].tries_remaining = 7;
    bootctl->slots[1].successful_boot = 0;

    return bootctl_update_crc32(bootctl);
}

EFI_STATUS
bootctl_validate(UkiBootCtrl *bootctl)
{
    EFI_STATUS rc;

    if (bootctl->magic != UKI_BOOT_CTRL_MAGIC) {
	return EFI_INVALID_PARAMETER;
    }

    UINT32 crc;

    rc = efi_calculate_crc32(bootctl, sizeof(UkiBootCtrl) - sizeof(UINT32), &crc);
    if (EFI_ERROR(rc)) {
	log(L"Can't compute CRC32");
	return rc;
    }

    if (bootctl->CRC32 != crc) {
	log(L"Invalid bootctrl CRC");
	return EFI_INVALID_PARAMETER;
    }

    return EFI_SUCCESS;
}

EFI_STATUS
bootctl_find_active_slot(UkiBootCtrl *bootctl, UINT32 *out_slot)
{
    INT32 active_slot = -1;
    for (INT32 i = 0; i < 2; i++) {
	if (bootctl_slot_valid(&bootctl->slots[i])) {
	    if (active_slot == -1) {
		active_slot = i;
	    } else {
		if (bootctl->slots[i].priority > bootctl->slots[active_slot].priority) {
		    active_slot = i;
		}
	    }
	}
    }

    if (active_slot != -1) {
	*out_slot = active_slot;
	return EFI_SUCCESS;
    }

    return EFI_NOT_FOUND;
}

EFI_STATUS
find_partition(EFI_HANDLE loaded_image,
	       EFI_GUID *partition_type,
	       CHAR16 *name,
	       UINT8 mbr_partition_type,
	       EFI_HANDLE *partition_out)
{
    UINTN n_handles;
    EFI_STATUS rc;
    _cleanup_free EFI_HANDLE *partitions = NULL;
    EFI_GUID pip = PARTITION_INFO_PROTOCOL;

    rc = uefi_call_wrapper(BS->LocateHandleBuffer,
			   5,
			   ByProtocol,
			   &pip,
			   NULL,
			   &n_handles,
			   (EFI_HANDLE **)&partitions);
    if (EFI_ERROR(rc)) {
	log(L"Could not find partition info");
	return rc;
    }

    EFI_HANDLE found_partition = 0;

    for (UINTN i = 0; i < n_handles; i++) {
	EFI_HANDLE partition = partitions[i];
	EFI_PARTITION_INFO_PROTOCOL *info;
	_cleanup_protocol Protocol info_p = PROTOCOL_INIT;

	rc = efi_open_protocol_exclusive(partition, &pip, loaded_image, &info_p);
	if (EFI_ERROR(rc)) {
	    log(L"Could not open partition info protocol");
	    continue;
	}
	info = info_p.iface;

	if (info->Revision == EFI_PARTITION_INFO_PROTOCOL_REVISION) {
	    if (info->Type == PARTITION_TYPE_GPT) {
		EFI_PARTITION_ENTRY *gpt = &info->Info.Gpt;

		if (CompareGuid(&gpt->PartitionTypeGUID, partition_type) == 0 &&
		    StrCmp(name, gpt->PartitionName) == 0) {
		    found_partition = partition;
		    break;
		}
	    } else if (info->Type == PARTITION_TYPE_MBR) {
		MBR_PARTITION_RECORD *mbr = &info->Info.Mbr;
		if (mbr->OSIndicator == mbr_partition_type) {
		    found_partition = partition;
		    break;
		}
	    }
	}
    }

    if (found_partition) {
	*partition_out = found_partition;
	return EFI_SUCCESS;
    }

    return EFI_NOT_FOUND;
}

UINT32
round_up_to_blocksize(UINT32 size, UINT32 block_size)
{
    UINT32 n_blocks = (size + (block_size - 1)) / block_size;
    return n_blocks * block_size;
}

VOID *
allocate_aligned_buffer(UINT32 buf_size, EFI_MEMORY_TYPE mem_type)
{
    EFI_STATUS rc;
    EFI_PHYSICAL_ADDRESS addr = 0;

    rc = uefi_call_wrapper(BS->AllocatePages,
			   4,
			   AllocateAnyPages,
			   mem_type,
			   (buf_size + EFI_PAGE_SIZE - 1) / EFI_PAGE_SIZE,
			   &addr);
    if (EFI_ERROR(rc))
	return NULL;

    return (VOID *)addr;
}

EFI_STATUS
find_active_slot(EFI_HANDLE loaded_image, UINT32 *slot_out)
{
    EFI_STATUS rc;
    EFI_HANDLE bootctl_partition;
    _cleanup_protocol Protocol proto = PROTOCOL_INIT;
    EFI_GUID partition_type = UKIBOOTCTL_PARTITION_TYPE_UUID;
    EFI_BLOCK_IO_PROTOCOL *blockio;

    rc = find_partition(loaded_image,
			&partition_type,
			L"ukibootctl",
			UKIBOOTCTL_PARTITION_TYPE_MBR_ID,
			&bootctl_partition);
    if (EFI_ERROR(rc)) {
	log(L"Unable to find ukibootctl partition");
	return rc;
    }

    rc = efi_open_protocol_exclusive(bootctl_partition, &BlockIoProtocol, loaded_image, &proto);
    if (EFI_ERROR(rc)) {
	log(L"Unable to open partition blockio");
	return rc;
    }
    blockio = proto.iface;

    UINT32 block_size = blockio->Media->BlockSize;
    UINT32 buf_size = round_up_to_blocksize(sizeof(UkiBootCtrl), block_size);

    VOID *buffer = allocate_aligned_buffer(buf_size, EfiLoaderData);
    if (buffer == NULL) {
	log(L"Unable to allocate memory to read ukibootctl data");
	return EFI_OUT_OF_RESOURCES;
    }

    rc = efi_read_blocks(blockio, 0, buf_size, buffer);
    if (EFI_ERROR(rc)) {
	log(L"Unable to read ukibootctl partition");
	return rc;
    }

    UkiBootCtrl *bootctl = (UkiBootCtrl *)buffer;

    rc = bootctl_validate(bootctl);
    if (EFI_ERROR(rc)) {
	log(L"Uninitialized uki_bootctl partition, reseting to default");

	rc = bootctl_init(bootctl);
	if (EFI_ERROR(rc)) {
	    log(L"Unable to init ukibootctl partition");
	    return rc;
	}

	rc = efi_write_blocks(blockio, 0, buf_size, buffer);
	if (EFI_ERROR(rc)) {
	    log(L"Unable to write initialized ukibootctl partition");
	    return rc;
	}
    }

    UINT32 active_slot;
    rc = bootctl_find_active_slot(bootctl, &active_slot);
    if (EFI_ERROR(rc)) {
	log(L"Unable to find active slot");
	return rc;
    }

    if (bootctl->slots[active_slot].successful_boot == 0) {
	log(L"Test booting slot %s, tries_remaining: %d",
	    slot_names[active_slot],
	    bootctl->slots[active_slot].tries_remaining);
	bootctl->slots[active_slot].tries_remaining -= 1;

	rc = bootctl_update_crc32(bootctl);
	if (EFI_ERROR(rc)) {
	    log(L"Unable to update tries_remaining");
	    return rc;
	}

	rc = efi_write_blocks(blockio, 0, buf_size, buffer);
	if (EFI_ERROR(rc)) {
	    log(L"Unable to write updated ukibootctl partition");
	    return rc;
	}
    }

    *slot_out = active_slot;
    return EFI_SUCCESS;
}

UINTN
get_filepath_size(EFI_DEVICE_PATH_PROTOCOL *dp_root)
{
    UINTN total_size = 0;

    EFI_DEVICE_PATH_PROTOCOL *dp;
    UINT8 *ptr = (UINT8 *)dp_root;
    do {
	dp = (EFI_DEVICE_PATH_PROTOCOL *)ptr;
	UINTN size = DevicePathNodeLength(dp);

	total_size += size;

	ptr += size;
    } while (!IsDevicePathEndType(dp));

    return total_size;
}

INT32
find_last_separator(CHAR16 *path)
{
    INT32 last_separator = -1;
    INT32 i = 0;

    while (path[i] != 0) {
	if (path[i] == L'\\') {
	    last_separator = i;
	}
	i++;
    }

    if (last_separator >= 0)
	return last_separator;

    /* No separator, return after path end */
    return i;
}

EFI_STATUS
get_fake_device_path_for_slot(EFI_HANDLE loaded_image,
			      UINT32 active_slot,
			      EFI_DEVICE_PATH **path_out)
{
    EFI_STATUS rc;
    _cleanup_protocol Protocol proto = PROTOCOL_INIT;
    EFI_GUID lidpp = EFI_LOADED_IMAGE_DEVICE_PATH_PROTOCOL_GUID;

    CHAR16 *fake_filename = slot_efi_name[active_slot];

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

    EFI_DEVICE_PATH_PROTOCOL *dp_root = proto.iface;

    /* We should have at least non-end node */
    if (IsDevicePathEndType(dp_root)) {
	log(L"Invalid loaded image device path");
	return EFI_INVALID_PARAMETER;
    }

    UINTN orig_size = get_filepath_size(dp_root);
    /* Ensure we have space for the original filepath + modified basename */
    _cleanup_free VOID *new_dp =
	AllocatePool(orig_size + (1 + StrLen(fake_filename) + 1) * sizeof(CHAR16));
    if (new_dp == NULL) {
	log(L"Unable to allocate memory for device path");
	return EFI_OUT_OF_RESOURCES;
    }

    UINT8 *ptr = (UINT8 *)dp_root;
    UINT8 *dst_ptr = (UINT8 *)new_dp;
    EFI_DEVICE_PATH_PROTOCOL *src_dp;

    src_dp = (EFI_DEVICE_PATH_PROTOCOL *)ptr;
    while (!IsDevicePathEndType(src_dp)) {
	UINTN src_size = DevicePathNodeLength(src_dp);
	EFI_DEVICE_PATH_PROTOCOL *next_src_dp = (EFI_DEVICE_PATH_PROTOCOL *)(ptr + src_size);
	EFI_DEVICE_PATH_PROTOCOL *dst_dp = (EFI_DEVICE_PATH_PROTOCOL *)dst_ptr;

	if (IsDevicePathEndType(next_src_dp)) {
	    /* This is the last (non-end) node, rewrite it */

	    if (DevicePathType(src_dp) != MEDIA_DEVICE_PATH ||
		DevicePathSubType(src_dp) != MEDIA_FILEPATH_DP) {
		log(L"Final device path element is not a file path node. Type=0x%x SubType=0x%x",
		    DevicePathType(src_dp),
		    DevicePathSubType(src_dp));
		return EFI_INVALID_PARAMETER;
	    }

	    FILEPATH_DEVICE_PATH *dst_fp = (FILEPATH_DEVICE_PATH *)dst_dp;
	    dst_dp->Type = src_dp->Type;
	    dst_dp->SubType = src_dp->SubType;

	    UINT32 new_path_len =
		(StrLen(ESP_DIR) + 1 + StrLen(fake_filename) + 1) * sizeof(CHAR16);

	    SetDevicePathNodeLength(dst_dp, new_path_len + DEVICE_PATH_HEADER_SIZE);
	    dst_ptr += new_path_len + DEVICE_PATH_HEADER_SIZE;

	    StrCpy(dst_fp->PathName, ESP_DIR);
	    StrCat(dst_fp->PathName, L"/");
	    StrCat(dst_fp->PathName, fake_filename);

	    log(L"Using uki pathname: %s", dst_fp->PathName);
	} else {
	    CopyMem(dst_dp, src_dp, src_size);
	    dst_ptr += src_size;
	}

	ptr += src_size;
	src_dp = (EFI_DEVICE_PATH_PROTOCOL *)ptr;
    }

    /* This is the end of device path node */
    EFI_DEVICE_PATH_PROTOCOL *end_dp = (EFI_DEVICE_PATH_PROTOCOL *)dst_ptr;
    end_dp->Type = END_DEVICE_PATH_TYPE;
    end_dp->SubType = END_ENTIRE_DEVICE_PATH_SUBTYPE;
    SetDevicePathNodeLength(end_dp, DEVICE_PATH_HEADER_SIZE);

    *path_out = (EFI_DEVICE_PATH *)_steal_pointer(&new_dp);

    return EFI_SUCCESS;
}

EFI_STATUS
boot(EFI_HANDLE loaded_image)
{
    EFI_STATUS rc;
    UINT32 active_slot;
    EFI_GUID partition_type = UKIBOOT_PARTITION_TYPE_UUID;
    UINT8 mbr_partition_ids[] = UKIBOOT_PARTITION_TYPE_MBR_IDS;

    rc = bootloader_interface_init(loaded_image);
    if (EFI_ERROR(rc)) {
	log(L"Warning: Unable to initialize bootloader interface: %d", rc);
    }

    rc = find_active_slot(loaded_image, &active_slot);
    if (EFI_ERROR(rc)) {
	log(L"Unable to find active slot");
	return rc;
    }

    CHAR16 *partition_name = slot_partition_names[active_slot];
    UINT8 mbr_partition_id = mbr_partition_ids[active_slot];

    EFI_HANDLE boot_partition;
    rc = find_partition(loaded_image,
			&partition_type,
			partition_name,
			mbr_partition_id,
			&boot_partition);
    if (EFI_ERROR(rc)) {
	log(L"Unable to find %s partition", partition_name);
	return rc;
    }

    _cleanup_protocol Protocol proto = PROTOCOL_INIT;
    EFI_BLOCK_IO_PROTOCOL *blockio;

    rc = efi_open_protocol_exclusive(boot_partition, &BlockIoProtocol, loaded_image, &proto);
    if (EFI_ERROR(rc)) {
	log(L"Unable to open partition blockio");
	return rc;
    }
    blockio = proto.iface;

    UINT32 block_size = blockio->Media->BlockSize;
    UINT32 buf_size = round_up_to_blocksize(PE_FILE_HEADER_BUFFER_SIZE, block_size);

    VOID *buffer = allocate_aligned_buffer(buf_size, EfiLoaderData);
    if (buffer == NULL) {
	log(L"Unable to allocate memory to read pe header");
	return EFI_OUT_OF_RESOURCES;
    }

    rc = efi_read_blocks(blockio, 0, buf_size, buffer);
    if (EFI_ERROR(rc)) {
	log(L"Unable to read pe header");
	return rc;
    }

    UINTN pe_size;
    rc = get_pe_size(buffer, buf_size, &pe_size);
    if (EFI_ERROR(rc)) {
	log(L"Unable to parse pe header");
	return rc;
    }

    buf_size = round_up_to_blocksize(pe_size, block_size);

    VOID *image_buf = allocate_aligned_buffer(buf_size, EfiLoaderCode);
    if (image_buf == NULL) {
	log(L"Unable to allocate memory for uki image");
	return EFI_OUT_OF_RESOURCES;
    }

    log(L"Loading UKI from partition %s", partition_name);

    rc = efi_read_blocks(blockio, 0, buf_size, image_buf);
    if (EFI_ERROR(rc)) {
	log(L"Unable to read uki partition");
	return rc;
    }

    _cleanup_free EFI_DEVICE_PATH *fake_path = NULL;
    rc = get_fake_device_path_for_slot(loaded_image, active_slot, &fake_path);
    if (EFI_ERROR(rc)) {
	log(L"Unable to construct device path for slot %d: %r", active_slot, rc);
	return rc;
    }

    log(L"Calling LoadImage");

    EFI_HANDLE image;
    rc = uefi_call_wrapper(BS->LoadImage,
			   6,
			   FALSE,
			   loaded_image,
			   fake_path,
			   image_buf,
			   pe_size,
			   &image);
    if (EFI_ERROR(rc)) {
	log(L"Unable to load uki image: %d", rc);
	return rc;
    }

    log(L"Starting UKI");

    rc = uefi_call_wrapper(BS->StartImage, 3, image, 0, NULL);
    if (EFI_ERROR(rc)) {
	log(L"Unable to load uki image: %d", rc);
	return rc;
    }

    /* Unexpected error if StartImage returns */
    return EFI_LOAD_ERROR;
}

EFI_STATUS
efi_main(EFI_HANDLE loaded_image, UNUSED EFI_SYSTEM_TABLE *SystemTable)
{
    EFI_STATUS rc;

    log(L"Starting ukiboot version %s", PACKAGE_VERSION_L);

    rc = boot(loaded_image);
    if (EFI_ERROR(rc)) {
	log(L"Error: Unable to boot (%r)", rc);
    }

    return rc;
}
