#include "writeout.h"

#include <sys/types.h>
#include <sys/stat.h>
#include <stdio.h>
#include <stdlib.h>
#include <iostream>
#include <time.h>
#include <sstream>
#include <dirent.h>
#include <errno.h>
#include <unistd.h>

#ifndef MPI_DISABLED
#include <mpi.h>
#endif

using namespace std;

string writeout::longToStr (long arg)
{
  stringstream ss;
  ss << arg;
  return ss.str();
}

void writeout::newsub(string subname) {
  of[subname] = new ofstream;

  if ( fulldir != "" ) {
    if(rank>0) of[subname]->open( (fulldir + "/rank" + cRank + "_" + subname + ".tmp").c_str() );
    else of[subname]->open( (fulldir + "/" + signature + "_" + subname + ".dat").c_str() );

    if ( !of[subname]->is_open() ) {
      logf << "WRITEOUT: Could not open output-file!" << endl << flush;
      exit(1);
    }

    buf[subname] = of[subname]->rdbuf();
  }
  else
    buf[subname] = cout.rdbuf();

  out[subname] = new ostream(buf[subname]);
}

writeout::writeout(const string& wdir, const string& _signature, 
		   const int& _rank, const int& procs)
{
  long timestamp;
  int iseq=0;

  fulldir = "";
  signature = _signature;
  rank = _rank;

  if(wdir != ""){
    numprocs = procs;
    sprintf(cRank, "%d", rank);

    if (rank == 0) {
      timestamp = time(NULL);

      while (true) {
	fulldir = wdir + "/" + longToStr(timestamp) + "." + longToStr(iseq) + "_" + signature + ".tmp";
	if ( mkdir(fulldir.c_str(), 0775) == 0 )
	  break;
	else if ( errno != EEXIST ) {
	  cerr << "WRITEOUT: Could not create outdir!" << endl << flush;
	  break;
	}
	iseq++;
      }

#ifndef MPI_DISABLED
      for(int idest=1; idest<numprocs; idest++) {
	MPI_Send(&timestamp, 1, MPI_LONG, idest, 123, MPI_COMM_WORLD);
	MPI_Send(&iseq, 1, MPI_LONG, idest, 124, MPI_COMM_WORLD);
      }
#endif

    }

#ifndef MPI_DISABLED
    else if(rank>0) {
      MPI_Recv(&timestamp, 1, MPI_LONG, 0, 123, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
      MPI_Recv(&iseq, 1, MPI_LONG, 0, 124, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
      fulldir = wdir + "/" + longToStr(timestamp) + "." + longToStr(iseq) + "_" + signature + ".tmp";
    }
#endif  
    
    logf.open( (fulldir + "/rank" + cRank + ".log").c_str() );

    if ( !logf.is_open() ) {
      cerr << "WRITEOUT: Could not open log-file!" << endl << flush;
      exit(1);
    }

    logf << "[ " << timestring() << " ] Log starts here." << endl;
    
    logbuf = logf.rdbuf();
  }
  else{
    logbuf = cerr.rdbuf();
  }
  log = new ostream(logbuf);
}

string writeout::timestring()
{
  time_t rawtime;
  string timestring;
  time( &rawtime );
  timestring = asctime( localtime( &rawtime ) );
  return timestring.substr(0, timestring.size()-1);;
}

writeout::~writeout()
{
  if(fulldir != "") {
    for (map<string,ofstream*>::iterator ofit = of.begin(); ofit != of.end(); ++ofit) {
      if( cRank[0] == '0' ) {
	int jobsdone=0;
	while(jobsdone<numprocs-1) {
	  string nextfile;
	  if( (nextfile = getdatfile(ofit->first)) == "" ) 
	    sleep(1);
	  else {
	    logf << "collecting " << nextfile << endl;
		    
	    ifstream myfile( (fulldir + "/" + nextfile).c_str() );
	    while(true) {
	      string line;
	      getline(myfile, line);
	      if( !myfile.good() ) break;
	      *ofit->second << line << endl << flush;
	    }
	    myfile.close();
	    remove( (fulldir + "/" + nextfile).c_str() );
	    jobsdone++;
	  }
	}
	*ofit->second << "#end" << endl << flush;
	ofit->second->close();
      }
      else {
	ofit->second->close();
	rename((fulldir + "/rank" + cRank + "_" + ofit->first + ".tmp").c_str(),
	       (fulldir + "/rank" + cRank + "_" + ofit->first + ".part").c_str());
      }
    }
    if( cRank[0] == '0' )
      rename( fulldir.c_str(), fulldir.substr(0, fulldir.length()-4).c_str() );
  }
  logf << "[ " << timestring() << " ] Log ends here." << endl;
  logf.close();
}

string writeout::getdatfile(string subname)
{
  string myfile;
  DIR *dp;
  struct dirent *dirp;

  if((dp  = opendir(fulldir.c_str())) == NULL) {
    logf << "Error(" << errno << ") opening " << fulldir << endl;
    closedir(dp);
    return "";
  }
  
  while ((dirp = readdir(dp)) != NULL)
    {
      myfile = string(dirp->d_name);

      if(myfile.length() > 3 && myfile.substr(myfile.length()-4) == "part" &&
	 subname == myfile.substr( myfile.find("_")+1, myfile.rfind(".")-myfile.find("_")-1 ) ) {
	closedir(dp);
	return myfile;
      }
    }

  closedir(dp);
  return "";
}