#include "clustermanager.h"

void *pthread_listen( void *arg )
{
  clustermanager *a;
  a = (clustermanager*) arg;
  a->workerlisten();
  return NULL;
}

void *pthread_processworker( void *a )
{
  processworker_data *arg;
  arg = (processworker_data*)a;
  arg->arg1->processworker( arg->arg2 );
  delete arg;
  return NULL;
}

clustermanager::clustermanager()
{
  pthread_mutex_init( &workerlist_lock, NULL );
  pthread_mutex_init( &queuelock, NULL );  
  pthread_mutex_init( &finishedlock, NULL );
  pthread_mutex_init( &wcountlock, NULL );
  pthread_mutex_init( &workdatalock, NULL );
  pthread_cond_init( &workerlist_cond, NULL );
}

clustermanager::~clustermanager()
{
  pthread_mutex_destroy( &workerlist_lock );
  pthread_mutex_destroy( &queuelock );
  pthread_mutex_destroy( &finishedlock );
  pthread_mutex_destroy( &wcountlock );  
  pthread_mutex_destroy( &workdatalock );
  pthread_cond_destroy( &workerlist_cond );
}

int clustermanager::init( std::string add, int p )
{
  serv.init( add, p );
  return 0;
}

int clustermanager::getfinished( std::vector< workunit* > &lst )
{
  lst = finished;
  finished.clear();
  return 0;
}

int clustermanager::start()
{
  finish = 0;
  activeworkers = 0;
  workerlist.clear();
  queue.clear();
  finished.clear();

  int ret;
  ret = pthread_create( &tid, 0, &pthread_listen, this );
  std::cout << "waiting for worker connection\n";
  // wait for at least one worker to job
  pthread_mutex_lock( &workerlist_lock );
  while( workerlist.size() == 0 )
    pthread_cond_wait( &workerlist_cond, &workerlist_lock );
  pthread_mutex_unlock( &workerlist_lock );
  std::cout << "got worker\n";
  return ret;
}

int clustermanager::runjoblist()
{

  waitworkers();
  for( std::list<worker*>::iterator it = workerlist.begin();
       it != workerlist.end(); it++ )
    {
      pthread_mutex_lock( &((*it)->joblock) );
      (*it)->curjob = 2;
      pthread_cond_broadcast( &((*it)->jobcond) );
      pthread_mutex_unlock( &((*it)->joblock) );
    }

  waitworkers();
  return 0;
}

int clustermanager::end()
{
  finish = 1;
  void *ret;

  serv.closesocket();

  pthread_join( tid, &ret );
  waitworkers();
  // join all active workers
  for( std::list< worker* >::iterator it = workerlist.begin();
       it != workerlist.end(); it++ )
    {
      pthread_mutex_lock( &((*it)->joblock) );
      (*it)->curjob = 1;
      pthread_cond_broadcast( &((*it)->jobcond) );
      pthread_mutex_unlock( &((*it)->joblock) );
    }
  waitworkers();
  std::list<worker*>::iterator w;
  w = workerlist.begin();
  while( w != workerlist.end()) 
    {
      if( (*w)->status == 0 )
        {
          std::cout << "worker finished joining\n";
          void *myret;
          pthread_join( (*w)->id, &myret );
          pthread_mutex_destroy( &( (*w)->joblock ) );
          pthread_cond_destroy( &( (*w)->jobcond ) );
          std::list<worker*>::iterator p;
          delete *w;
          p = w;
          w++;
          workerlist.erase(p);
        }
      else
        w++;
    }
  std::cout << workerlist.size() << " workers left\n";

  return 0;
}

int clustermanager::addwork( workunit *wu )
{
  pthread_mutex_lock( &queuelock );
  queue.push_back( wu );
  pthread_mutex_unlock( &queuelock );
  return 0;
}

