#include <cassert>
#include <cstdio>
#include <cstdlib>
#include <cstring>

#include <queue>
#include <set>
#include <utility>
#include <vector>

#include <dlfcn.h>
#include <signal.h>
#include <time.h>
#include <unistd.h>
#include <sys/epoll.h>
#include <sys/signalfd.h>
#include <sys/wait.h>

using std::pair;
using std::queue;
using std::set;
using std::vector;

double LIMIT_THREAD_TIME = 1.0; // 1.0 second
long long LIMIT_THREAD_MEMORY = 256 << 20; // 256 mebibytes
double LIMIT_IDLENESS = 10; // 10.0 seconds
long long LIMIT_TRAFFIC = 8 << 20; // 8 mebibytes

int load_and_solve () {
  void *solution = dlopen ("./solution.so", RTLD_NOW);
  if (!solution) {
    fprintf (stderr, "[worker] dlopen ('./solution.so') failed: %s\n", dlerror ());
    exit (1);
  }
  // fprintf (stderr, "[debug] loaded solution: %p\n", solution);
  int (*solve)();
  *(void **) (&solve) = dlsym (solution, "solve");
  if (!solve) {
    *(void **) (&solve) = dlsym (solution, "main");
  }
  if (!solve) {
    fprintf (stderr, "[master] dlsym ('solve') failed: %s\n", dlerror ());
    exit (1);
  }
  return solve ();
}

