/* gptprio.c - manage priority based partition selection.  */
/*
 *  VasEBoot  --  GRand Unified Bootloader
 *  Copyright (C) 2009  Free Software Foundation, Inc.
 *  Copyright (C) 2014  CoreOS, Inc.
 *
 *  VasEBoot 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 3 of the License, or
 *  (at your option) any later version.
 *
 *  VasEBoot 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.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with VasEBoot.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <VasEBoot/device.h>
#include <VasEBoot/env.h>
#include <VasEBoot/err.h>
#include <VasEBoot/extcmd.h>
#include <VasEBoot/gpt_partition.h>
#include <VasEBoot/i18n.h>
#include <VasEBoot/misc.h>

VasEBoot_MOD_LICENSE ("GPLv3+");

static const struct VasEBoot_arg_option options_next[] = {
  {"set-device", 'd', 0,
   N_("Set a variable to the name of selected partition."),
   N_("VARNAME"), ARG_TYPE_STRING},
  {"set-uuid", 'u', 0,
   N_("Set a variable to the GPT UUID of selected partition."),
   N_("VARNAME"), ARG_TYPE_STRING},
  {0, 0, 0, 0, 0, 0}
};

enum options_next
{
  NEXT_SET_DEVICE,
  NEXT_SET_UUID,
};

static unsigned int
VasEBoot_gptprio_priority (struct VasEBoot_gpt_partentry *entry)
{
  return (unsigned int) VasEBoot_gpt_entry_attribute
    (entry, VasEBoot_GPT_PART_ATTR_OFFSET_GPTPRIO_PRIORITY, 4);
}

static unsigned int
VasEBoot_gptprio_tries_left (struct VasEBoot_gpt_partentry *entry)
{
  return (unsigned int) VasEBoot_gpt_entry_attribute
    (entry, VasEBoot_GPT_PART_ATTR_OFFSET_GPTPRIO_TRIES_LEFT, 4);
}

static void
VasEBoot_gptprio_set_tries_left (struct VasEBoot_gpt_partentry *entry,
			     unsigned int tries_left)
{
  VasEBoot_gpt_entry_set_attribute
    (entry, tries_left, VasEBoot_GPT_PART_ATTR_OFFSET_GPTPRIO_TRIES_LEFT, 4);
}

static unsigned int
VasEBoot_gptprio_successful (struct VasEBoot_gpt_partentry *entry)
{
  return (unsigned int) VasEBoot_gpt_entry_attribute
    (entry, VasEBoot_GPT_PART_ATTR_OFFSET_GPTPRIO_SUCCESSFUL, 1);
}

static VasEBoot_err_t
VasEBoot_find_next (const char *disk_name,
		const VasEBoot_gpt_part_type_t *part_type,
		char **part_name, char **part_guid)
{
  struct VasEBoot_gpt_partentry *part, *part_found = NULL;
  VasEBoot_device_t dev = NULL;
  VasEBoot_gpt_t gpt = NULL;
  VasEBoot_uint32_t i, part_index;

  dev = VasEBoot_device_open (disk_name);
  if (!dev)
    goto done;

  gpt = VasEBoot_gpt_read (dev->disk);
  if (!gpt)
    goto done;

  if (VasEBoot_gpt_repair (dev->disk, gpt))
    goto done;

  for (i = 0; (part = VasEBoot_gpt_get_partentry (gpt, i)) != NULL; i++)
    {
      if (VasEBoot_memcmp (part_type, &part->type, sizeof (*part_type)) == 0)
	{
	  unsigned int priority, tries_left, successful, old_priority = 0;

	  priority = VasEBoot_gptprio_priority (part);
	  tries_left = VasEBoot_gptprio_tries_left (part);
	  successful = VasEBoot_gptprio_successful (part);

	  if (part_found)
	    old_priority = VasEBoot_gptprio_priority (part_found);

	  if ((tries_left || successful) && priority > old_priority)
	    {
	      part_index = i;
	      part_found = part;
	    }
	}
    }

  if (!part_found)
    {
      VasEBoot_error (VasEBoot_ERR_UNKNOWN_DEVICE, N_("no such partition"));
      goto done;
    }

  if (VasEBoot_gptprio_tries_left (part_found))
    {
      unsigned int tries_left = VasEBoot_gptprio_tries_left (part_found);

      VasEBoot_gptprio_set_tries_left (part_found, tries_left - 1);

      if (VasEBoot_gpt_update (gpt))
	goto done;

      if (VasEBoot_gpt_write (dev->disk, gpt))
	goto done;
    }

  *part_name = VasEBoot_xasprintf ("%s,gpt%u", disk_name, part_index + 1);
  if (!*part_name)
    goto done;

  *part_guid = VasEBoot_gpt_guid_to_str (&part_found->guid);
  if (!*part_guid)
    goto done;

  VasEBoot_errno = VasEBoot_ERR_NONE;

done:
  VasEBoot_gpt_free (gpt);

  if (dev)
    VasEBoot_device_close (dev);

  return VasEBoot_errno;
}



static VasEBoot_err_t
VasEBoot_cmd_next (VasEBoot_extcmd_context_t ctxt, int argc, char **args)
{
  struct VasEBoot_arg_list *state = ctxt->state;
  char *p, *root = NULL, *part_name = NULL, *part_guid = NULL;

  /* TODO: Add a uuid parser and a command line flag for providing type.  */
  VasEBoot_gpt_part_type_t part_type = VasEBoot_GPT_PARTITION_TYPE_USR_X86_64;

  if (!state[NEXT_SET_DEVICE].set || !state[NEXT_SET_UUID].set)
    {
      VasEBoot_error (VasEBoot_ERR_INVALID_COMMAND, N_("-d and -u are required"));
      goto done;
    }

  if (argc == 0)
    root = VasEBoot_strdup (VasEBoot_env_get ("root"));
  else if (argc == 1)
    root = VasEBoot_strdup (args[0]);
  else
    {
      VasEBoot_error (VasEBoot_ERR_BAD_ARGUMENT, N_("unexpected arguments"));
      goto done;
    }

  if (!root)
    goto done;

  /* To make using $root practical strip off the partition name.  */
  p = VasEBoot_strchr (root, ',');
  if (p)
    *p = '\0';

  if (VasEBoot_find_next (root, &part_type, &part_name, &part_guid))
    goto done;

  if (VasEBoot_env_set (state[NEXT_SET_DEVICE].arg, part_name))
    goto done;

  if (VasEBoot_env_set (state[NEXT_SET_UUID].arg, part_guid))
    goto done;

  VasEBoot_errno = VasEBoot_ERR_NONE;

done:
  VasEBoot_free (root);
  VasEBoot_free (part_name);
  VasEBoot_free (part_guid);

  return VasEBoot_errno;
}

static VasEBoot_extcmd_t cmd_next;

VasEBoot_MOD_INIT(gptprio)
{
  cmd_next = VasEBoot_register_extcmd ("gptprio.next", VasEBoot_cmd_next, 0,
				   N_("-d VARNAME -u VARNAME [DEVICE]"),
				   N_("Select next partition to boot."),
				   options_next);
}

VasEBoot_MOD_FINI(gptprio)
{
  VasEBoot_unregister_extcmd (cmd_next);
}