#include <iostream>
#include <fstream>
#include <string>
#include <stdlib.h>
#include <sstream>
#include <vector>
#include <sys/types.h>
#include <sys/wait.h>
#include <stdio.h>
#include <unistd.h>
#include <assert.h>
#include <sys/select.h>
#include <map>
#include <set>
#include <iterator>
using namespace std;

static map<pid_t, int> pidsToDescs;
static map<pid_t, string> pidsToFiles;
// map pipe to content from it
static map<int, string> childrenOutput;
static string outputExtension;

string cmd = "";
  
inline bool patchMode() {
  return outputExtension.empty();
}

static void readFileList(const string &file, vector<string> &list) {
  ifstream fin(file.c_str());
  if(!fin.is_open()) {
    cerr << "Could not open file list: " << file << endl;
    exit(1);
  }
  string line;
  while(getline(fin, line)) {
    list.push_back(line);
  }
}

static void run_squash(int out_fd, vector<string> &vArgv) {
  // closing STDOUT_FILENO
  // makes the next open descriptor replace it
  // in this case that would be out_fd
  dup2(out_fd, STDOUT_FILENO);
  close(out_fd);
  char *argv[vArgv.size() + 1];
  for(unsigned int i = 0;i < vArgv.size();i++) {
    argv[i] = (char*)vArgv[i].c_str();
  }
  argv[vArgv.size()] = 0;
  nice(19);
  execvp(argv[0], argv);
  perror("execvp");
}

typedef map<int, string> ismap;
typedef map<string, ismap> patch_container;
static void combine_chunk(ismap &existing, string const &hunk, string const &filename, string const &inputFile) {
  int line = atoi(hunk.c_str() + 4);
  ismap::iterator it = existing.find(line);
  if(it == existing.end()) {
    existing[line] = hunk;
    return;
  }
  if(it->second != hunk) {
    cerr << filename << ":" << line 
	 << ": Patches disagree on what the hunk should look like:" << endl;
    cerr << it->second << " vs " << hunk;
    cerr << filename << ":" << line
	 << ": caused by " << inputFile << endl;
    //abort();
  }
}

static void combine_hunk(patch_container &existing, string hunk, string inputFile) {
  string::size_type pos = hunk.find("\n");
  string filename;
  filename.append(hunk, 4, pos - 4);
  ismap &this_file = existing[filename];
  pos = hunk.find("\n@@", pos);
  hunk.erase(0, pos + 1);
  while(pos != string::npos) {
    string::size_type pos = hunk.find("\n@@ -");
    if(pos == string::npos) {
      pos = hunk.size();
    } else {
      pos++;
    }
    // existing.insert(s.substr(0, pos));
    combine_chunk(this_file, hunk.substr(0, pos), filename, inputFile);

    if(pos == hunk.size()) return;

    hunk.erase(0, pos);
  }
 // cerr << '"' << filename << '"' << endl;
}

/* separates out stuff from the new patch and puts it into a set for later
*/
static void combine(patch_container &existing, const string &additional,
		    string const &inputFile) {
  string s = additional;
  while(!s.empty()) {
    string::size_type pos = s.find("\n---");
    if(pos == string::npos) {
      pos = s.size();
    } else {
      pos++;
    }
    // existing.insert(s.substr(0, pos));
    combine_hunk(existing, s.substr(0, pos), inputFile);

    if(pos == s.size()) return;

    s.erase(0, pos);
  }
}

// if a child died point out the file that caused the error
// and terminate once the rest of the children have quit
// unless SQUASH_CRASH env var is defined
// return true if a child died
static bool quitIfChildDied(patch_container &patch) {
  int status(0);
  // this will provide queued up exit info
  pid_t child = wait(&status);
  int exit_code = 0;
  int badexit=false;
  // this our child that died?
  // execing from perl occasionally leaves us with other children
  if(pidsToFiles.find(child) == pidsToFiles.end()) return false;

  map<pid_t, int>::iterator it = pidsToDescs.find(child);
  if (it != pidsToDescs.end())
    close(it->second);

  if(WIFEXITED(status) && WEXITSTATUS(status) != 0) {
    cerr << cmd << " exited with code " << WEXITSTATUS(status) << endl;
    exit_code = WEXITSTATUS(status);
    badexit = true;
  } else if (WIFSIGNALED(status)) {
    cerr << cmd << " killed by signal " << WSTOPSIG(status) << endl;
    badexit=true;
  }
  std::string inputFile = pidsToFiles[child];
  if(!badexit) {
    std::string const &output = childrenOutput[pidsToDescs[child]];
    if (patchMode()) {
      combine(patch, output, inputFile);
    } else {
      inputFile += outputExtension;
      ofstream fos(inputFile.c_str());
      fos.write(output.c_str(), output.size());
      fos.close();
    }
    return false;
  }

  cerr << "File that caused crash: " << inputFile << endl;
  if(true /*getenv("SQUASH_CRASH")*/) {
    cerr << "IMMORTAL MODE: Continuing" << endl;
    return true;
  }
  // there usually wont be any output, but we can try :)
  cout << childrenOutput[pidsToDescs[child]];
  // wait for other processes to finish
  while(wait(&status) > 0);
  exit(exit_code);
  // special treat to keep the compiler happy
  return false;
}

