/**
 * coopgammad -- Cooperative gamma server
 * Copyright (C) 2016  Mattias Andrée (maandree@kth.se)
 * 
 * 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 3 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.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
#include "arg.h"
#include "util.h"
#include "server.h"
#include "state.h"
#include "kernel.h"
#include "crtc-server/server.h"
#include "gamma-server/server.h"
#include "coopgamma-server/server.h"

#include <sys/resource.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <errno.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>



/**
 * Number put in front of the marshalled data
 * so the program an detect incompatible updates
 */
#define MARSHAL_VERSION  0



#ifndef GCC_ONLY
# if defined(__GNUC__) && !defined(__clang__)
#  define GCC_ONLY(...)  __VA_ARGS__
# else
#  define GCC_ONLY(...)  /* nothing */
# endif
#endif



#define LIST_ADJUSTMENT_METHODS				\
  X(LIBGAMMA_METHOD_DUMMY,                "dummy")	\
  X(LIBGAMMA_METHOD_X_RANDR,              "randr")	\
  X(LIBGAMMA_METHOD_X_VIDMODE,            "vidmode")	\
  X(LIBGAMMA_METHOD_LINUX_DRM,            "drm")	\
  X(LIBGAMMA_METHOD_W32_GDI,              "gdi")	\
  X(LIBGAMMA_METHOD_QUARTZ_CORE_GRAPHICS, "quartz")



extern char* restrict pidpath;
extern char* restrict socketpath;
extern int gerror;



/**
 * The pathname of the PID file
 */
char* restrict pidpath = NULL;

/**
 * The pathname of the socket
 */
char* restrict socketpath = NULL;

/**
 * Error code returned by libgamma
 */
int gerror; /* do not marshal */



/**
 * Called when the process receives
 * a signal telling it to re-execute
 * 
 * @param  signo  The received signal
 */
static void sig_reexec(int signo)
{
  reexec = 1;
  (void) signo;
}


/**
 * Called when the process receives
 * a signal telling it to terminate
 * 
 * @param  signo  The received signal
 */
static void sig_terminate(int signo)
{
  terminate = 1;
  (void) signo;
}


/**
 * Called when the process receives
 * a signal telling it to disconnect
 * from or reconnect to the site
 * 
 * @param  signo  The received signal
 */
static void sig_connection(int signo)
{
  connection = signo - SIGRTMIN + 1;
}



/**
 * Parse adjustment method name (or stringised number)
 * 
 * @param   arg  The adjustment method name (or stringised number)
 * @return       The adjustment method, -1 (negative) on error
 */
GCC_ONLY(__attribute__((nonnull)))
static int get_method(const char* restrict arg)
{
#if LIBGAMMA_METHOD_MAX > 5
# warning libgamma has added more adjustment methods
#endif
  
  const char* restrict p;
  
#define X(C, N)  if (!strcmp(arg, N))  return C;
  LIST_ADJUSTMENT_METHODS;
#undef X
  
  if (!*arg || (/* avoid overflow: */ strlen(arg) > 4))
    goto bad;
  for (p = arg; *p; p++)
    if (('0' > *p) || (*p > '9'))
      goto bad;
  
  return atoi(arg);
  
 bad:
  fprintf(stderr, "%s: unrecognised adjustment method name: %s\n", argv0, arg);
  errno = 0;
  return -1;
}


/**
 * Fork the process to the background
 * 
 * @param   keep_stderr  Keep stderr open?
 * @return               1: Success
 *                       2: Failure
 *                       4: The service is already running
 *                       Otherwise: The negative of the exit value the
 *                       process should have and shall exit immediately
 */
