/*
 * Copyright (c) 2026 The Fuchsia Authors
 *
 * SPDX-License-Identifier:	BSD-3-Clause
 */

#include <common.h>
#include <efi_loader.h>
#include <gbl_efi_fastboot_protocol.h>
#include <mmc.h>
#include <amlogic/aml_mmc.h>
#include <version.h>
#include <g_dnl.h>
#include <fb_fastboot.h>
#include <partition_table.h>

static efi_status_t EFIAPI get_var(struct GblEfiFastbootProtocol *self,
				   size_t num_args, const char *const *args,
				   size_t *buffer_size, char *buffer)
{
	EFI_ENTRY("%p, %zu, %p, %p, %p", self, num_args, args, buffer_size, buffer);
	if (num_args == 0 || !args || !buffer_size || (*buffer_size > 0 && !buffer))
		return EFI_EXIT(EFI_INVALID_PARAMETER);

	char cmd_buf[RESPONSE_LEN];
	char resp_buf[RESPONSE_LEN] = "";
	snprintf(cmd_buf, sizeof(cmd_buf), "%s", args[0]);

	if (fastboot_get_var(cmd_buf, resp_buf, sizeof(resp_buf) - 1) != 0)
		return EFI_EXIT(EFI_NOT_FOUND);

	size_t len = strlen(resp_buf);
	if (*buffer_size < len + 1) {
		*buffer_size = len + 1;
		return EFI_EXIT(EFI_BUFFER_TOO_SMALL);
	}

	strncpy(buffer, resp_buf, *buffer_size);
	*buffer_size = len;
	return EFI_EXIT(EFI_SUCCESS);
}

static efi_status_t EFIAPI get_var_all(struct GblEfiFastbootProtocol *self,
				       void *context, GetVarAllCallback cb)
{
	EFI_ENTRY("%p, %p, %p", self, context, cb);
	if (!cb)
		return EFI_EXIT(EFI_INVALID_PARAMETER);

	const char (*var_list)[RESPONSE_LEN];
	size_t var_list_size;

	if (has_boot_slot == 1) {
		var_list = fastboot_getvar_list_ab_gbl_limited;
		var_list_size = fastboot_getvar_list_ab_gbl_limited_size;
	} else if (dynamic_partition) {
		var_list = fastboot_getvar_list_dynamic;
		var_list_size = fastboot_getvar_list_dynamic_size;
	} else {
		var_list = fastboot_getvar_list;
		var_list_size = fastboot_getvar_list_size;
	}

	for (size_t i = 0; i < var_list_size; i++) {
		const char *name = var_list[i];
		char cmd_buf[RESPONSE_LEN];
		char resp_buf[RESPONSE_LEN] = "";
		snprintf(cmd_buf, sizeof(cmd_buf), "%s", name);
		if (fastboot_get_var(cmd_buf, resp_buf, sizeof(resp_buf) - 1) == 0 &&
		    strncmp(resp_buf, "FAIL", 4) != 0) {
			const char *const args[] = { name };
			cb(context, 1, args, resp_buf);
		}
	}
	return EFI_EXIT(EFI_SUCCESS);
}

static efi_status_t EFIAPI get_staged(struct GblEfiFastbootProtocol *self,
				      size_t *buffer_size, size_t *buffer_remains,
				      uint8_t *buffer)
{
	return EFI_UNSUPPORTED;
}

static efi_status_t EFIAPI get_partition_type(struct GblEfiFastbootProtocol *self,
					      const char *part_name,
					      size_t *part_type_len, char *part_type)
{
	EFI_ENTRY("%p, %s, %p, %p", self, part_name, part_type_len, part_type);
	if (!part_name || !part_type_len || (*part_type_len > 0 && !part_type))
		return EFI_EXIT(EFI_INVALID_PARAMETER);

	char cmd_buf[RESPONSE_LEN];
	snprintf(cmd_buf, sizeof(cmd_buf), "partition-type:%s", part_name);
	const char *part = parse_fb_cmd_partition(cmd_buf);
	if (!part)
		return EFI_EXIT(EFI_NOT_FOUND);

	const char *type_str = "raw";
	size_t len = strlen(type_str);
	if (*part_type_len < len + 1) {
		*part_type_len = len + 1;
		return EFI_EXIT(EFI_BUFFER_TOO_SMALL);
	}

	strncpy(part_type, type_str, *part_type_len);
	*part_type_len = len;
	return EFI_EXIT(EFI_SUCCESS);
}

static efi_status_t handle_flash_bootloader(size_t download_buffer_used_size,
					    uint8_t *download_buffer,
					    FastbootMessageSender sender, void *context)
{
	if (sender) {
		const char *msg = "Flashing bootloader via amlmmc...\n";
		sender(context, GBL_EFI_FASTBOOT_MESSAGE_TYPE_INFO, strlen(msg), msg);
	}

	// Call amlmmc_write_bootloader
	// We use AML_BL_BOOT to write to both boot0 and boot1
	int ret = amlmmc_write_bootloader(CONFIG_FASTBOOT_FLASH_MMC_DEV,
					  AML_BL_BOOT,
					  download_buffer_used_size,
					  download_buffer);
	if (ret) {
		if (sender) {
			const char *msg = "Flash bootloader failed\n";
			sender(context, GBL_EFI_FASTBOOT_MESSAGE_TYPE_FAIL, strlen(msg), msg);
		}
		return EFI_DEVICE_ERROR;
	}

	if (sender) {
		const char *msg = "Flash bootloader success\n";
		sender(context, GBL_EFI_FASTBOOT_MESSAGE_TYPE_OKAY, strlen(msg), msg);
	}
	return EFI_SUCCESS;
}