int clustermanager::workerlisten()
{
  tcpconnection *ntc;
  ntc = NULL;
  while( finish == 0 )
    {

#ifdef __DEBUG_CLUSTER_4__
      std::cout << "waiting for worker\n";
#endif

      ntc = new tcpconnection;
      int ret;
      ret = serv.wait( *ntc );
      // new worker join 
      if( ret == 0 )
        {
          std::cout << "new worker\n";
          worker newworker;
          worker *w;
          w = new worker;
          w->status = 1;
          w->curjob = 0;
          w->cn = ntc;
          pthread_mutex_init( &(w->joblock), NULL );
          pthread_cond_init( &(w->jobcond), NULL );

          processworker_data *myargs;
          myargs = new processworker_data;
          myargs->arg1 = this;
          myargs->arg2 = (void*)w;
          pthread_create( &(w->id), 0, &pthread_processworker, myargs);

          pthread_mutex_lock( &workerlist_lock );
          workerlist.push_back( w );
          pthread_cond_broadcast( &workerlist_cond );
          pthread_mutex_unlock( &workerlist_lock );
        }
      else
        delete ntc ;
      /// time out should do some cleanup work
      pthread_mutex_lock( &workerlist_lock );
      std::list<worker*>::iterator w;
      w = workerlist.begin();
      while( w != workerlist.end()) 
        {
          if( (*w)->status == 0 )
            {
              std::cout << "worker finished joining\n";
              void *myret;
              pthread_join( (*w)->id, &myret );
              pthread_mutex_destroy( &( (*w)->joblock ) );
              pthread_cond_destroy( &( (*w)->jobcond ) );
              std::list<worker*>::iterator p;
              delete *w;
              p = w;
              w++;
              workerlist.erase(p);
            }
          else
            w++;
        }
      pthread_mutex_unlock( &workerlist_lock );
    }
  return 0;
}

int clustermanager::senddata( worker *w )
{
  // send message to delete data
  int wsig;
  wsig = CM_DELDATA;
  w->cn->write( (char*)&wsig, sizeof(int) );
  // read finish signal from worker
  w->cn->read( (char*)&wsig, sizeof(int) );
  if( wsig != CM_REQUEST_NEWJOB )
    {
      std::cout << "ERROR worker\n";
      exit(0);
    }
  int msg[5];
  msg[0] = (int) imsa.size();
  msg[1] = imsa[0]->getheight();
  msg[2] = imsa[0]->getwidth();
  msg[3] = imsb[0]->getheight();
  msg[4] = imsb[0]->getwidth();

  w->cn->sendmessage( CM_NEWDATA, sizeof(int)*5, (char*)msg );
  // send maps
  w->cn->write( (char*) mapa->getimageptr(), sizeof(double)*mapa->getwidth()*mapa->getheight() );
  w->cn->write( (char*) mapb->getimageptr(), sizeof(double)*mapb->getwidth()*mapb->getheight() );

  for(unsigned i = 0; i < imsa.size(); i++ )
    w->cn->write( (char*)( imsa[i]->getimageptr() ),
                  sizeof(double)*msg[1]*msg[2] );
  for(unsigned i = 0; i < imsa.size(); i++ )
    w->cn->write( (char*)( imsb[i]->getimageptr() ),
                  sizeof(double)*msg[3]*msg[4] );

  return 0;
}

int clustermanager::sendqueue( worker *w )
{
  int flag;
  flag = 1;
  while( flag == 1 )
    {
      w->wu = NULL;
      pthread_mutex_lock( &queuelock );
      if( queue.size() == 0 )
        flag = 0;
      else
        {
          w->wu = queue.back();
          queue.pop_back();
        }
      pthread_mutex_unlock( &queuelock );
      if( w->wu != NULL )
        {
          char *msg;
          int msz;
          int wst;
          w->wu->pack( &msg, msz );
          w->cn->sendmessage( CM_NEWWORK, msz, msg );
          delete [] msg;
          w->cn->read( (char*)&wst, sizeof(int) );
          if( wst != CM_JOBRESULT )
            {
              std::cout << "worker expect job result error\n";
              exit(0);
            }
          double res;
          w->cn->read( (char*)&wst, sizeof(int) );
          w->cn->read( (char*)&res, sizeof(double) );
          w->wu->setresult( res );
          pthread_mutex_lock( &finishedlock );
          finished.push_back( w->wu );
          pthread_mutex_unlock( &finishedlock );
          w->wu = NULL;
          w->cn->read( (char*)&wst, sizeof(int) );
          if( wst != CM_REQUEST_NEWJOB )
            {
              std::cout << "worker expect job result error\n";
              exit(0);
            }
        }
    }
  flag = CM_QUEUE_FINISHED;
  w->cn->write( (char*)(&flag), sizeof(int) );
  return 0;
}