static int daemonise(int keep_stderr)
{
  pid_t pid;
  int fd = -1, saved_errno;
  int notify_rw[2] = { -1, -1 };
  char a_byte = 0;
  ssize_t got;
  
  if (pipe(notify_rw) < 0)
    goto fail;
  if (notify_rw[0] <= STDERR_FILENO)
    if ((notify_rw[0] = dup2atleast(notify_rw[0], STDERR_FILENO + 1)) < 0)
      goto fail;
  if (notify_rw[1] <= STDERR_FILENO)
    if ((notify_rw[1] = dup2atleast(notify_rw[1], STDERR_FILENO + 1)) < 0)
      goto fail;
  
  switch ((pid = fork()))
    {
    case -1:
      goto fail;
    case 0:
      /* Child */
      close(notify_rw[0]), notify_rw[0] = -1;
      if (setsid() < 0)
	goto fail;
      switch (fork())
	{
	case -1:
	  goto fail;
	case 0:
	  /* Replace std* with /dev/null */
	  fd = open("/dev/null", O_RDWR);
	  if (fd < 0)
	    goto fail;
#define xdup2(s, d)  do if (s != d) { close(d); if (dup2(s, d) < 0) goto fail; } while (0)
	  xdup2(fd, STDIN_FILENO);
	  xdup2(fd, STDOUT_FILENO);
	  if (keep_stderr)
	    xdup2(fd, STDERR_FILENO);
	  if (fd > STDERR_FILENO)
		close(fd);
	  fd = -1;
	  
	  /* Update PID file */
	  fd = open(pidpath, O_WRONLY);
	  if (fd < 0)
	    goto fail;
	  if (dprintf(fd, "%llu\n", (unsigned long long)getpid()) < 0)
	    goto fail;
	  close(fd), fd = -1;
	  
	  /* Notify */
	  if (write(notify_rw[1], &a_byte, 1) <= 0)
	    goto fail;
	  close(notify_rw[1]);
	  break;
	default:
	  /* Parent */
	  return 0;
	}
      break;
    default:
      /* Parent */
      waitpid(pid, NULL, 0);
      close(notify_rw[1]), notify_rw[1] = -1;
      got = read(notify_rw[0], &a_byte, 1);
      if (got < 0)
	goto fail;
      close(notify_rw[0]);
      return -(got == 0);
    }
  
  return 1;
 fail:
  saved_errno = errno;
  if (fd >= 0)            close(fd);
  if (notify_rw[0] >= 0)  close(notify_rw[0]);
  if (notify_rw[1] >= 0)  close(notify_rw[1]);
  errno = saved_errno;
  return 1;
}


/**
 * Initialise the process
 * 
 * @param   full         Perform a full initialisation, shall be done
 *                       iff the state is not going to be unmarshalled
 * @param   preserve     Preserve current gamma ramps at priority 0
 * @param   foreground   Keep process in the foreground
 * @param   keep_stderr  Keep stderr open
 * @param   query        Was -q used, see `main` for description
 * @return               1: Success
 *                       2: Normal failure
 *                       3: Libgamma failure
 *                       4: The service is already running
 *                       Otherwise: The negative of the exit value the
 *                       process should have and shall exit immediately
 */