int main ( int argc, char **argv ) {
  int r;
  assert (argc >= 3 && argc <= 7);
  int threads = atoi (argv[1]);
  char *stdout_name = argv[2];
  if (argc >= 4)
    sscanf (argv[3], "%lf", &LIMIT_THREAD_TIME);
  if (argc >= 5)
    sscanf (argv[4], "%lld", &LIMIT_THREAD_MEMORY);
  if (argc >= 6)
    sscanf (argv[5], "%lf", &LIMIT_IDLENESS);
  if (argc >= 7)
    sscanf (argv[6], "%lld", &LIMIT_TRAFFIC);
  int report_fd = dup (1);
  assert (freopen (stdout_name, "w", stdout));

  int linux_page_size = getpagesize ();
  assert (linux_page_size > 0);

  int linux_clockticks = sysconf (_SC_CLK_TCK);
  assert (linux_clockticks > 0);
  // fprintf (stderr, "[debug] linux_clockticks = %d\n", linux_clockticks);

  int epoll_fd = epoll_create1 (0);
  assert (epoll_fd >= 0);

  sigset_t mask;
  sigemptyset (&mask);
  sigaddset (&mask, SIGCHLD);
  
  r = sigprocmask (SIG_BLOCK, &mask, NULL);
  assert (r == 0);

  int sfd = signalfd (-1, &mask, SFD_NONBLOCK);
  assert (sfd >= 0);

  struct epoll_event event = {
    .events = EPOLLIN,
    .data = { .u32 = (unsigned int)-2 }
  };
  r = epoll_ctl (epoll_fd, EPOLL_CTL_ADD, sfd, &event);
  assert (r == 0);

  int child_pipe_r[threads];
  int child_pipe_w[threads];
  int child_pid[threads];
  for (int i = 0; i < threads; i++) {
    int fds[4];
    r = pipe (fds);
    assert (r == 0);
    r = pipe (fds + 2);
    assert (r == 0);
    int fds2[2];
    r = pipe (fds2);
    assert (r == 0);

    r = fork ();
    assert (r >= 0);
    if (r == 0) {
      r = dup2 (fds2[0], 0);
      close (fds2[0]);
      close (fds2[1]);
      close (fds[0]);
      close (fds[3]);
      int value = load_and_solve ();
      close (fds[1]);
      close (fds[2]);
      exit (value);
    }

    close (fds2[0]);
    dprintf (fds2[1], "%d %d %d %d\n", threads, i, fds[1], fds[2]);
    close (fds[1]);
    close (fds[2]);

    child_pid[i] = r;
    child_pipe_r[i] = fds[0];
    child_pipe_w[i] = fds[3];
  }

  struct timespec start;
  r = clock_gettime (CLOCK_MONOTONIC_RAW, &start);

  for (int i = 0; i < threads; i++) {
    struct epoll_event event = {
      .events = EPOLLIN | EPOLLHUP | EPOLLRDHUP,
      .data = { .u32 = (unsigned)i }
    };
    r = epoll_ctl (epoll_fd, EPOLL_CTL_ADD, child_pipe_r[i], &event);
    assert (r == 0);
  }

  char buffer[1024];

  int child_buffer_index[threads];
  int child_want[threads];
  int child_clock[threads];
  // long long child_time[threads];
  double child_time[threads];
  vector <char> child_buffer_from[threads];
  queue <vector <char> > child_queue[threads][threads];
  set <int> child_queue_hint[threads];
  memset (child_buffer_index, 0, sizeof (child_buffer_index[0]) * threads);
  memset (child_want, -1, sizeof (child_want[0]) * threads);

  for (int i = 0; i < threads; i++) {
    r = clock_getcpuclockid (child_pid[i], &child_clock[i]);
    assert (r == 0);
    child_time[i] = 0.0;
  }

  struct timespec last_full_check = {
    .tv_sec = 0,
    .tv_nsec = 0
  };

  bool kill_all_childs = false;
  long long rss_peak_one = 0;
  int rss_peak_id = -1;
  long long rss_peak_total = 0;
  long long thread_total_traffic = 0;
  while (1) {
    struct epoll_event events[10];
    int cnt = epoll_wait (epoll_fd, events, 10, 5/*ms*/); // need to check ML and TL many times per second
    bool check_childs = false;
    set <int> workers_aware;
    for (int i = 0; i < cnt; i++) {
      if (events[i].data.u32 == (unsigned)-2) {
        check_childs = true;
        while (read (sfd, buffer, 1024) > 0)
          ;
        continue;
      }
      
      // fprintf (stderr, "event! %d, %d\n", events[i].events, events[i].data.u32);
      int source = (int)events[i].data.u32;
      assert (0 <= source && source < threads);
      r = read (child_pipe_r[source], buffer, 1024);
      assert (r >= 0);
      for (int i = 0; i < r; i++)
        child_buffer_from[source].push_back (buffer[i]);

      while (true) {
        struct {
          unsigned target;
          unsigned length;
        } header;
        if (child_buffer_index[source] + sizeof (header) > child_buffer_from[source].size ())
          break;
        memcpy (&header, child_buffer_from[source].data () + child_buffer_index[source], sizeof (header));
        if (header.length == (unsigned)-1) {
          // fprintf (stderr, "[debug master] child #%d wants data from %d\n", source, (int)header.target);
          if (header.target == (unsigned)-1)
            child_want[source] = -2;
          else
            child_want[source] = header.target;
          workers_aware.insert (source);
          child_buffer_index[source] += sizeof (header);
          continue;
        }
        if (child_buffer_index[source] + sizeof (header) + header.length > child_buffer_from[source].size ())
          break;
        thread_total_traffic += header.length;
        if (thread_total_traffic > LIMIT_TRAFFIC) {
          fprintf (stderr, "traffic limit exceeded: [worker %d] tried to send %d bytes leads to total of %lld bytes\n", source, (int)header.length, thread_total_traffic);
          dprintf (report_fd, "NL send %lld bytes is more than %lld\n", thread_total_traffic, LIMIT_TRAFFIC);
          kill_all_childs = true;
          break;
        }
        // fprintf (stderr, "packet from %d: target=%d length=%d\n", source, header.target, header.length);
        vector <char> packet (header.length);
        for (int i = 0; i < (int)header.length; i++)
          packet[i] = child_buffer_from[source][child_buffer_index[source] + sizeof (header) + i];
        child_buffer_index[source] += sizeof (header) + header.length;
        assert (0 <= (int)header.target && (int)header.target < threads);
        child_queue[header.target][source].push (packet);
        child_queue_hint[header.target].insert (source);
        workers_aware.insert (header.target);
      }
      if (kill_all_childs)
        break;
    }

    if (kill_all_childs)
      break;

    // for (int i = 0; i < threads; i++) {
    for (auto i: workers_aware) {
      if (child_want[i] == -1)
        continue;
      if (child_want[i] == -2) {
        while (child_queue_hint[i].size ()) {
          auto it = child_queue_hint[i].begin ();
          if (child_queue[i][*it].size ()) {
            child_want[i] = *it;
            break;
          } else {
            child_queue_hint[i].erase (it);
          }
        }
      }
      if (child_want[i] == -2)
        continue;
      assert (0 <= child_want[i] && child_want[i] < threads);
      if (child_queue[i][child_want[i]].size () == 0)
        continue;
      auto packet = child_queue[i][child_want[i]].front ();
      child_queue[i][child_want[i]].pop ();
      struct {
        unsigned target;
        unsigned length;
      } header = {
        .target = (unsigned)child_want[i],
        .length = (unsigned)packet.size ()
      };
      // fprintf (stderr, "[debug] send packet to %d from %d\n", i, child_want[i]);
      r = write (child_pipe_w[i], &header, sizeof (header));
      assert (r == (int)sizeof (header));
      r = write (child_pipe_w[i], packet.data (), packet.size ());
      assert (r == (int)packet.size ());
      child_want[i] = -1;
    }

    struct timespec now;
    clock_gettime (CLOCK_MONOTONIC_RAW, &now);
    if ((now.tv_sec - start.tv_sec) * 1000000000ll + (now.tv_nsec - start.tv_nsec) > LIMIT_IDLENESS * 1000000000ll) {
      fprintf (stderr, "idleness limit exceeded (%lld of %lf)\n", (now.tv_sec - start.tv_sec) * 1000000000ll + (now.tv_nsec - start.tv_nsec),  LIMIT_IDLENESS * 1000000000ll);
      dprintf (report_fd, "IL %.2lfs of %.2lfs\n", now.tv_sec - start.tv_sec + 0.000000001 * (now.tv_nsec - start.tv_nsec), 1.0 * LIMIT_IDLENESS);
      kill_all_childs = true;
      break;
    }

    if (cnt == 0 || check_childs || (now.tv_sec - last_full_check.tv_sec) * 1000000000ll + (now.tv_nsec - last_full_check.tv_nsec) > 4000000) { // full check every 4 ms
      last_full_check = now;

      long long rss_total = 0;
      for (int i = 0; i < threads; i++) {
        if (child_pid[i] == -1)
          continue;

        static char stat_name[1024];

        snprintf (stat_name, 1000, "/proc/%d/statm", child_pid[i]);
        FILE *stat_file = fopen (stat_name, "r");
        if (!stat_file) {
          fprintf (stderr, "[master] failed to open: %s: %m\n", stat_name);
          continue;
        }

        long long nonce0, rss;
        r = fscanf (stat_file, "%lld%lld", &nonce0, &rss);
        if (r != 2) {
          fprintf (stderr, "[master] failed to read statm: r=%d %m\n", r);
          continue;
        }
        fclose (stat_file);
        rss *= linux_page_size;

        if (rss > LIMIT_THREAD_MEMORY) {
          fprintf (stderr, "memory limit exceeded: child #%d: RSS: %lld bytes\n", i, rss);
          dprintf (report_fd, "ML [worker #%d] %.2lf MiB of %.2lf MiB\n", i, rss / 1024.0 / 1024.0, LIMIT_THREAD_MEMORY / 1024.0 / 1024.0);
          kill_all_childs = true;
          break;
        }
        if (rss > rss_peak_one) {
          rss_peak_one = rss;
          rss_peak_id = i;
        }
        rss_total += rss;

        /*
        snprintf (stat_name, 1000, "/proc/%d/stat", child_pid[i]);
        stat_file = fopen (stat_name, "r");
        if (!stat_file) {
          fprintf (stderr, "[master] failed to open: %s: %m\n", stat_name);
          continue;
        }

        int nonce3, nonce4, nonce5, nonce6, nonce7;
        static char nonce1[1024];
        char nonce2;
        unsigned nonce8;
        unsigned long nonce9, nonce10, nonce11, nonce12;
        unsigned long time_user, time_system;

        r = fscanf (stat_file,
          "%lld%s %c%d%d%d%d%d%u%lu%lu%lu%lu"
          "%lu%lu",
          &nonce0, nonce1, &nonce2, &nonce3, &nonce4, &nonce5, &nonce6, &nonce7, &nonce8, &nonce9, &nonce10, &nonce11, &nonce12,
          &time_user, &time_system
        );
        if (r != 15) {
          fprintf (stderr, "[master] failed to read stat (r=%d)\n", r);
          continue;
        }
        // fprintf (stderr, "[worker #%d] time = %lu + %lu\n", i, time_user, time_system);
        if ((long long)(time_user + time_system) < child_time[i]) {
          fprintf (stderr, "[master] child #%d cpu time goes down %lld -> %lu\n", i, child_time[i], time_user + time_system);
        } else {
          child_time[i] = time_user + time_system;
        }
        */
        struct timespec process_time;
        r = clock_gettime (child_clock[i], &process_time);
        assert (r == 0);
        // fprintf (stderr, "[debug] worker #%d clock_id: %d, time: (%ld, %ld)\n", i, child_clock[i], process_time.tv_sec, process_time.tv_nsec);

        if ((process_time.tv_sec + 0.000000001 * process_time.tv_nsec) < child_time[i]) {
          fprintf (stderr, "[master] child #%d cpu time goes down %.4lfs -> %.4lfs\n", i, child_time[i], process_time.tv_sec + 0.000000001 * process_time.tv_nsec);
        } else {
          child_time[i] = process_time.tv_sec + 0.000000001 * process_time.tv_nsec;
        }

        if (child_time[i] > LIMIT_THREAD_TIME) {
          fprintf (stderr, "cpu time limit exceeded: child #%d: %.4lfs of %.4lfs\n", i, child_time[i], LIMIT_THREAD_TIME);
          dprintf (report_fd, "TL [worker #%d]: %.4lfs of %.4lfs\n", i, child_time[i], LIMIT_THREAD_TIME);
          kill_all_childs = true;
          break;
        }
      }
      if (rss_total > rss_peak_total) {
        rss_peak_total = rss_total;
      }

      if (kill_all_childs)
        break;
    }

    if (!check_childs)
      continue;

    int threads_alive = 0;
    for (int i = 0; i < threads; i++) {
      if (child_pid[i] == -1)
        continue;
      int wstatus;
      r = waitpid (child_pid[i], &wstatus, WNOHANG);
      if (r == 0) {
        threads_alive++;
        continue;
      }
      int pid = child_pid[i];
      child_pid[i] = -1;
      if (WIFEXITED (wstatus) && WEXITSTATUS (wstatus) == 0)
        continue;
      if (WIFEXITED (wstatus)) {
        fprintf (stderr, "runtime error: child #%d [pid %d]: exit code %d\n", i, pid, WEXITSTATUS (wstatus));
        dprintf (report_fd, "RE [worker #%d] exit code %d\n", i, WEXITSTATUS (wstatus));
      } else if (WIFSIGNALED (wstatus)) {
        fprintf (stderr, "runtime error: child #%d [pid %d]: exit by signal %d\n", i, pid, WTERMSIG (wstatus));
        dprintf (report_fd, "RE [worker #%d] signal %d\n", i, WTERMSIG (wstatus));
      } else {
        fprintf (stderr, "check failed: child #%d [pid %d]: unknown wstatus=%d\n", i, pid, wstatus);
        dprintf (report_fd, "CF [worker #%d] wstatus %d\n", i, wstatus);
      }
      kill_all_childs = true;
      break;
      // fprintf (stderr, "child #%d exited: %d\n", child_pid[i], WEXITSTATUS (wstatus));
    }

    if (threads_alive == 0 || kill_all_childs)
      break;
  }
  if (kill_all_childs) {
    for (int i = 0; i < threads; i++) {
      if (child_pid[i] == -1)
        continue;
      r = kill (child_pid[i], SIGKILL);
    }
  } else {
    if (thread_total_traffic >= 1000 * 1024)
      fprintf (stderr, "total traffic: %.2lf MiB\n", thread_total_traffic / 1024.0 / 1024.0);
    else if (thread_total_traffic >= 1000)
      fprintf (stderr, "total traffic: %.2lf KiB\n", thread_total_traffic / 1024.0);
    else
      fprintf (stderr, "total traffic: %lld bytes\n", thread_total_traffic);
    
    struct timespec now;
    clock_gettime (CLOCK_MONOTONIC_RAW, &now);
    fprintf (stderr, "real time: %.4lfs\n", now.tv_sec - start.tv_sec + 0.000000001 * (now.tv_nsec - start.tv_nsec));

    fprintf (stderr, "memory peak (one): %.2lf MiB (worker #%d)\n", rss_peak_one / 1024.0 / 1024.0, rss_peak_id);
    fprintf (stderr, "memory peak (total): %.2lf MiB\n", rss_peak_total / 1024.0 / 1024.0);

    double cpu_max = 0.0, cpu_sum = 0.0;
    int cpu_max_id = 0;
    for (int i = 0; i < threads; i++) {
      cpu_sum += child_time[i];
      if (child_time[i] > cpu_max) {
        cpu_max = child_time[i];
        cpu_max_id = i;
      }
    }
    fprintf (stderr, "cpu time (max): %.4lfs (worker #%d)\n", cpu_max, cpu_max_id);
    fprintf (stderr, "cpu time (total): %.4lfs\n", cpu_sum);

    dprintf (report_fd, "OK\n");
  }
  return 0;
}