void randomize(vector<string>& v) {
  srand(time(NULL));
  for(unsigned int len = v.size();len;len--) {
    unsigned int const j = (unsigned int) (((float)len) * (rand() / (RAND_MAX + 1.0)));
    swap(v[j], v[len-1]);
  }
}

void error_usage(const char *self) {
  cerr << "Usage: " << self 
       << "[--save-output .file-extension] process-count /path/to/input_filelist.txt /path/to/tool [options to pass to the tool]" << endl;
  exit(1);
}

int main(unsigned int argc, char** argv) {
  unsigned int n = 1;
  int crash_count = 0;
  if (argc > 2) {
    if (!strcmp(argv[1], "--save-output")) {
      outputExtension = argv[2];
      n += 2;
    }
  }
  if(argc - n + 1 < 4) {
    error_usage(argv[0]);
  }
  time_t oldTime;
  time(&oldTime);
  const int maxProcessCount = atoi(argv[n++]);
  int processCount = 0;
  vector<string> fileList;
  readFileList(argv[n++], fileList);
  cmd = argv[n++];
  //if(!true /*getenv("SQUASH_CRASH")*/) {
  // without randomizing, every time I fix a bug and do another run
  // it runs longer until next error, this way errors should be raised sooner on average
  //randomize(fileList);
    //}
  map<int, string>::iterator it;
  vector<string> child_argv;
  child_argv.push_back(cmd);
  while(n < argc) {
    child_argv.push_back(argv[n++]);
  }
  patch_container patch;
  for(n = 0; n < fileList.size(); n++) {
    int pipes[2];
    if(pipe(pipes) == -1) {
      perror("pipe");
      exit(1);
    }

    pid_t child = fork();
    if(child == 0) {
      // close the reading pipe
      close(pipes[0]);
      child_argv.push_back(fileList[n]);
      run_squash(pipes[1], child_argv);
      // shouldn't get here
      abort();
    }
    // close the writing pipe
    close(pipes[1]);
    childrenOutput[pipes[0]] = "";
    pidsToDescs[child] = pipes[0];
    pidsToFiles[child] = fileList[n];
    processCount++;
    while (processCount == maxProcessCount || (n + 1 == fileList.size() && processCount)) {
      fd_set rfds;
      FD_ZERO(&rfds);
      int maxfd = 0;
      for(it = childrenOutput.begin(); it != childrenOutput.end(); it++) {
        maxfd = max(maxfd, it->first + 1);
        FD_SET(it->first, &rfds);
      }
      int rval = select(maxfd, &rfds, NULL, NULL, NULL);
      if(rval == -1) {
        perror("select");
        exit(2);
      }
      char buf[1024];
      for(it = childrenOutput.begin(); it != childrenOutput.end(); it++) {
        if(!FD_ISSET(it->first, &rfds)) continue;
        ssize_t n = read(it->first, buf, sizeof(buf) - 1);
        if(n == 0) {
          close(it->first);
          processCount--;
          crash_count += quitIfChildDied(patch);
          childrenOutput.erase(it);
          continue;
        }
        buf[n] = 0;
        it->second += buf;
      }
    }
  }
  for(patch_container::iterator it = patch.begin(); it != patch.end(); it++) {
    cout << "--- " << it->first << "\n+++ " << it->first << "\n";
    for(ismap::iterator it2 = it->second.begin(); it2 != it->second.end(); it2++) {
      cout << it2->second;
    }
  }
//   copy(patch.begin(), patch.end(), ostream_iterator<string>(cout,""));
  time_t diff = time(NULL) - oldTime;
  if (diff) {
    cerr << argv[0] << " processed " << fileList.size() << " files in "
         << diff << " seconds. " 
         << (((float)fileList.size())/diff) << " files/second ";
    if (crash_count) 
      cerr << "(" << crash_count << " crashes)";
    cerr << endl;
  }
  return 0;
}