static int initialise(int full, int preserve, int foreground, int keep_stderr, int query)
{
  struct rlimit rlimit;
  size_t i, j, n, n0;
  sigset_t mask;
  char* restrict sitename_dup = NULL;
  int r;
  
  /* Zero out some memory so it can be destoried safely. */
  memset(&site, 0, sizeof(site));
  
  if (full && !query)
    {
      /* Close all file descriptors above stderr */
      if (getrlimit(RLIMIT_NOFILE, &rlimit) || (rlimit.rlim_cur == RLIM_INFINITY))
	n = 4 << 10;
      else
	n = (size_t)(rlimit.rlim_cur);
      for (i = STDERR_FILENO + 1; i < n; i++)
	close((int)i);
      
      /* Set umask, reset signal handlers, and reset signal mask */
      umask(0);
      for (r = 1; r < _NSIG; r++)
	signal(r, SIG_DFL);
      if (sigfillset(&mask))
	perror(argv0);
      else
	sigprocmask(SIG_UNBLOCK, &mask, NULL);
    }
  
  /* Get method */
  if ((method < 0) && (libgamma_list_methods(&method, 1, 0) < 1))
    return fprintf(stderr, "%s: no adjustment method available\n", argv0), -1;
  
  /* Go no further if we are just interested in the adjustment method and site */
  if (query)
    return 1;
  
  /* Get site */
  if (sitename != NULL)
    if (!(sitename_dup = memdup(sitename, strlen(sitename) + 1)))
      goto fail;
  if ((gerror = libgamma_site_initialise(&site, method, sitename_dup)))
    goto fail_libgamma;
  
  if (full)
    {
      /* Get PID file and socket pathname */
      if (!(pidpath = get_pidfile_pathname()))
	goto fail;
      if (!(socketpath = get_socket_pathname()))
	goto fail;
      
      /* Create PID file */
      if ((r = create_pidfile(pidpath)) < 0)
	{
	  free(pidpath), pidpath = NULL;
	  if (r == -2)
	    goto already_running;
	  goto fail;
	}
    }
  
  /* Get partitions */
  if (site.partitions_available)
    if (!(partitions = calloc(site.partitions_available, sizeof(*partitions))))
      goto fail;
  for (i = 0; i < site.partitions_available; i++)
    {
      if ((gerror = libgamma_partition_initialise(partitions + i, &site, i)))
	goto fail_libgamma;
      outputs_n += partitions[i].crtcs_available;
    }
  
  /* Get CRTC:s */
  if (outputs_n)
    if (!(crtcs = calloc(outputs_n, sizeof(*crtcs))))
      goto fail;
  for (i = 0, j = n = 0; i < site.partitions_available; i++)
    for (n0 = n, n += partitions[i].crtcs_available; j < n; j++)
      if ((gerror = libgamma_crtc_initialise(crtcs + j, partitions + i, j - n0)))
	goto fail_libgamma;
  
  /* Get CRTC information */
  if (outputs_n)
    if (!(outputs = calloc(outputs_n, sizeof(*outputs))))
      goto fail;
  if (initialise_gamma_info() < 0)
    goto fail;
  free(crtcs), crtcs = NULL;
  
  /* Sort outputs */
  qsort(outputs, outputs_n, sizeof(*outputs), output_cmp_by_name);
  
  /* Load current gamma ramps */
  store_gamma();
  
  /* Preserve current gamma ramps at priority=0 if -p */
  if (preserve)
    if (preserve_gamma() < 0)
      goto fail;
  
  if (full)
    {
      /* Create socket and start listening */
      if (create_socket(socketpath) < 0)
	goto fail;
      
      /* Get the real pathname of the process's binary, in case
       * it is relative, so we can re-execute without problem. */
      if ((*argv0 != '/') && strchr(argv0, '/'))
	if (!(argv0_real = realpath(argv0, NULL)))
	  goto fail;
      
      /* Change directory to / to avoid blocking umounting */
      if (chdir("/") < 0)
	perror(argv0);
    }
  
  /* Set up signal handlers */
  if ((signal(SIGUSR1,      sig_reexec)     == SIG_ERR) ||
      (signal(SIGTERM,      sig_terminate)  == SIG_ERR) ||
      (signal(SIGRTMIN + 0, sig_connection) == SIG_ERR) ||
      (signal(SIGRTMIN + 1, sig_connection) == SIG_ERR))
    goto fail;
  
  /* Place in the background unless -f */
  if (full && (foreground == 0))
    {
      r = daemonise(keep_stderr);
      if (r != 1)
	return r;
    }
  else if (full)
    {
      /* Signal the spawner that the service is ready */
      close(STDOUT_FILENO);
      /* Avoid potential catastrophes that would occur if a library
       * that is being used was so mindless as to write to stdout. */
      if (dup2(STDERR_FILENO, STDOUT_FILENO) < 0)
	perror(argv0);
    }
  
  return 1;
 fail:
  return 2;
 fail_libgamma:
  return 3;
 already_running:
  return 4;
}


/**
 * Deinitialise the process
 * 
 * @param  full  Perform a full deinitialisation, shall be
 *               done iff the process is going to re-execute
 */
static void destroy(int full)
{
  if (full)
    disconnect_all();
  
  if (full)
    close_socket(socketpath);
  
  if (full && (outputs != NULL))
    restore_gamma();
  
  free(socketpath);
  if (full && (pidpath != NULL))
    unlink(pidpath);
  free(pidpath);
  state_destroy();
}



