/*
 * Copyright (C) 2012 Jan Luebbe <j.luebbe@pengutronix.de>
 * Copyright (C) 2015 Marc Kleine-Budde <mkl@pengutronix.de>
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program 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 General Public License for more details.
 */

#include <bootstate.h>
#include <common.h>
#include <envfs.h>
#include <environment.h>
#include <errno.h>
#include <fcntl.h>
#include <fs.h>
#include <globalvar.h>
#include <init.h>
#include <ioctl.h>
#include <libbb.h>
#include <libfile.h>
#include <malloc.h>
#include <net.h>
#include <printk.h>
#include <state.h>
#include <stdio.h>
#include <watchdog.h>
#include <xfuncs.h>

#include <linux/err.h>
#include <linux/kernel.h>
#include <linux/list.h>
#include <linux/mtd/mtd-abi.h>
#include <linux/mtd/mtd.h>

#include <asm/unaligned.h>

/* list of all registered bootstate instances */
static LIST_HEAD(bootstate_list);

struct state_backend;

struct bootstate {
	struct device_d dev;
	const char *name;
	struct list_head list;
	struct list_head targets;
	struct list_head targets_unsorted;
	struct bootstate_backend *backend;
	bool dirty;
};

struct bootstate_backend {
	int (*load)(struct bootstate_backend *backend, struct bootstate *bootstate);
	int (*save)(struct bootstate_backend *backend, struct bootstate *bootstate);
	const char *name;
	const char *path;
};

struct bootstate_target {
	struct list_head list;
	struct list_head list_unsorted;

	/* state */
	unsigned int priority;
	unsigned int remaining_attempts;
	bool ok;

	/* spec */
	const char *name;
	unsigned int default_attempts;
};

static void pr_target(struct bootstate *bootstate, struct bootstate_target *target)
{
	printf("%s: target: name=%s prio=%u, ok=%d, rem=%u, def=%u\n",
	       bootstate->name, target->name, target->priority, target->ok,
	       target->remaining_attempts, target->default_attempts);
}

static struct bootstate *bootstate_new(const char *name)
{
	struct bootstate *bootstate;
	int ret;

	bootstate = xzalloc(sizeof(*bootstate));
	safe_strncpy(bootstate->dev.name, name, MAX_DRIVER_NAME);
	bootstate->name = bootstate->dev.name;
	bootstate->dev.id = DEVICE_ID_DYNAMIC;
	INIT_LIST_HEAD(&bootstate->targets);
	INIT_LIST_HEAD(&bootstate->targets_unsorted);

	ret = register_device(&bootstate->dev);
	if (ret) {
		free(bootstate);
		return ERR_PTR(ret);
	}

	list_add_tail(&bootstate->list, &bootstate_list);

	return bootstate;
}

static void bootstate_release(struct bootstate *bootstate)
{
	unregister_device(&bootstate->dev);
	free(bootstate);
}

static int bootstate_target_compare(struct list_head *a, struct list_head *b)
{
	struct bootstate_target *bootstate_a = list_entry(a, struct bootstate_target, list);
	struct bootstate_target *bootstate_b = list_entry(b, struct bootstate_target, list);

	/* order descending */
	return bootstate_a->priority >= bootstate_b->priority ? -1 : 1;
}

static void bootstate_target_add(struct bootstate *bootstate, struct bootstate_target *target)
{
	list_del(&target->list);
	list_add_sort(&target->list, &bootstate->targets, bootstate_target_compare);
}

static int bootstate_variable_read_u32(const struct bootstate *bootstate,
				       const char *name, uint32_t *out_val)
{
	char *var;
	int ret;

	var = asprintf("%s.%s.%s", bootstate->backend->path, bootstate->name, name);
	ret = getenv_uint(var, out_val);
	free(var);

	return ret;
}

static int bootstate_backend_variable_read_target_u32(const struct bootstate_backend *backend,
						      const struct bootstate *bootstate,
						      const struct bootstate_target *target,
						      const char *name, uint32_t *out_val)
{
	char *var;
	int ret;

	var = asprintf("%s.%s.%s.%s", backend->path, bootstate->name,
		       target->name, name);
	ret = getenv_uint(var, out_val);
	free(var);

	return ret;
}

