#include <errno.h>
#include <signal.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/select.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>

#if defined(__FreeBSD__) || defined(__DragonFly__) || defined(__NetBSD__) || defined(__OpenBSD__)
#include <libutil.h>
#else
#include <pty.h>
#endif

static bool wait_for_substr(int fd, const char *needle, int timeout_ms) {
  char buf[8192];
  size_t len = 0;
  buf[0] = '\0';

  const int step_ms = 20;
  int remaining = timeout_ms;

  while (remaining > 0) {
    fd_set rfds;
    FD_ZERO(&rfds);
    FD_SET(fd, &rfds);

    struct timeval tv;
    tv.tv_sec = step_ms / 1000;
    tv.tv_usec = (step_ms % 1000) * 1000;

    int rc = select(fd + 1, &rfds, NULL, NULL, &tv);
    if (rc < 0) {
      if (errno == EINTR) {
        continue;
      }
      return false;
    }
    if (rc == 0) {
      remaining -= step_ms;
      continue;
    }

    if (!FD_ISSET(fd, &rfds)) {
      remaining -= step_ms;
      continue;
    }

    ssize_t n = read(fd, buf + len, sizeof(buf) - 1 - len);
    if (n <= 0) {
      return false;
    }
    len += (size_t)n;
    buf[len] = '\0';

    if (strstr(buf, needle)) {
      return true;
    }

    if (len > sizeof(buf) / 2) {
      const size_t keep = sizeof(buf) / 3;
      memmove(buf, buf + (len - keep), keep);
      len = keep;
      buf[len] = '\0';
    }
  }

  return false;
}

static void drain_fd(int fd) {
  char tmp[4096];
  for (;;) {
    fd_set rfds;
    FD_ZERO(&rfds);
    FD_SET(fd, &rfds);

    struct timeval tv;
    tv.tv_sec = 0;
    tv.tv_usec = 0;

    int rc = select(fd + 1, &rfds, NULL, NULL, &tv);
    if (rc <= 0 || !FD_ISSET(fd, &rfds)) {
      return;
    }
    ssize_t n = read(fd, tmp, sizeof(tmp));
    if (n <= 0) {
      return;
    }
  }
}

static bool wait_for_exit(int fd, pid_t pid, int timeout_ms, int *status_out) {
  const int step_ms = 10;
  int remaining = timeout_ms;

  while (remaining > 0) {
    int st = 0;
    pid_t w = waitpid(pid, &st, WNOHANG);
    if (w == pid) {
      if (status_out) {
        *status_out = st;
      }
      return true;
    }
    if (w < 0) {
      if (errno == EINTR) {
        continue;
      }
      return false;
    }
    drain_fd(fd);
    usleep(step_ms * 1000);
    remaining -= step_ms;
  }

  return false;
}

int main(int argc, char **argv) {
  if (argc < 2 || !argv[1] || !*argv[1]) {
    fprintf(stderr, "usage: %s /path/to/flowlog\n", argv[0]);
    return 2;
  }

  const char *bin = argv[1];

  int master = -1;
  pid_t pid = forkpty(&master, NULL, NULL, NULL);
  if (pid < 0) {
    perror("forkpty");
    return 2;
  }

  if (pid == 0) {
    execl(bin, bin, NULL);
    perror("exec");
    _exit(127);
  }

  if (master < 0) {
    kill(pid, SIGKILL);
    (void)waitpid(pid, NULL, 0);
    fprintf(stderr, "no pty master\n");
    return 2;
  }

  if (!wait_for_substr(master, "?- ", 5000)) {
    kill(pid, SIGKILL);
    (void)waitpid(pid, NULL, 0);
    fprintf(stderr, "did not see prompt\n");
    return 1;
  }

  // Ctrl-T: Flowlog SIGINFO keybinding (BSD SIGINFO equivalent).
  const unsigned char ctrl_t = 0x14;
  if (write(master, &ctrl_t, 1) != 1) {
    kill(pid, SIGKILL);
    (void)waitpid(pid, NULL, 0);
    fprintf(stderr, "failed to write Ctrl-T\n");
    return 1;
  }

  // In raw mode, the terminal must still translate '\n' to CRLF, otherwise
  // the status report "staircases" to the right.
  if (!wait_for_substr(master, "\r\nFlowlog SIGINFO (status)\r\n", 5000)) {
    kill(pid, SIGKILL);
    (void)waitpid(pid, NULL, 0);
    fprintf(stderr, "siginfo output missing CRLF framing\n");
    return 1;
  }

  // Exit via Ctrl-D at a fresh prompt.
  const unsigned char ctrl_d = 0x04;
  if (write(master, &ctrl_d, 1) != 1) {
    kill(pid, SIGKILL);
    (void)waitpid(pid, NULL, 0);
    fprintf(stderr, "failed to write Ctrl-D\n");
    return 1;
  }

  int st = 0;
  if (!wait_for_exit(master, pid, 3000, &st)) {
    kill(pid, SIGKILL);
    (void)waitpid(pid, NULL, 0);
    fprintf(stderr, "timeout waiting for exit\n");
    return 1;
  }

  if (WIFEXITED(st) && WEXITSTATUS(st) == 0) {
    return 0;
  }
  if (WIFSIGNALED(st)) {
    fprintf(stderr, "exited via signal: %d\n", WTERMSIG(st));
  } else if (WIFEXITED(st)) {
    fprintf(stderr, "unexpected exit code: %d\n", WEXITSTATUS(st));
  } else {
    fprintf(stderr, "unexpected exit status: %d\n", st);
  }
  return 1;
}