static efi_status_t handle_run_staged_efi(size_t download_buffer_used_size,
					  uint8_t *download_buffer,
					  FastbootMessageSender sender, void *context)
{
	if (download_buffer_used_size == 0 || !download_buffer) {
		if (sender) {
			const char *msg = "No EFI image staged\n";
			sender(context, GBL_EFI_FASTBOOT_MESSAGE_TYPE_FAIL, strlen(msg), msg);
		}
		return EFI_NOT_FOUND;
	}

	if (sender) {
		const char *msg = "Running staged EFI image...\n";
		sender(context, GBL_EFI_FASTBOOT_MESSAGE_TYPE_INFO, strlen(msg), msg);
	}

	// We need to run the image.
	// Similar to cb_run_staged_efi, we can use bootefi command.
	char cmd[128];
	snprintf(cmd, sizeof(cmd), "bootefi 0x%lx:0x%lx ${dtb_mem_addr}",
		 (unsigned long)download_buffer,
		 (unsigned long)download_buffer_used_size);

	printf("Executing: %s\n", cmd);

	// This might not return if successful
	run_command(cmd, 0);

	// If it returns, it failed or exited
	if (sender) {
		const char *msg = "EFI image returned\n";
		sender(context, GBL_EFI_FASTBOOT_MESSAGE_TYPE_FAIL, strlen(msg), msg);
	}
	return EFI_DEVICE_ERROR;
}

typedef efi_status_t (*fastboot_cmd_handler_t)(size_t download_buffer_used_size,
					      uint8_t *download_buffer,
					      FastbootMessageSender sender, void *context);

struct fastboot_cmd_match {
	const char *const *args;
	size_t num_args;
	fastboot_cmd_handler_t handler;
};

#define CMD_ENTRIES_2(cb, w1, w2) \
	{ \
		.args = (const char*const[]){w1, w2}, \
		.num_args = 2, \
		.handler = cb, \
	}, \
	{ \
		.args = (const char*const[]){w1 " " w2}, \
		.num_args = 1, \
		.handler = cb, \
	}

static const struct fastboot_cmd_match cmd_matches[] = {
	CMD_ENTRIES_2(handle_flash_bootloader, "flash", "bootloader"),
	CMD_ENTRIES_2(handle_run_staged_efi, "oem", "run-staged-efi"),
};

static efi_status_t EFIAPI command_exec(struct GblEfiFastbootProtocol *self,
					size_t num_args, const char *const *args,
					size_t download_buffer_size,
					size_t download_buffer_used_size,
					uint8_t *download_buffer,
					GblEfiFastbootCommandExecResult *implementation,
					FastbootMessageSender sender, void *context)
{
	EFI_ENTRY("%p, %zu, %p, %zu, %zu, %p", self, num_args, args,
		  download_buffer_size, download_buffer_used_size, download_buffer);

	if (num_args == 0 || !args || !implementation)
		return EFI_EXIT(EFI_INVALID_PARAMETER);

	// Default to not implemented by us
	*implementation = GBL_EFI_FASTBOOT_COMMAND_EXEC_RESULT_DEFAULT_IMPL;

	for (size_t i = 0; i < ARRAY_SIZE(cmd_matches); i++) {
		const struct fastboot_cmd_match *match = &cmd_matches[i];
		if (num_args < match->num_args)
			continue;

		bool found = true;
		for (size_t j = 0; j < match->num_args; j++) {
			if (strcmp(args[j], match->args[j]) != 0) {
				found = false;
				break;
			}
		}

		if (found) {
			*implementation = GBL_EFI_FASTBOOT_COMMAND_EXEC_RESULT_CUSTOM_IMPL;
			return EFI_EXIT(match->handler(download_buffer_used_size,
						      download_buffer,
						      sender, context));
		}
	}

	return EFI_EXIT(EFI_SUCCESS);
}

static const efi_guid_t guid = GBL_EFI_FASTBOOT_PROTOCOL_GUID;

static GblEfiFastbootProtocol protocol = {
	.revision = GBL_EFI_FASTBOOT_PROTOCOL_REVISION,
	.serial_number = "VIM3-GBL-FASTBOOT",
	.get_var = get_var,
	.get_var_all = get_var_all,
	.get_staged = get_staged,
	.command_exec = command_exec,
	.get_partition_type = get_partition_type,
};

efi_status_t efi_gbl_fastboot_register(void)
{
	char *serial = getenv("serial#");
	if (serial) {
		strncpy(protocol.serial_number, serial, sizeof(protocol.serial_number) - 1);
		protocol.serial_number[sizeof(protocol.serial_number) - 1] = '\0';
	}

	efi_status_t ret = efi_add_protocol(efi_root, &guid, (void *)&protocol);
	if (ret != EFI_SUCCESS)
		printf("Cannot install GBL_EFI_FASTBOOT_PROTOCOL\n");
	return ret;
}