static int bootstate_backend_variable_write_target_u32(const struct bootstate_backend *backend,
						       const struct bootstate *bootstate,
						       const struct bootstate_target *target,
						       const char *name, uint32_t in_val)
{
	char *var;
	char *val;
	int ret;

	var = asprintf("%s.%s.%s.%s", backend->path, bootstate->name,
		       target->name, name);
	val = asprintf("%d", in_val);
	ret = setenv(var, val);
	free(val);
	free(var);

	return ret;
}

static int bootstate_variable_nv_init_u32(const struct bootstate_backend *backend,
					  const struct bootstate *bootstate,
					  const struct bootstate_target *target,
					  const char *name)
{
	char *var;
	int ret;

	var = asprintf("%s.%s.%s", bootstate->name, target->name, name);
	ret = nvvar_add(var, "0");
	free(var);

	return ret;
}

static struct bootstate_target *bootstate_target_find(const struct bootstate *bootstate,
						      const char *name)
{
	struct bootstate_target *target;

	list_for_each_entry(target, &bootstate->targets, list) {
		if (!strcmp(target->name, name))
			return target;
	}

	return ERR_PTR(-ENOENT);
}

static int bootstate_target_from_node(struct bootstate *bootstate, const struct device_node *node, bool create)
{
	struct bootstate_target *target;
	char *name, *indexs;
	int ret;

	name = xstrdup(node->name);
	indexs = strchr(name, '@');
	if (indexs)
		*indexs++ = 0;

	if (create) {
		/* create*/
		target = xzalloc(sizeof(*target));

		target->name = xstrdup(name);
		list_add_tail(&target->list, &bootstate->targets);
		list_add_tail(&target->list_unsorted,
			      &bootstate->targets_unsorted);
	} else {
		target = bootstate_target_find(bootstate, name);
		if (IS_ERR(target)) {
			int ret = PTR_ERR(target);
			pr_err("no such boot target: %s: %s\n",
			       name, strerror(-ret));
			return ret;
		}
	}

	/* init */
	ret = of_property_read_u32(node, "default_attempts",
				   &target->default_attempts);
	if (ret)
		return ret;

	free(name);

	return 0;
}

static int bootstate_from_node(struct bootstate *bootstate,
			       const struct device_node *node, bool create)
{
	struct device_node *child;
	int ret;

	for_each_child_of_node(node, child) {
		ret = bootstate_target_from_node(bootstate, child, create);
		if (ret)
			return ret;
	}

	return 0;
}

static int bootstate_backend_load_one(const struct bootstate_backend *backend,
				      const struct bootstate *bootstate,
				      struct bootstate_target *target)
{
	uint32_t tmp;
	int ret;

	ret = bootstate_backend_variable_read_target_u32(backend, bootstate, target,
							 "remaining_attempts",
							 &target->remaining_attempts);
	if (ret)
		return ret;

	ret = bootstate_backend_variable_read_target_u32(backend, bootstate, target,
							 "priority", &target->priority);
	if (ret)
		return ret;

	ret = bootstate_backend_variable_read_target_u32(backend, bootstate, target,
							 "ok", &tmp);
	if (ret)
		return ret;

	target->ok = !!tmp;

	return ret;
}

static int bootstate_backend_load(struct bootstate_backend *backend,
				  struct bootstate *bootstate)
{
	struct bootstate_target *target;
	int ret;

	list_for_each_entry(target, &bootstate->targets_unsorted, list_unsorted) {
		ret = bootstate_backend_load_one(backend, bootstate, target);
		if (ret)
			return ret;
		bootstate_target_add(bootstate, target);
	}

	return 0;
}

static int bootstate_backend_save_one(const struct bootstate_backend *backend,
				      const struct bootstate *bootstate,
				      struct bootstate_target *target)
{
	int ret;

	ret = bootstate_backend_variable_write_target_u32(backend, bootstate, target,
							  "remaining_attempts",
							  target->remaining_attempts);
	if (ret)
		return ret;

	ret = bootstate_backend_variable_write_target_u32(backend, bootstate, target,
							  "priority", target->priority);
	if (ret)
		return ret;

	ret = bootstate_backend_variable_write_target_u32(backend, bootstate, target,
							  "ok", target->ok);
	if (ret)
		return ret;

	return 0;
}

