#define _POSIX_C_SOURCE 1
#include <unistd.h>
#include <stdlib.h>
#include <stdio.h>
#include <sys/wait.h>
#include <sys/types.h>
#include <sys/select.h>
#include <errno.h>
#include <signal.h>
#include <string.h>
#include <fcntl.h>

#define BUFSIZE 1024

typedef enum{
 kSignal, ///< Killed by signal
 kExit,   ///< Called exit
 kAlive   ///< Still alive
} Status;

typedef struct{
  int in;        ///< process stdin(w)
  int out;       ///< process stdout(r)
  int err;       ///< process stderr(r)
  pid_t pid;     ///< process id
  Status status; ///< process status
  int code;      ///< exit/signal code
} proc;

/**
 * Exec with fd dup for stdin, stdout and stderr.
 * argv should usually contain cmd.
 */
void execdup2(int in, int out, int err, const char* cmd, char* const argv[]){
  dup2(in,  STDIN_FILENO);
  dup2(out,STDOUT_FILENO);
  dup2(err,STDERR_FILENO);
  execv(cmd,argv);
  exit(1);
}

/**
 * Exec cmd with args and get pid and pipes to stdin, stdout and stderr.
 */
void exec_piped(proc* s, const char* cmd, char* const argv[]){
  int in[2];
  int out[2];
  int err[2];
  pipe(in);
  pipe(out);
  pipe(err);
  pid_t pid = fork();
  if(pid==0){
    close(in[1]);
    close(out[0]);
    close(err[0]);
    execdup2(in[0], out[1], err[1], cmd, argv);
  }
  close(in[0]);
  close(out[1]);
  close(err[1]);
  s->in = in[1];
  s->out = out[0];
  s->err = err[0];
  s->pid = pid;
  fcntl(s->err,F_SETFL,O_NONBLOCK);
  fcntl(s->out,F_SETFL,O_NONBLOCK);
}

int mx(int a, int b){return a<b?b:a;}

/**
 * Wait for process exit and update the structure.
 */
int wpd(proc* v){
  int status;
  if(waitpid(v->pid, &status, 0)==-1)return -1;
  if(WIFEXITED(status)) v->status = kExit, v->code = WEXITSTATUS(status);
  else if(WIFSIGNALED(status)) v->status = kSignal, v->code = WTERMSIG(status);
  return 0;
}

char errbuf[BUFSIZE];

