#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <fcntl.h>
#include <unistd.h>
#include <signal.h>
#include <ctype.h>
#include <errno.h>
#include <sys/epoll.h>
#include <sys/signalfd.h>

bool keep_running = true;

void asserno(bool ok, const char* name){
  if(ok) return;
  perror(name);
  exit(2);
}

typedef struct epoll_handler_s {
  int fd;
  void(*callback)(void*);
} epoll_handler_t;

typedef struct handle_signalfd_data_s {
  epoll_handler_t eph;
  pid_t child_pid;
} handle_signalfd_data_t;

void handle_signalfd_data(void* x){
  handle_signalfd_data_t*restrict d = x;
  struct signalfd_siginfo fdsi;
  asserno(read(d->eph.fd, &fdsi, sizeof(fdsi)) == sizeof(fdsi), "read");
  if(fdsi.ssi_signo == SIGCHLD){
    if((pid_t)fdsi.ssi_pid == d->child_pid){
      puts("Child program exited");
      keep_running = false;
    }
  }
}

typedef struct handle_pty_data_s {
  epoll_handler_t eph;
} handle_pty_data_t;

void handle_pty_data(void* x){
  handle_pty_data_t*restrict pty = x;
  char buf[256];
  int s = read(pty->eph.fd, buf, sizeof(buf));
  if(s<0) return;
  for(int i=0; i<s; i++)
    if(!isprint(buf[i]))
      buf[i] = '.';
  printf("Got data %d: %.*s\n", s, s, buf);
}

int main(int argc, char* argv[]){
  if(argc < 2){
    fprintf(stderr, "Usage: %s program [args]\n", argv[0]);
    return 1;
  }

  sigset_t mask;
  sigemptyset(&mask);
  sigaddset(&mask, SIGCHLD);
  asserno(sigprocmask(SIG_BLOCK, &mask, 0) != -1, "sigprocmask");
  int sfd = signalfd(-1, &mask, 0);
  asserno(sfd != -1, "signalfd");

  int epollfd = epoll_create1(0);
  asserno(epollfd != -1, "epoll_create1");

  int pts = open("/dev/ptmx", O_RDONLY);
  asserno(pts != -1, "open");
  asserno(grantpt(pts) != -1, "grantpt");
  asserno(unlockpt(pts) != -1, "unlockpt");

  pid_t pid = fork();
  asserno(pid != -1, "fork");
  if(!pid){
    setenv("MY_DEVICE_PATH", ptsname(pts), true);
    execvp(argv[1], argv+1);
    exit(1);
  }

  handle_signalfd_data_t eph_signalfd = {
    {
      .fd = sfd,
      .callback = handle_signalfd_data,
    },
    .child_pid = pid
  };
  asserno(epoll_ctl(epollfd, EPOLL_CTL_ADD, eph_signalfd.eph.fd, &(struct epoll_event){
    .events = EPOLLIN,
    .data.ptr = &eph_signalfd,
  }) != -1, "epoll_ctl");

  handle_pty_data_t eph_pty = {{
    .fd = pts,
    .callback = handle_pty_data,
  }};
  asserno(epoll_ctl(epollfd, EPOLL_CTL_ADD, eph_pty.eph.fd, &(struct epoll_event){
    .events = EPOLLIN,
    .data.ptr = &eph_pty,
  }) != -1, "epoll_ctl");

  enum { MAX_EVENTS = 10 };
  struct epoll_event events[MAX_EVENTS];
  while(keep_running){
    int nfds = epoll_wait(epollfd, events, MAX_EVENTS, -1);
    if(nfds == -1 && errno == EINTR)
      continue;
    asserno(nfds != -1, "epoll_wait");
    for(int i=0; i<nfds; i++){
      epoll_handler_t* eh = events[i].data.ptr;
      if(events[i].events & EPOLLERR){
        // We could delete it from the epoll set here, but for now, let's just exit
        fprintf(stderr, "fd %d has an error state!\n", eh->fd);
        keep_running = false;
        break;
      }
      eh->callback(eh);
    }
  }
  puts("Stopping program!");
}