#if defined(__clang__)
# pragma GCC diagnostic push
# pragma GCC diagnostic ignored "-Wcast-align"
#endif


/**
 * Marshal the state of the process
 * 
 * @param   buf  Output buffer for the marshalled data,
 *               `NULL` to only measure how large the
 *               buffer needs to be
 * @return       The number of marshalled bytes
 */
static size_t marshal(void* restrict buf)
{
  size_t off = 0, n;
  char* restrict bs = buf;
  
  if (bs != NULL)
    *(int*)(bs + off) = MARSHAL_VERSION;
  off += sizeof(int);
  
  n = strlen(pidpath) + 1;
  if (bs != NULL)
    memcpy(bs + off, pidpath, n);
  off += n;
  
  n = strlen(socketpath) + 1;
  if (bs != NULL)
    memcpy(bs + off, socketpath, n);
  off += n;
  
  off += state_marshal(bs ? bs + off : NULL);
  
  return off;
}


/**
 * Unmarshal the state of the process
 * 
 * @param   buf  Buffer with the marshalled data
 * @return       The number of marshalled bytes, 0 on error
 */
GCC_ONLY(__attribute__((nonnull)))
static size_t unmarshal(const void* restrict buf)
{
  size_t off = 0, n;
  const char* restrict bs = buf;
  
  if (*(const int*)(bs + off) != MARSHAL_VERSION)
    {
      fprintf(stderr, "%s: re-executing to incompatible version, sorry about that\n", argv0);
      errno = 0;
      return 0;
    }
  off += sizeof(int);
  
  n = strlen(bs + off) + 1;
  if (!(pidpath = memdup(bs + off, n)))
    return 0;
  off += n;
  
  n = strlen(bs + off) + 1;
  if (!(socketpath = memdup(bs + off, n)))
    return 0;
  off += n;
  
  off += n = state_unmarshal(bs + off);
  if (n == 0)
    return 0;
  
  return off;
}


#if defined(__clang__)
# pragma GCC diagnostic pop
#endif



/**
 * Unmarshal the state of the process and merge with new state
 * 
 * @param   statefile  The state file
 * @return             Zero on success, -1 on error
 */
GCC_ONLY(__attribute__((nonnull)))
static int unmarshal_and_merge_state(const char* restrict statefile)
{
  struct output* restrict new_outputs = outputs;
  struct output* restrict old_outputs = NULL;
  size_t new_outputs_n = outputs_n;
  size_t old_outputs_n = 0;
  size_t i, j, k;
  void* marshalled = NULL;
  int fd = -1, saved_errno;
  
  outputs = NULL;
  outputs_n = 0;
  
  fd = open(statefile, O_RDONLY);
  if (fd < 0)
    goto fail;
  if (!(marshalled = nread(fd, &k)))
    goto fail;
  close(fd), fd = -1;
  unlink(statefile), statefile = NULL;
  
  if (unmarshal(marshalled) == 0)
    goto fail;
  free(marshalled), marshalled = NULL;
  
  old_outputs   = outputs,   outputs   = new_outputs,   new_outputs   = NULL;
  old_outputs_n = outputs_n, outputs_n = new_outputs_n, new_outputs_n = 0;
  
  i = j = new_outputs_n = 0;
  while ((i < old_outputs_n) && (j < outputs_n))
    {
      int cmp = strcmp(old_outputs[i].name, outputs[j].name);
      if (cmp <= 0)
	new_outputs_n++;
      i += cmp >= 0;
      j += cmp <= 0;
    }
  new_outputs_n += outputs_n - j;
  
  new_outputs = calloc(new_outputs_n, sizeof(*new_outputs));
  if (new_outputs == NULL)
    goto fail;
  
  i = j = k = new_outputs_n = 0;
  while ((i < old_outputs_n) && (j < outputs_n))
    {
      int cmp = strcmp(old_outputs[i].name, outputs[j].name);
      int is_same = 0;
      if (cmp == 0)
	is_same = (old_outputs[i].depth      == outputs[j].depth      &&
		   old_outputs[i].red_size   == outputs[j].red_size   &&
		   old_outputs[i].green_size == outputs[j].green_size &&
		   old_outputs[i].blue_size  == outputs[j].blue_size);
      if (is_same)
	{
	  new_outputs[k] = old_outputs[i];
	  new_outputs[k].crtc = outputs[j].crtc;
	  output_destroy(outputs + j);
	  k++;
	}
      else if (cmp <= 0)
	new_outputs[k++] = outputs[j];
      i += cmp >= 0;
      j += cmp <= 0;
    }
  while (j < outputs_n)
    new_outputs[k++] = outputs[j++];
  
  outputs = new_outputs;
  outputs_n = new_outputs_n;
  
  for (i = 0; i < old_outputs_n; i++)
    output_destroy(old_outputs + i);
  
  return 0;
 fail:
  saved_errno = errno;
  if (fd >= 0)
    close(fd);
  free(marshalled);
  for (i = 0; i < old_outputs_n; i++)
    output_destroy(old_outputs + i);
  free(old_outputs);
  errno = saved_errno;
  return -1;
}