static int bootstate_backend_save(const struct bootstate_backend *backend,
				  const struct bootstate *bootstate)
{
	struct bootstate_target *target;
	int ret;

	list_for_each_entry(target, &bootstate->targets, list) {
		ret = bootstate_backend_save_one(backend, bootstate, target);
		if (ret)
			return ret;
	}

	return 0;
}

static int bootstate_backend_nv_init_one(const struct bootstate_backend *backend,
					 const struct bootstate *bootstate,
					 struct bootstate_target *target)
{
	int ret;

	ret = bootstate_variable_nv_init_u32(backend, bootstate, target,
					     "remaining_attempts");
	if (ret)
		return ret;

	ret = bootstate_variable_nv_init_u32(backend, bootstate, target,
					     "priority");
	if (ret)
		return ret;

	ret = bootstate_variable_nv_init_u32(backend, bootstate, target,
					     "ok");
	if (ret)
		return ret;

	return 0;
}

static int bootstate_backend_nv_init(struct bootstate_backend *backend,
				     struct bootstate *bootstate)
{
	struct bootstate_target *target;
	int ret;

	list_for_each_entry(target, &bootstate->targets_unsorted, list_unsorted) {
		ret = bootstate_backend_nv_init_one(backend, bootstate, target);
		if (ret)
			return ret;
	}

	return 0;
}

static int bootstate_backend_nv_save(struct bootstate_backend *backend,
				     struct bootstate *bootstate)
{
	int ret;

	ret = bootstate_backend_save(backend, bootstate);
	if (ret)
		return ret;

	return envfs_save(NULL, NULL, 0);
}

static int bootstate_backend_nv_load(struct bootstate_backend *backend,
				     struct bootstate *bootstate)
{
	return bootstate_backend_load(backend, bootstate);
}

struct bootstate_backend_nv {
	struct bootstate_backend backend;
};

int bootstate_backend_nv(struct bootstate *bootstate)
{
	struct bootstate_backend_nv *backend_nv;
	struct bootstate_backend *backend;

	if (bootstate->backend)
		return -EBUSY;

	backend_nv = xzalloc(sizeof(*backend_nv));
	backend = &backend_nv->backend;

	backend->load = bootstate_backend_nv_load;
	backend->save = bootstate_backend_nv_save;
	backend->name = "nv";
	backend->path = "nv";

	bootstate->backend = backend;

	return bootstate_backend_nv_init(backend, bootstate);
}

struct bootstate_backend_state {
	struct bootstate_backend backend;
	struct state *state;
};

static int bootstate_backend_state_save(struct bootstate_backend *backend,
					struct bootstate *bootstate)
{
	struct bootstate_backend_state *backend_state =
		container_of(backend, struct bootstate_backend_state, backend);
	int ret;

	ret = bootstate_backend_save(backend, bootstate);
	if (ret)
		return ret;

	return state_save(backend_state->state);
}

static int bootstate_backend_state_load(struct bootstate_backend *backend,
					struct bootstate *bootstate)
{
	return bootstate_backend_load(backend, bootstate);
}

int bootstate_backend_state(struct bootstate *bootstate, const struct device_node *node)
{
	struct bootstate_backend_state *backend_state;
	struct bootstate_backend *backend;
	const struct device_node *state_node;

	if (bootstate->backend)
		return -EBUSY;

	backend_state = xzalloc(sizeof(*backend_state));
	backend = &backend_state->backend;

	backend->load = bootstate_backend_state_load;
	backend->save = bootstate_backend_state_save;
	backend->name = "state";

	bootstate->backend = backend;

	state_node = of_parse_phandle(node, "backend", 0);
	if (!state_node)
		return -EINVAL;

	backend_state->state = state_by_node(state_node);
	if (!backend_state->state)
		return -EINVAL;

	return state_get_name(backend_state->state, &backend->path);
}

/*
 * bootstate_new_from_node - create a new bootstate instance from a device_node
 *
 * @name	The name of the new bootstate instance
 * @node	The device_node describing the new bootstate instance
 */
struct bootstate *bootstate_new_from_node(const char *name, const struct device_node *node)
{
	struct bootstate *bootstate;
	int ret;

	pr_debug("%s: node=%s, name=%s\n", __func__, node->full_name, name);