int clustermanager::waitworkers()
{
  std::vector< worker* > curworkers;
  pthread_mutex_lock( &workerlist_lock );
  for( std::list< worker* >::iterator it = workerlist.begin(); 
       it != workerlist.end(); it++ )
    {
      if( (*it)->status != 0 )
        curworkers.push_back( *it );
    }
  pthread_mutex_unlock( &workerlist_lock );
  for( unsigned i = 0; i < curworkers.size(); i++ )
    {
      std::cout << "waiting for worker: " << i << "\n";
      pthread_mutex_lock( &(curworkers[i]->joblock) );
      while( curworkers[i]->curjob != 0 )
        {
          std::cout << "worker " << i << " not finished\n";
          pthread_cond_wait( &(curworkers[i]->jobcond),
                             &(curworkers[i]->joblock) );
        }
      pthread_mutex_unlock( &(curworkers[i]->joblock) );
    }
  
  return 0;
}

int clustermanager::setdata( std::vector< image2d* > &sa,
                             std::vector< image2d* > &sb,
                             image2d *ma, image2d *mb )
{
  std::cout << "getting workdata lock\n";
  pthread_mutex_lock( &workdatalock );
  waitworkers();
  imsa = sa;
  imsb = sb;
  mapa = ma;
  mapb = mb;
  std::cout << "setting data\n";
  for( std::list<worker*>::iterator it = workerlist.begin();
       it != workerlist.end(); it++ )
    {
      pthread_mutex_lock( &((*it)->joblock) );
      (*it)->curjob = 3;
      pthread_cond_broadcast( &((*it)->jobcond) );
      pthread_mutex_unlock( &((*it)->joblock) );
    }
  pthread_mutex_unlock( &workdatalock );
  std::cout << "waiting for data to be sent\n";
  waitworkers();
  std::cout << "send done\n";
  return 0; 
}

int clustermanager::sendexit( worker *w )
{
  int id;
  id = CM_EXIT;
  w->cn->write( (char*)(&id), sizeof(int) );
  return 0;
}

int clustermanager::workerrequest( worker *w )
{
  int msg;
  w->cn->read((char*)(&msg), sizeof(int) );
  if( msg == CM_EXIT )
    return 0;
  else if( msg == CM_REQUEST_NEWJOB )
    return 1;
  else
    {
      std::cout << "UNKNOWN MESSAGEID workerresquest\n";
      exit(0);
    }
  return -1;
}

int clustermanager::processworker( void  *arg)
{
  worker *w;
  w = (worker*) arg;
  unsigned myid;
  myid = (unsigned int) w->id;
  //int messageid;
  //int messagesize;
  int flag;
  flag = 1;
  //char *buffer;
  int ttt;
  w->cn->read((char*)(&ttt), sizeof(int) );
  std::cout << "worker signaled ready\n";
  pthread_mutex_lock( &wcountlock );
  activeworkers++;
  pthread_mutex_unlock( &wcountlock );
  while( flag == 1 )
    {
      pthread_mutex_lock( &(w->joblock) );
      while( w->curjob == 0 )
        {
          std::cout << "worker waiting for job\n";
          pthread_cond_wait( &(w->jobcond),
                             &(w->joblock) );
        }
      std::cout << "worker got job\n";
      int myjob;
      myjob = w->curjob;
      pthread_mutex_unlock( &(w->joblock) );
      if( w->curjob == 2 )
        {
          // finish queue
          sendqueue(w);
        }
      else if( w->curjob == 3 )
        {
          // send data
          senddata(w);
          std::cout << "finished sending\n";
        }
      else if( w->curjob == 1 )
        {
          // exit
          std::cout << "sending exit signal\n";
          sendexit(w);
          flag = 0;
        }
      if( flag == 1 )
        {
          flag = workerrequest(w);
        }
      std::cout << "waiting for joblock\n";
      pthread_mutex_lock( &(w->joblock) );
      w->curjob = 0;
      std::cout << "ready to signal\n";
      pthread_cond_signal( &(w->jobcond) );
      pthread_mutex_unlock( &(w->joblock) );
      std::cout << "signaled\n";
    }
  w->cn->closesocket();
  pthread_mutex_lock( &wcountlock );
  w->status = 0;
  activeworkers--;
  pthread_mutex_unlock( &wcountlock );
  delete w->cn;
  std::cout << "worker exiting\n";
  return 0;
}
