#define _POSIX_SOURCE
#include <errno.h>
#include <regex.h>
#include <signal.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/un.h>
#include <sys/wait.h>
#include <unistd.h>

#define _MONITOR_REGEXP "unix:\\([-/.[:alnum:]]*\\)"
#define _SHUTDOWN_CMD "system_powerdown\n"
#define _GRACE_PERIOD 15
#define _MAX_SOCKET_TIMEOUT 5

struct _runtime {
  char *socket_path;
  int *pid;
};

struct _runtime _opts;

static void _check_code_and_exit(int status) {
  int result = WEXITSTATUS(status);
  if (result == 0) {
    puts("child process exited normally");
    exit(EXIT_SUCCESS);
  } else {
    printf("child process exited with code: %d\n", result);
    exit(EXIT_FAILURE);
  }
}

static void _send_shutdown(char *path) {
  int s, len;
  struct sockaddr_un remote = {
      .sun_family = AF_UNIX,
  };
  char str[100];
  if ((s = socket(AF_UNIX, SOCK_STREAM, 0)) == -1) {
    perror("socket");
    exit(1);
  }
  struct timeval tv;
  tv.tv_sec = _MAX_SOCKET_TIMEOUT;
  tv.tv_usec = 0;
  setsockopt(s, SOL_SOCKET, SO_RCVTIMEO, (const char *)&tv, sizeof tv);
  strcpy(remote.sun_path, _opts.socket_path);
  len = strlen(remote.sun_path) + sizeof(remote.sun_family);
  if (connect(s, (struct sockaddr *)&remote, len) == -1) {
    perror("connect");
    return;
  }
  printf("connected to %s\n", _opts.socket_path);
  if (send(s, _SHUTDOWN_CMD, strlen(_SHUTDOWN_CMD) + 1, 0) == -1) {
    perror("send");
    return;
  }

  for (;;) {
    len = recv(s, str, sizeof(str) - 1, 0);
    switch (len) {
    case -1:
      perror("recv");
      close(s);
      return;
    case 0:
      close(s);
      return;
    }
  }

  close(s);
}

static void _shutdown() {
  _send_shutdown(_opts.socket_path);
  printf("process %d has %d seconds to shutdown\n", *_opts.pid, _GRACE_PERIOD);
  int n_seconds = 0;
  while (n_seconds < _GRACE_PERIOD) {
    int stat;
    if (waitpid(*_opts.pid, &stat, WNOHANG) == -1) {
      printf("child process stopped, good bye!\n");
      _check_code_and_exit(stat);
    };
    sleep(2);
    n_seconds = n_seconds + 2;
  }
  printf("grace period ended, sending SIGTERM\n");
  kill(*_opts.pid, SIGTERM);
}

static void signal_handler(int sig) {
  printf("handling signal %d\n", sig);
  switch (sig) {
  case SIGTERM:
    _shutdown();
  case SIGINT:
    _shutdown();
  }
}

char *_get_arg_by_name(int argc, char *argv[], char *name) {
  char *res;
  int i;
  for (i = 0; i < argc; i++) {
    if (strcmp(argv[i], name) == 0 && (i < argc)) {
      res = argv[i + 1];
      return res;
    }
  }
  return NULL;
}

char *_get_arg_with_regexp(int argc, char *argv[], char *name,
                           char *expression) {
  char *arg_raw;
  arg_raw = _get_arg_by_name(argc, argv, name);
  if (arg_raw == NULL) {
    return NULL;
  }

  regex_t regex;
  regmatch_t groups[2];

  if (regcomp(&regex, expression, 0) != 0) {
    fprintf(stderr, "regexp failed to compile\n");
    exit(1);
  }

  char *result = malloc(strlen(arg_raw) + 1);

  if (regexec(&regex, arg_raw, 2, groups, 0) == 0) {
    if (groups[1].rm_so == (size_t)-1) {
      fprintf(stderr, "cannot match against flag:'%s'\n", arg_raw);
      exit(1);
    }
    strcpy(result, arg_raw);
    result[groups[1].rm_eo] = 0;
    result = result + groups[1].rm_so;
  } else {
    fprintf(stderr, "cannot match against flag:'%s'\n", arg_raw);
    exit(1);
  }

  regfree(&regex);
  return result;
}

void _do_fork(char *argv[], int *pid) {
  pid_t child_pid;
  if ((child_pid = fork()) == -1) {
    perror("fork");
    exit(1);
  }
  if (child_pid == 0) {
    setpgid(child_pid, child_pid);
    execvp(argv[0], argv);
  } else {
    *pid = child_pid;
  }
}

int main(int argc, char *argv[]) {
  char *monitor_path = NULL;
  monitor_path = _get_arg_with_regexp(argc, argv, "-monitor", _MONITOR_REGEXP);
  if (monitor_path == NULL) {
    printf("-monitor flag was not specified\n");
    return EXIT_FAILURE;
  }
  _opts.socket_path = monitor_path;
  char *_copy_args[argc];
  _copy_args[argc - 1] = NULL;
  for (int i = 1; i < argc; ++i) {
    _copy_args[i - 1] = argv[i];
  }

  signal(SIGTERM, signal_handler);
  signal(SIGINT, signal_handler);

  int pid = 0;
  int stat;

  _do_fork(_copy_args, &pid);
  _opts.pid = &pid;
  printf("monitoring pid %d\n", pid);
  while (waitpid(pid, &stat, 0) == -1) {
    if (errno != EINTR) {
      stat = -1;
      break;
    }
  }
  _check_code_and_exit(stat);
}