/* ukibootctl
   Copyright (C) 2025 Alexander Larsson <alexl@redhat.com>

   This file is free software: you can redistribute it and/or modify
   it under the terms of the GNU Lesser General Public License as
   published by the Free Software Foundation; either version 2.1 of the
   License, or (at your option) any later version.

   This file is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Lesser General Public License for more details.

   You should have received a copy of the GNU Lesser General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/
#include "config.h"

#include "bootctl.h"

#include <fcntl.h>
#include <getopt.h>
#include <stdio.h>

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

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

    return 0;
}

static void
bootctl_update_crc32(UkiBootCtrl *bootctl)
{
    uint32_t crc = crc32((uint8_t *)bootctl, sizeof(UkiBootCtrl) - sizeof(uint32_t));
    bootctl->CRC32 = htole32(crc);
}

static void
bootctl_prepare_switch(UkiBootCtrl *bootctl, uint32_t new_slot)
{
    uint32_t old_slot = bootctl_slot_other(new_slot);

    // Ensure old slot is bootable
    bootctl->slots[old_slot].successful_boot = 1;
    bootctl->slots[old_slot].tries_remaining = 0;

    // Mark new slot non-bootable
    bootctl->slots[new_slot].successful_boot = 0;
    bootctl->slots[new_slot].tries_remaining = 0;
    bootctl->slots[new_slot].priority = 0;

    bootctl_update_crc32(bootctl);
}

static void
bootctl_finalize_switch(UkiBootCtrl *bootctl, uint32_t new_slot)
{
    uint32_t old_slot = bootctl_slot_other(new_slot);

    // Ensure old slot lower prio
    bootctl->slots[old_slot].priority = 14;

    // Mark new slot as active for 7 tries
    bootctl->slots[new_slot].successful_boot = 0;
    bootctl->slots[new_slot].tries_remaining = 7;
    bootctl->slots[new_slot].priority = 15;

    bootctl_update_crc32(bootctl);
}

static void
bootctl_mark_successful(UkiBootCtrl *bootctl, uint32_t slot)
{
    bootctl->slots[slot].successful_boot = 1;
    bootctl->slots[slot].tries_remaining = 0;

    bootctl_update_crc32(bootctl);
}

static bool
bootctl_read(UkiBootCtrl *bootctl, const char *path)
{
    autoclose int fd = open(path, O_RDONLY | O_CLOEXEC);
    if (fd < 0) {
	error("Failed to open %s: %s\n", path, strerror(errno));
	return false;
    }

    ssize_t res = read(fd, (void *)bootctl, sizeof(UkiBootCtrl));
    if (res < 0) {
	error("Failed to read %s: %s\n", path, strerror(errno));
	return false;
    }
    if (res != sizeof(UkiBootCtrl)) {
	error("Failed to read %s: short read\n", path);
	return false;
    }

    return true;
}

static bool
bootctl_read_and_validate(UkiBootCtrl *bootctl, const char *path)
{
    if (!bootctl_read(bootctl, path)) {
	return false;
    }

    if (!bootctl_validate(bootctl)) {
	error("Bootctl partition is invalid\n");
	return false;
    }

    return true;
}

static bool
bootctl_write(UkiBootCtrl *bootctl, const char *path)
{
    autoclose int fd = open(path, O_WRONLY | O_CLOEXEC);
    if (fd < 0) {
	error("Failed to open %s: %s\n", path, strerror(errno));
	return false;
    }

    ssize_t res = write(fd, (void *)bootctl, sizeof(UkiBootCtrl));
    if (res < 0) {
	error("Failed to write %s: %s\n", path, strerror(errno));
	return false;
    }
    if (res != sizeof(UkiBootCtrl)) {
	error("Failed to write %s: short write\n", path);
	return false;
    }

    return true;
}

#define OPT_VERSION   102
#define OPT_VERBOSE   103
#define OPT_PARTITION 104
#define OPT_HELP      105

const char *opt_partition = "/dev/disk/by-partlabel/ukibootctl";

static int
cmd_init(int argc, arg_unused char **argv)
{
    if (argc != 1) {
	error("Unexpected arguments to init\n");
	return 1;
    }

    UkiBootCtrl bootctl;

    memset(&bootctl, 0, sizeof(UkiBootCtrl));
    bootctl.magic = htole32(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;

    bootctl_update_crc32(&bootctl);

    if (!bootctl_write(&bootctl, opt_partition)) {
	return 1;
    }

    return 0;
}

static int
cmd_dump(int argc, arg_unused char **argv)
{
    if (argc != 1) {
	error("Unexpected arguments to dump\n");
	return 1;
    }

    UkiBootCtrl bootctl;
    if (!bootctl_read_and_validate(&bootctl, opt_partition)) {
	return 1;
    }

    for (int i = 0; i < 2; i++) {
	UkiBootSlotInfo *slot = &bootctl.slots[i];
	printf("Slot %d: prio: %d, tries_remaining: %d, successful_boot: %d\n",
	       i,
	       slot->priority,
	       slot->tries_remaining,
	       slot->successful_boot);
    }

    return 0;
}

static int
cmd_get_active(int argc, arg_unused char **argv)
{
    if (argc != 1) {
	error("Unexpected arguments to get-active\n");
	return 1;
    }

    UkiBootCtrl bootctl;
    if (!bootctl_read_and_validate(&bootctl, opt_partition)) {
	return 1;
    }

    uint32_t slot;
    if (!bootctl_find_active_slot(&bootctl, &slot)) {
	error("No active slot\n");
	return 1;
    }

    printf("%d\n", slot);

    return 0;
}

static int
get_booted_slot(void)
{
    int slot = 0;

    autofree char *cmdline = read_proc_cmdline();
    if (cmdline == NULL)
	fatal("Failed to read /proc/cmdline\n");

    autofree char *slot_suffix = find_proc_cmdline_key(cmdline, "androidboot.slot_suffix");

    if (slot_suffix == NULL) {
	error("No slot suffix in command line, assuming slot 0\n");
    } else {
	if (strcmp(slot_suffix, "_a") == 0)
	    slot = 0;
	else if (strcmp(slot_suffix, "_b") == 0)
	    slot = 1;
	else
	    fatal("androidboot.slot_suffix invalid: %s\n", slot_suffix);
    }

    return slot;
}

static int
cmd_get_booted(int argc, arg_unused char **argv)
{
    if (argc != 1) {
	error("Unexpected arguments to get-booted\n");
	return 1;
    }

    int slot = get_booted_slot();
    printf("%d\n", slot);

    return 0;
}

static int
parse_slot(const char *slot_s)
{
    char *slot_s_end;
    uint32_t slot = strtol(slot_s, &slot_s_end, 10);
    if (*slot_s == 0 || *slot_s_end != 0) {
	return -1;
    }
    if (slot != 0 && slot != 1)
	return -1;

    return slot;
}

static int
cmd_prepare_switch(int argc, arg_unused char **argv)
{
    if (argc < 2) {
	error("No slot specified\n");
	return 1;
    }

    if (argc > 2) {
	error("Unexpected arguments to prepare-switch\n");
	return 1;
    }

    char *slot_s = argv[1];
    int32_t slot = parse_slot(slot_s);
    if (slot < 0) {
	error("Invalid slot %s\n", slot_s);
	return 1;
    }

    UkiBootCtrl bootctl;
    if (!bootctl_read_and_validate(&bootctl, opt_partition)) {
	return 1;
    }

    bootctl_prepare_switch(&bootctl, slot);

    if (!bootctl_write(&bootctl, opt_partition)) {
	return 1;
    }

    return 0;
}

static int
cmd_finalize_switch(int argc, arg_unused char **argv)
{
    if (argc < 2) {
	error("No slot specified\n");
	return 1;
    }

    if (argc > 2) {
	error("Unexpected arguments to finalize-switch\n");
	return 1;
    }

    char *slot_s = argv[1];
    int32_t slot = parse_slot(slot_s);
    if (slot < 0) {
	error("Invalid slot %s\n", slot_s);
	return 1;
    }

    UkiBootCtrl bootctl;
    if (!bootctl_read_and_validate(&bootctl, opt_partition)) {
	return 1;
    }

    bootctl_finalize_switch(&bootctl, slot);

    if (!bootctl_write(&bootctl, opt_partition)) {
	return 1;
    }

    return 0;
}

static int
cmd_mark_successful(int argc, arg_unused char **argv)
{
    int32_t slot = 0;

    if (argc < 2) {
	/* No slot specified, use the booted */
	slot = get_booted_slot();
    } else if (argc == 2) {
	char *slot_s = argv[1];
	slot = parse_slot(slot_s);
	if (slot < 0) {
	    error("Invalid slot %s\n", slot_s);
	    return 1;
	}
    } else {
	error("Unexpected arguments to mark-successful\n");
	return 1;
    }

    UkiBootCtrl bootctl;
    if (!bootctl_read_and_validate(&bootctl, opt_partition)) {
	return 1;
    }

    bootctl_mark_successful(&bootctl, slot);

    if (!bootctl_write(&bootctl, opt_partition)) {
	return 1;
    }

    return 0;
}