	bootstate = bootstate_new(name);
	if (!bootstate)
		return ERR_PTR(-EINVAL);

	ret = bootstate_from_node(bootstate, node, true);
	if (ret) {
		bootstate_release(bootstate);
		return ERR_PTR(ret);
	}

	return bootstate;
}

/*
 * bootstate_by_name - find a bootstate instance by name
 *
 * @name	The name of the state instance
 */
struct bootstate *bootstate_by_name(const char *name)
{
	struct bootstate *bs;

	list_for_each_entry(bs, &bootstate_list, list) {
		if (!strcmp(name, bs->name))
			return bs;
	}

	return NULL;
}

/*
 * bootstate_load - load a bootstate from the backing store
 *
 * @bootstate	The state instance to load
 */
static int bootstate_load(struct bootstate *bootstate)
{
	int ret;

	if (!bootstate->backend)
		return -ENOSYS;

	ret = bootstate->backend->load(bootstate->backend, bootstate);
	if (ret)
		bootstate->dirty = 1;
	else
		bootstate->dirty = 0;

	return ret;
}

/*
 * bootstate_save - save a bootstate to the backing store
 *
 * @bootstate	The bootstate instance to save
 */
static int bootstate_save(struct bootstate *bootstate)
{
	int ret;

	if (!bootstate->dirty)
		return 0;

	if (!bootstate->backend)
		return -ENOSYS;

	ret = bootstate->backend->save(bootstate->backend, bootstate);
	if (ret)
		return ret;

	bootstate->dirty = 0;

	return 0;
}

void bootstate_info(void)
{
	struct bootstate *bootstate;

	printf("registered bootstate instances:\n");

	list_for_each_entry(bootstate, &bootstate_list, list) {
		printf("%-20s ", bootstate->name);
		printf("(backend: %s, path: %s)\n",
		       bootstate->backend->name, bootstate->backend->path);
	}
}

#define __BF(arg) [__BOOTCHOOSER_FLAG_##arg##_SHIFT] = __stringify(arg)

static const char * const bootstate_flags_str[] = {
	__BF(ATTEMPTS_KEEP),
	__BF(ATTEMPTS_DEC),
	__BF(ATTEMPTS_RESET),
	__BF(DEACTIVATE_ON_ZERO_ATTEMPTS),
	__BF(VERBOSE),
	__BF(DRYRUN),
	__BF(RETRY_WITH_DEC),
	__BF(WATCHDOG_ENABLE),
	__BF(WATCHDOG_TIMEOUT_FROM_STATE),
};

#undef __BF