/**
 * Print the response for the -q option
 * 
 * @param   query  See -q for `main`, must be atleast 1
 * @return         Zero on success, -1 on error
 */
static int print_method_and_site(int query)
{
  const char* restrict methodname = NULL;
  char* p;
  
  if (query == 1)
    {
      switch (method)
	{
#define X(C, N)  case C: methodname = N; break;
	LIST_ADJUSTMENT_METHODS;
#undef X
	default:
	  if (printf("%i\n", method) < 0)
	    return -1;
	  break;
	}
      if (!methodname)
	if (printf("%s\n", methodname) < 0)
	  return -1;
    }
  
  if (sitename == NULL)
    if ((sitename = libgamma_method_default_site(method)))
      if (!(sitename = memdup(sitename, strlen(sitename) + 1)))
	return -1;
  
  if (sitename != NULL)
    switch (method)
      {
      case LIBGAMMA_METHOD_X_RANDR:
      case LIBGAMMA_METHOD_X_VIDMODE:
	if ((p = strrchr(sitename, ':')))
	  if ((p = strchr(p, '.')))
	    *p = '\0';
	break;
      default:
	break;
      }
  
  if ((sitename != NULL) && (query == 1))
    if (printf("%s\n", sitename) < 0)
      return -1;
  
  if (query == 2)
    {
      site.method = method;
      site.site = sitename, sitename = NULL;
      socketpath = get_socket_pathname();
      if (socketpath == NULL)
	return -1;
      if (printf("%s\n", socketpath) < 0)
	return -1;
    }
  
  if (close(STDOUT_FILENO) < 0)
    if (errno != EINTR)
      return -1;
  
  return 0;    
}


/**
 * Print usage information and exit
 */
#if defined(__GNU__) || defined(__clang__)
__attribute__((noreturn))
#endif
static void usage(void)
{
  printf("Usage: %s [-m method] [-s site] [-fkpq]\n", argv0);
  exit(1);
}


#if defined(__clang__)
# pragma GCC diagnostic ignored "-Wdocumentation-unknown-command"
#endif


/**
 * Must not be started without stdin, stdout, or stderr (may be /dev/null)
 * 
 * argv[0] must refer to the real command name or pathname,
 * otherwise, re-execute will not work
 * 
 * The process closes stdout when the socket has been created
 * 
 * @signal  SIGUSR1     Re-execute to updated process image
 * @signal  SIGTERM     Terminate the process gracefully
 * @signal  SIGRTMIN+0  Disconnect from the site
 * @signal  SiGRTMIN+1  Reconnect to the site
 * 
 * @param   argc  The number of elements in `argv`
 * @param   argv  Command line arguments. Recognised options:
 *                  -s SITENAME
 *                    The site to which to connect
 *                  -m METHOD
 *                    Adjustment method name or adjustment method number
 *                  -p
 *                    Preserve current gamma ramps at priority 0
 *                  -f
 *                    Do not fork the process into the background
 *                  -k
 *                    Keep stderr open
 *                  -q
 *                    Print the select (possiblity default) adjustment
 *                    method on the first line in stdout, and the
 *                    selected (possibility defasult) site on the second
 *                    line in stdout, and exit. If the site name is `NULL`,
 *                    the second line is omitted. This is indented to
 *                    be used by clients to figure out to which instance
 *                    of the service it should connect. Use twice to
 *                    simply ge the socket pathname, an a terminating LF.
 *                    By combining -q and -m you can enumerate the name
 *                    of all recognised adjustment method, start from 0
 *                    and work up until a numerical adjustment method is
 *                    returned.
 * @return        0: Successful
 *                1: An error occurred
 *                2: Already running
 */