typedef struct {
    char *name;
    int (*fn)(int argc, char **argv);
} CmdHandler;

CmdHandler commands[] = {
    {"init", cmd_init},
    {"dump", cmd_dump},
    {"get-active", cmd_get_active},
    {"get-booted", cmd_get_booted},
    {"prepare-switch", cmd_prepare_switch},
    {"finalize-switch", cmd_finalize_switch},
    {"mark-successful", cmd_mark_successful},
};

static CmdHandler *
find_command_handler(const char *name)
{
    for (size_t i = 0; i < N_ELEMENTS(commands); i++) {
	CmdHandler *handler = &commands[i];
	if (strcmp(name, handler->name) == 0)
	    return handler;
    }

    return NULL;
}

static void
usage(const char *argv0)
{
    const char *bin = gnu_basename(argv0);
    fprintf(stderr,
	    "Usage: %s [OPTIONS] COMMAND [...]\n"
	    "Options:\n"
	    "  --verbose               Show verbose output\n"
	    "  --help                  Show help\n"
	    "  --partition=PATH        Path to ukibootctrl partition\n"
	    "                          Default: /dev/disk/by-partlabel/ukibootctl\n"
	    "Commands:\n"
	    "  init                    Initialize boot options with default values\n"
	    "  dump                    Dump current boot options\n"
	    "  get-active              Show current active slot (on disk)\n"
	    "  get-booted              Show current booted slot\n"
	    "  prepare-switch SLOT     Prepare to switch to SLOT (before update)\n"
	    "  finalize-switch SLOT    Finalize switch to SLOT (after update)\n"
	    "  mark-successful [SLOT]  Mark boot successful (after update & boot)\n",
	    bin);
}