int main(int argc, char* const argv[]){
  signal(SIGPIPE, SIG_IGN);

  // Make logging to stderr buffered
  setvbuf(stderr, errbuf, _IOFBF, BUFSIZE);

  //Processes and buffer.
  proc v[2];
  char buff[BUFSIZE], feedback[BUFSIZE], scores[BUFSIZE];

  feedback[0] = '\0';
  scores[0] = '\0';

  char** argv0 = malloc((argc-1)*sizeof(char*));
  for(int i=2;i<argc;++i)argv0[i-2]=argv[i];
  argv0[argc-2]=NULL;

  char* argv1[10];
  char* pp = strtok(argv[1], " ");
  int pc = 0;
  while (pp != NULL) {
    argv1[pc++] = pp;
    pp = strtok(NULL, " ");
  }
  argv1[pc] = NULL;

  v[0].status = kAlive;
  v[1].status = kAlive;

  exec_piped(v+0,argv[2],argv0);
  exec_piped(v+1,argv1[0],argv1);

  free(argv0);

  int m = mx(mx(v[0].out,v[0].err),mx(v[1].out,v[1].err));
  ++m;

  int rd_fds[] = {v[0].out,v[0].err,v[1].out,v[1].err};

  int rd_start=0;
  int rd_end=4;

  int ex_start=0;
  int ex_end=6;
  int should_restart = 0;

  {
    ssize_t s;
    do{
      s = read(STDIN_FILENO, buff, BUFSIZE);
      if(s>0)write(v[0].in, buff, s);
    }while(s==BUFSIZE);
  }

  while(v[0].status == kAlive && (v[1].status == kAlive || (should_restart && v[1].status == kExit && v[1].code == 0))) {
    fd_set rd;
    FD_ZERO(&rd);
    for(int i=rd_start;i<rd_end;++i)FD_SET(rd_fds[i],&rd);

    // We wait for any input.
    int k = select(m, &rd, NULL, NULL, NULL);

    // Dump grader stderr to avoid hangups
    if(FD_ISSET(v[0].err, &rd)){
      ssize_t s;
      int killed=0;
      do{
        s = read(v[0].err, buff, BUFSIZE);
        if(s==-1 && (errno == EAGAIN || errno == EWOULDBLOCK)) break;
        // Control interface here
        ssize_t ii=0;
        while(ii<s){
          if(s-ii>=7 && memcmp(buff+ii,"encoder",7) == 0) {
            fprintf(stderr, "Activating encoder mode\n");
            fflush(stderr);
            should_restart = 1;
            ii+=6;
          }
          else if(s-ii>=7 && memcmp(buff+ii,"decoder",7) == 0) {
            fprintf(stderr, "Deactivating encoder mode\n");
            fflush(stderr);
            should_restart = 0;
            ii+=6;
          }
          else if(s-ii>=5 && memcmp(buff+ii,"reset",5) == 0) {
            fprintf(stderr,"Grader restarted graded software\n");
            fflush(stderr);
            /*managed term sig here*/
            if(v[1].status == kAlive){
              kill(v[1].pid,SIGTERM);
              if(wpd(v+1)==-1)exit(1);
              close(v[1].in);
            }
            v[1].status=kAlive;
            v[1].code=0;
            exec_piped(v+1,argv1[0],argv1);
            rd_fds[2]=v[1].out;
            rd_fds[3]=v[1].err;
            m = mx(mx(v[0].out,v[0].err),mx(v[1].out,v[1].err));
            ++m;
            ii+=4;
            killed=1;
          }
          else if(s-ii>=4 && memcmp(buff+ii,"kill",4) == 0) {
            fprintf(stderr,"Grader killed graded software\n");
            fflush(stderr);
            /*kill sig here, this is basically early exit by grader*/
            // if our graded software has already exited, we won't need this
            if(v[1].status == kAlive){
              kill(v[1].pid,SIGKILL);
              if(wpd(v+1)==-1)exit(1);
              close(v[1].in);
              // we are going to exit grading so we don't want RTE in this case
              v[1].status=kExit;
              v[1].code=0;
            }
            ii+=3;
            killed=1;
          }
          else if(s-ii > 9 && memcmp(buff+ii,"feedback ",9) == 0) {
            ii+=9;
            int fi = ii;
            for(; fi+1 < s; ++fi){
              if(buff[fi] == '\n') break;
            }
            buff[fi] = '\0';
            strncpy(feedback, buff + ii, fi-ii+1);
            fprintf(stderr, "Grader gave feedback: %s\n", feedback);
            fflush(stderr);
            ii = fi+1;
            continue;
          }
          else if(s-ii > 7 && memcmp(buff+ii,"scores ",7) == 0) {
            ii+=7;
            int fi = ii;
            for(; fi+1 < s; ++fi){
              if(buff[fi] == '\n') break;
            }
            buff[fi] = '\0';
            strncpy(scores, buff + ii, fi-ii+1);
            fprintf(stderr, "Grader gave subtask scores: %s\n", scores);
            fflush(stderr);
            ii = fi+1;
            continue;
          }
          ++ii;
        }
      }while(s==BUFSIZE);
      if(killed)goto end;
    }

    // Grader program wants to output to graded program.
    if(FD_ISSET(v[0].out, &rd)){
      ssize_t s;
      do{
        s = read(v[0].out, buff, BUFSIZE);
        if(s==-1 && (errno == EAGAIN || errno == EWOULDBLOCK)) break;
        if(s==-1||s==0)wpd(v+0);
        else{
          write(v[1].in, buff, s);
          fprintf(stderr, "%.*s", (int)s, buff);
        }
      }while(s==BUFSIZE);
    }

    // Graded program wants to output to grader program
    if(FD_ISSET(v[1].out, &rd)){
      ssize_t s;
      do{
        s = read(v[1].out, buff, BUFSIZE);
        if(s==-1 && (errno == EAGAIN || errno == EWOULDBLOCK)) break;
        if (s==-1) wpd(v+1);
        else if (s==0) {
            // Propagate EOF
            close(v[0].in);
            wpd(v+1);
        }
        else{
          write(v[0].in, buff, s);
          fprintf(stderr, "%.*s", (int)s, buff);
        }
      }while(s==BUFSIZE);
    }

    // Dump graded program stderr to avoid hangups
    if(FD_ISSET(v[1].err, &rd)){
      ssize_t s;
      do{
        s = read(v[1].err, buff, BUFSIZE);
        if(s==-1 && (errno == EAGAIN || errno == EWOULDBLOCK)) break;
      }while(s==BUFSIZE);
    }
end:;
  }

  if(v[1].status==kSignal){
    kill(v[0].pid,SIGKILL);
    close(v[0].in);
    if(wpd(v+0) == -1)exit(1);
    v[0].status=kExit;
    v[0].code=0;
  }

  // If grader is still alive, close it.
  if(v[0].status==kAlive){
    close(v[0].in);
    if(wpd(v+0)==-1)exit(1);
  }

  // If graded program is still alive, close it.
  // NOTE: in grader early exit, kill graded program before this through control channel
  //v[1].status = kExit;
  if(v[1].status==kAlive){
    if((v[0].status == kExit && v[0].code == 0) || v[0].status == kSignal) kill(v[1].pid,SIGKILL);
    close(v[1].in);
    if(wpd(v+1)==-1)exit(1);
    if(v[0].status == kExit && v[0].code == 0) v[1].status=kExit, v[1].code=0;
  }

  { // Read grader stderr in case we missed feedback
    ssize_t s;
    s = read(v[0].err, buff, BUFSIZE);
    if (s != -1) {
      ssize_t ii=0;
      while (ii < s) {
        if(s-ii > 9 && memcmp(buff+ii,"feedback ",9) == 0) {
          ii+=9;
          int fi = ii;
          for(; fi+1 < s; ++fi){
            if(buff[fi] == '\n') break;
          }
          buff[fi] = '\0';
          strncpy(feedback, buff + ii, fi-ii+1);
          fprintf(stderr, "Grader gave feedback: %s\n", feedback);
          fflush(stderr);
          ii = fi+1;
          continue;
        }
        else if(s-ii > 7 && memcmp(buff+ii,"scores ",7) == 0) {
          ii+=7;
          int fi = ii;
          for(; fi+1 < s; ++fi){
            if(buff[fi] == '\n') break;
          }
          buff[fi] = '\0';
          strncpy(scores, buff + ii, fi-ii+1);
          fprintf(stderr, "Grader gave subtask scores: %s\n", scores);
          fflush(stderr);
          ii = fi+1;
          continue;
        }
        ii++;
      }
    }
  }

  // Print grader status, graded status, grader code, graded code
  // Feedback on 2nd line (or empty)
  // Scores on 3rd line (or empty, only used for multipurpose scoring)
  printf("%d %d %d %d\n%s\n%s\n", v[0].status, v[1].status, v[0].code, v[1].code, feedback, scores);
  return 0;
}
