#include <error.h>
#include <errno.h>
#include <stdlib.h>
#include <unistd.h>
#include <libintl.h>
#include <stdio.h>
#include <string.h>

#ifndef STAP_SCRIPT_FILENAME
#error "STAP_SCRIPT_FILENAME is required"
#endif

#define STRINGIFY(lit) #lit
#define _(x) gettext (x)

static int opt_d, opt_f, opt_q, opt_p;

/* Script filename.  */
static char *opt_Z = STRINGIFY (STAP_SCRIPT_FILENAME);

static void *
xmalloc (size_t size)
{
  void *retval = malloc (size);

  if (retval == NULL)
    error (EXIT_FAILURE, errno, _("memory allocation of %zu bytes"), size);

  return retval;
}

static void
usage (int status)
{
  printf ("\
usage: staptrace [-dfq] [-o file] [-p pid] [command [arg ...]]\n\
                 [-Z <stap script>]\n\
\n\
Default -Z is: %s\n\
",
	  opt_Z);
  exit (status);
}

static void
dump_args (char **args)
{
  char **argp;

  for (argp = args; *argp; argp++)
    {
      if (argp > args)
	fputs (", ", stderr);
      fprintf (stderr, "'%s'", *argp);
    }
  fputc ('\n', stderr);
}

int
main (int argc, char **argv)
{
  char **args, **argp;

  args = xmalloc (sizeof (*args) * argc * 2);
  argp = args;
  *argp++ = argv[0];

  for (;;)
    {
      int i;

      i = getopt (argc, argv, "+dfqo:p:hZ:");
      if (i == EOF)
	break;
      switch (i)
	{
	case 'd':
	  opt_d = 1;
	  break;
	case 'f':
	  if (opt_f)
	    error (EXIT_FAILURE, 0, _("-ff is unsupported"));
	  *argp++ = "-G";
	  *argp++ = "opt_f=1";
	  opt_f = 1;
	  break;
	case 'q':
	  *argp++ = "-G";
	  *argp++ = "opt_q=1";
	  opt_q = 1;
	  break;
	case 'o':
	  *argp++ = "-o";
	  *argp++ = optarg;
	  break;
	case 'p':
	  *argp++ = "-x";
	  *argp++ = optarg;
	  opt_p = 1;
	  break;
	case 'Z':
	  opt_Z = optarg;
	  break;
	case 'h':
	  usage (EXIT_SUCCESS);
	default:
	  error (0, 0, _("Invalid option '%c'"), i);
	  usage (EXIT_FAILURE);
	}
    }

  if ((optind == argc) == !opt_p)
    usage (EXIT_FAILURE);

  if (optind < argc)
    {
      size_t len;
      int i;
      char *d;
      const char *s;

      *argp++ = "-c";

      len = 32;
      for (i = optind; i < argc; i++)
	len += 3 + strlen (argv[i]) * 4;

      d = *argp++ = xmalloc (len);
      d = stpcpy (d, "exec");
      for (i = optind; i < argc; i++)
	{
	  *d++ = ' ';
	  *d++ = '\'';
	  for (s = argv[optind]; *s; s++)
	    switch (*s)
	      {
	      case '\'':
		d = stpcpy (d, "'\\''");
		break;
	      case '\\':
		*d++ = '\\';
		/* FALLTHRU */
	      default:
		*d++ = *s;
	      }
	  *d++ = '\'';
	}
      *d++ = 0;
    }
  *argp++ = opt_Z;
  *argp = NULL;

  if (opt_d)
    dump_args (args);
  execvp ("stap", args);

  if (!opt_d)
    dump_args (args);
  error (EXIT_FAILURE, errno, "Error executing stap");
}