#define pr(verbose, format, args...) \
	({ \
		(verbose) ? pr_info((format), ##args) : 0; \
	})

void _pr_flags(struct bootstate *bootstate, unsigned flags)
{
	int i;

	pr_info("%s: flags=0x%08x\n", bootstate->name, flags);

	for (i = 0; i < ARRAY_SIZE(bootstate_flags_str); i++) {
		if (flags & (1 << i))
			pr_info("%s: -> %s\n", bootstate->name,
				bootstate_flags_str[i]);
	}
}

#define pr_flags(verbose, bootstate, flags) \
	({ \
		(verbose) ? _pr_flags(bootstate, flags) : 0; \
	})

/*
 * bootstate_get_target - create a new state instance from a device_node
 *
 * @bootstate	the bootstate instance to work in
 * @flags	supported flags:
 *		BOOTCHOOSER_FLAG_VERBOSE
 *		BOOTCHOOSER_FLAG_ATTEMPTS_KEEP
 *		BOOTCHOOSER_FLAG_ATTEMPTS_DEC
 *		BOOTCHOOSER_FLAG_ATTEMPTS_RESET
 *		BOOTCHOOSER_FLAG_DEACTIVATE_ON_ZERO_ATTEMPTS
 * @target_out	a string to the choosen boot target is returned via
 *		this paramater
 */
int bootstate_get_target(struct bootstate *bootstate, unsigned flags, char **target_out)
{
	struct bootstate_target *target;
	int ret;
	bool found = false;
	bool v = flags & BOOTCHOOSER_FLAG_VERBOSE;

	pr_flags(v, bootstate, flags);

	ret = bootstate_load(bootstate);
	if (ret)
		return ret;

	if (flags & BOOTCHOOSER_FLAG_ATTEMPTS_RESET) {
		list_for_each_entry(target, &bootstate->targets, list) {
			if (target->priority ==  0)
				continue;

			target->remaining_attempts = target->default_attempts;
			bootstate->dirty = true;

			pr(v, "%s: target: name=%s setting rem to %d due to %s\n",
			   bootstate->name, target->name, target->default_attempts,
			   bootstate_flags_str[__BOOTCHOOSER_FLAG_ATTEMPTS_RESET_SHIFT]);
		}
		pr(v, "%s: --------\n", bootstate->name);
	}

	list_for_each_entry(target, &bootstate->targets, list) {
		pr_target(bootstate, target);

		if (found)
			continue;

		if (target->priority == 0) {
			pr(v, "%s: name=%s prio=%d - trying next\n",
			   bootstate->name, target->name, target->priority);
			continue;
		}

		if (target->remaining_attempts == 0) {
			pr(v, "%s: name=%s remaining attempts == 0 - trying next\n",
			   bootstate->name, target->name);
			continue;
		}

		if (flags & BOOTCHOOSER_FLAG_ATTEMPTS_DEC) {
			bootstate->dirty = true;
			target->remaining_attempts--;

			pr(v, "%s: name=%s decrementing remaining_attempts to %d due to %s\n",
			   bootstate->name, target->name,
			   target->remaining_attempts,
			   bootstate_flags_str[__BOOTCHOOSER_FLAG_ATTEMPTS_DEC_SHIFT]);

			if ((target->remaining_attempts == 0) &&
			    (flags & BOOTCHOOSER_FLAG_DEACTIVATE_ON_ZERO_ATTEMPTS)) {
				target->priority = 0;

				pr(v, "%s: name=%s deactivating target (setting priority = 0) due to %s\n",
				   bootstate->name, target->name,
				   bootstate_flags_str[__BOOTCHOOSER_FLAG_DEACTIVATE_ON_ZERO_ATTEMPTS_SHIFT]);
			}
		}

		found = true;
		*target_out = strdup(target->name);
		pr_debug("%s: selected target '%s'\n", __func__, target->name);
		if (!v)
			goto out;

		pr(v, "%s: --- other bootsources ---\n", bootstate->name);
	}

 out:
	bootstate_save(bootstate);

	if (!found)
		return -ENOENT;

	return 0;
}

int bootstate_bootchooser(char *name, unsigned flags, unsigned timeout)
{
	struct bootstate *bootstate;
	bool v = flags & BOOTCHOOSER_FLAG_VERBOSE;
	char *target;
	int ret;

	if (!name)
		name = "bootstate";

	bootstate = bootstate_by_name(name);
	if (!bootstate)
		return -ENOENT;

	if (flags & BOOTCHOOSER_FLAG_WATCHDOG_ENABLE) {
		if (flags & BOOTCHOOSER_FLAG_WATCHDOG_TIMEOUT_FROM_STATE) {
			ret = bootstate_variable_read_u32(bootstate, "watchdog_timeout",
							  &timeout);
			if (ret)
				return ret;
		}

		if (timeout != 0) {
			pr(v, "%s: starting watchdog with timeout=%ds\n",
			   __func__, timeout);

			ret = watchdog_set_timeout(timeout);
			if (ret)
				return ret;
		}
	}

	while (true) {
		char *cmd;

		ret = bootstate_get_target(bootstate, flags, &target);
		if (ret)
			return ret;

		cmd = asprintf("boot %s", target);
		free(target);
		pr_info("%srunning: %s...\n",
			flags & BOOTCHOOSER_FLAG_DRYRUN ? "not " : "", cmd);
		if (!(flags & BOOTCHOOSER_FLAG_DRYRUN))
			ret = run_command(cmd);
		free(cmd);

		if (flags & BOOTCHOOSER_FLAG_RETRY_WITH_DEC) {
			flags |= BOOTCHOOSER_FLAG_ATTEMPTS_DEC;
			flags &= ~(BOOTCHOOSER_FLAG_ATTEMPTS_RESET |
				   BOOTCHOOSER_FLAG_ATTEMPTS_KEEP);
			continue;
		}

		return ret;
	}

	return -ENOENT;
}