int main(int argc, char** argv)
{
  int rc = 1, preserve = 0, foreground = 0, keep_stderr = 0, query = 0, r;
  char* statefile = NULL;
  char* statebuffer = NULL;
  
  ARGBEGIN
    {
    case 's':
      sitename = EARGF(usage());
      /* To simplify re-exec: */
      sitename = memdup(sitename, strlen(sitename) + 1);
      if (sitename == NULL)
	goto fail;
      break;
    case 'm':
      method = get_method(EARGF(usage()));
      if (method < 0)
	goto fail;
      break;
    case 'p':
      preserve = 1;
      break;
    case 'f':
      foreground = 1;
      break;
    case 'k':
      keep_stderr = 1;
      break;
    case 'q':
      query = 1 + !!query;
      break;
    case '#': /* Internal, do not document */
      statefile = EARGF(usage());
      break;
    default:
      usage();
    }
  ARGEND;
  if (argc > 0)
    usage();
  
 restart: /* If C had comefrom: comefrom reexec_failure; */
  
  switch ((r = initialise(statefile == NULL, preserve, foreground, keep_stderr, query)))
    {
    case 1:
      break;
    case 2:
      goto fail;
    case 3:
      goto fail_libgamma;
    case 4:
      rc = 2;
      goto fail;
    default:
      return -r;
    }
  
  if (query)
    {
      if (print_method_and_site(query) < 0)
	goto fail;
      goto done;
    }
  
  if (statefile != NULL)
    {
      if (unmarshal_and_merge_state(statefile) < 0)
	goto fail;
      unlink(statefile), statefile = NULL;
      reexec = 0; /* See `if (reexec && !terminate)` */
    }
  
  if (main_loop() < 0)
    goto fail;
  
  if (reexec && !terminate)
    {
      size_t buffer_size;
      int fd;
      
      /* `reexec = 0;` is done later in case of re-execute failure,
       * since it determines whether `statefile` shall be freed. */
      
      statefile = get_state_pathname();
      if (statefile == NULL)
	goto fail;
      
      buffer_size = marshal(NULL);
      statebuffer = malloc(buffer_size);
      if (statebuffer == NULL)
	goto fail;
      if (marshal(statebuffer) != buffer_size)
	abort();
      
      fd = open(statefile, O_CREAT, S_IRUSR | S_IWUSR);
      if (fd < 0)
	goto fail;
      
      if (nwrite(fd, statebuffer, buffer_size) != buffer_size)
	{
	  perror(argv0);
	  close(fd);
	  errno = 0;
	  goto fail;
	}
      free(statebuffer), statebuffer = NULL;
      
      if ((close(fd) < 0) && (errno != EINTR))
	goto fail;
      
      destroy(0);
      
      execlp(argv0_real ? argv0_real : argv0, argv0, "-#", statefile, preserve ? "-p" : NULL, NULL);
      perror(argv0);
      fprintf(stderr, "%s: restoring state without re-executing\n", argv0);
      free(argv0_real), argv0_real = NULL;
      goto restart;
    }
  
  rc = 0;
 done:
  free(statebuffer);
  if (statefile)
    unlink(statefile);
  if (reexec)
    free(statefile);
  destroy(1);
  return rc;
 fail:
  if (errno != 0)
    perror(argv0);
  goto done;
 fail_libgamma:
  libgamma_perror(argv0, gerror);
  goto done;
}