/*
 *  GRUB  --  GRand Unified Bootloader
 *  Copyright (C) 2013  Free Software Foundation, Inc.
 *
 *  GRUB 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.
 *
 *  GRUB 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 GRUB.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <grub/types.h>
#include <grub/misc.h>
#include <grub/mm.h>
#include <grub/err.h>
#include <grub/dl.h>
#include <grub/command.h>
#include <grub/i18n.h>
#include <grub/device.h>
#include <grub/mm.h>
#include <grub/fs.h>
#include <grub/env.h>
#include <grub/file.h>

GRUB_MOD_LICENSE ("GPLv3+");

static const char *modnames_def[] = { "pata", "ahci", "usbms", "ohci", "uhci", "ehci" };

static grub_err_t
get_uuid (const char *name, char **uuid)
{
  grub_device_t dev;
  grub_fs_t fs = 0;

  dev = grub_device_open (name);
  if (!dev)
    return grub_errno;
  if (dev)
    fs = grub_fs_probe (dev);
  if (!fs)
    {
      grub_device_close (dev);
      return grub_errno;
    }
  if (!fs->uuid || fs->uuid (dev, uuid))
    {
      grub_device_close (dev);

      if (!grub_errno)
	grub_error (GRUB_ERR_NOT_IMPLEMENTED_YET,
		    N_("%s does not support UUIDs"), fs->name);

      return grub_errno;
    }
  grub_device_close (dev);
  return GRUB_ERR_NONE;
}

struct search_ctx
{
  char *root_uuid;
  char *prefix_uuid;
  const char *prefix_path;
  int prefix_found, root_found;
};

static int
iterate_device (const char *name, void *data)
{
  struct search_ctx *ctx = data;
  char *cur_uuid;

  if (get_uuid (name, &cur_uuid))
    {
      if (grub_errno == GRUB_ERR_UNKNOWN_FS)
	grub_errno = 0;
      grub_print_error ();
      return 0;
    }
  if (grub_strcasecmp (cur_uuid, ctx->prefix_uuid) == 0)
    {
      char *prefix;
      prefix = grub_xasprintf ("(%s)/%s", name, ctx->prefix_path);
      grub_env_set ("prefix", prefix);
      grub_free (prefix);
      ctx->prefix_found = 1;
    }
  if (grub_strcasecmp (cur_uuid, ctx->root_uuid) == 0)
    {
      grub_env_set ("root", name);
      ctx->root_found = 1;
    }
  return ctx->prefix_found && ctx->root_found;
}

static grub_err_t
grub_cmd_nativedisk (grub_command_t cmd __attribute__ ((unused)),
		     int argc, char **args_in)
{
  char *uuid_root = 0, *uuid_prefix, *prefdev = 0;
  const char *prefix = 0;
  const char *path_prefix = 0;
  int mods_loaded = 0;
  grub_dl_t *mods;
  struct search_ctx ctx;
  const char **args;
  grub_fs_autoload_hook_t saved_autoload;
  int i;

  if (argc == 0)
    {
      argc = ARRAY_SIZE (modnames_def);
      args = modnames_def;
    }
  else
    args = (const char **) args_in;

  prefix = grub_env_get ("prefix");

  if (! prefix)
    return grub_error (GRUB_ERR_FILE_NOT_FOUND, N_("variable `%s' isn't set"), "prefix");

  if (prefix)
    path_prefix = (prefix[0] == '(') ? grub_strchr (prefix, ')') : NULL;
  if (path_prefix)
    path_prefix++;
  else
    path_prefix = prefix;

  mods = grub_malloc (argc * sizeof (mods[0]));
  if (!mods)
    return grub_errno;

  if (get_uuid (NULL, &uuid_root))
    return grub_errno;

  prefdev = grub_file_get_device_name (prefix);
  if (grub_errno)
    {
      grub_print_error ();
      prefdev = 0;
    }

  if (get_uuid (prefdev, &uuid_prefix))
    {
      grub_free (uuid_root);
      return grub_errno;
    }

  for (mods_loaded = 0; mods_loaded < argc; mods_loaded++)
    {
      char *filename;
      grub_dl_t mod;
      grub_file_t file = NULL;
      grub_ssize_t size;
      void *core = 0;

      mod = grub_dl_get (args[mods_loaded]);
      if (mod)
	{
	  mods[mods_loaded] = 0;
	  continue;
	}

      filename = grub_xasprintf ("%s/" GRUB_TARGET_CPU "-" GRUB_PLATFORM "/%s.mod",
				 prefix, args[mods_loaded]);
      if (! filename)
	goto fail;

      file = grub_file_open (filename);
      grub_free (filename);
      if (! file)
	goto fail;

      size = grub_file_size (file);
      core = grub_malloc (size);
      if (! core)
	{
	  grub_file_close (file);
	  goto fail;
	}

      if (grub_file_read (file, core, size) != (grub_ssize_t) size)
	{
	  grub_file_close (file);
	  grub_free (core);
	  goto fail;
	}

      grub_file_close (file);

      mods[mods_loaded] = grub_dl_load_core_noinit (core, size);
      if (! mods[mods_loaded])
	goto fail;
    }

  for (i = 0; i < argc; i++)
    if (mods[i])
      grub_dl_init (mods[i]);

  /* No need to autoload FS since obviously we already have the necessary fs modules.  */
  saved_autoload = grub_fs_autoload_hook;
  grub_fs_autoload_hook = 0;

  ctx.root_uuid = uuid_root;
  ctx.prefix_uuid = uuid_prefix;
  ctx.prefix_path = path_prefix;
  ctx.prefix_found = 0;
  ctx.root_found = 0;

  /* FIXME: try to guess the correct values.  */
  grub_device_iterate (iterate_device, &ctx);

  grub_fs_autoload_hook = saved_autoload;

  grub_free (uuid_root);
  grub_free (uuid_prefix);

  return GRUB_ERR_NONE;

 fail:
  grub_free (uuid_root);
  grub_free (uuid_prefix);

  for (i = 0; i < mods_loaded; i++)
    if (mods[i])
      {
	mods[i]->fini = 0;
	grub_dl_unload (mods[i]);
      }
  return grub_errno;
}

static grub_command_t cmd;

GRUB_MOD_INIT(nativedisk)
{
  cmd = grub_register_command ("nativedisk", grub_cmd_nativedisk, "[MODULE1 MODULE2 ...]",
			       N_("Switch to native disk drivers. If no modules are specified default set (pata,ahci,usbms,ohci,uhci,ehci) is used"));
}

GRUB_MOD_FINI(nativedisk)
{
  grub_unregister_command (cmd);
}