int
main(int argc, char **argv)
{
    const char *bin = argv[0];
    const struct option longopts[] = {
	{.name = "help", .has_arg = no_argument, .flag = NULL, .val = OPT_HELP},
	{.name = "version", .has_arg = no_argument, .flag = NULL, .val = OPT_VERSION},
	{.name = "verbose", .has_arg = no_argument, .flag = NULL, .val = OPT_VERBOSE},
	{.name = "partition", .has_arg = required_argument, .flag = NULL, .val = OPT_PARTITION},
	{},
    };
    int opt;

    while ((opt = getopt_long(argc, argv, "", longopts, NULL)) != -1) {
	switch (opt) {
	case OPT_HELP:
	    usage(bin);
	    exit(0);
	    break;
	case OPT_VERSION:
	    printf("ukibootctl %s\n", PACKAGE_VERSION);
	    exit(0);
	    break;
	case OPT_VERBOSE:
	    opt_verbose = true;
	    break;
	case OPT_PARTITION:
	    opt_partition = optarg;
	    break;
	case ':':
	    fprintf(stderr, "option needs a value\n");
	    exit(EXIT_FAILURE);
	default:
	    usage(bin);
	    exit(1);
	}
    }

    argv += optind;
    argc -= optind;

    if (argc < 1) {
	error("No command specified\n");
	usage(bin);
	exit(1);
    }
    const char *command = argv[0];

    CmdHandler *handler = find_command_handler(command);

    if (handler == NULL) {
	error("Unknown command '%s'\n", command);
	usage(bin);
	return 1;
    }

    if (handler->fn(argc, argv)) {
	return 1;
    }

    return 0;
}
