#include <stdlib.h>
#include <string.h> 
#include <unistd.h> 
#include "UCTPacketMgmt.h"
#include "UCTGlobal.h"

/* The following variables are used throughout the UCTuple environment
   for fast access to datatype sizes. */
int intsize = sizeof(int);  
int floatsize = sizeof(float);
int addrsize = sizeof(void *);
int headersize = 2*intsize;

char * HeaderBuffer = new char[headersize];	//global header buffer

/* WriteData() and ReadData() are wrapper functions for read() and write()
   calls.  Since these calls are not guaranteed to transfer all requested
   data in one transaction, these functions are used to ensure that all
   data requested by the caller is delivered to them. */
int WriteData (int SocketDesc, char* buffer, int size) {
  int totcount = 0,
      count = 0;

  while (totcount < size){
    count = write(SocketDesc, (buffer + totcount), (size - totcount));
    if (count > 0) totcount += count;
  }
  return totcount;
}

int ReadData (int SocketDesc, char* buffer, int size) {
  int totcount = 0,
      count = 0;

  while (totcount < size){
    count = read(SocketDesc, (buffer + totcount), (size - totcount));
    if (count > 0) totcount += count;
  }
  return totcount;
}

/* EncodePacketClnt() is used by the workers for in(), out(), and rd() 
   calls to build a UCTuplet packet for transmission to the server.  This
   packet consists of header and data segments.  The header segment
   consists of two ints, which are concatenated together and placed at the
   beginning of the buffer that is sent to the manager.  This header is as
   follows:

   <size of data segment><operation type>

   The data segment is constructed as follows:

   <size of 1st data element><1st data element>
   <size of 2nd data element><2nd data element>
   .
   .
   .

   The header and data segments are concatenated together and delivered to
   the manager as a single packet. */

Tuple* EncodePacketClnt(int op, char* key, int keysize,
                        va_list ap1, va_list ap2) {
  int lastint,	//stores the last int processed, for array mgmt
      datasize;	//stores a running sum of the size of the data segment
  float f;

  char* index,  //used to index into the send buffer 
      * buffer, //send buffer
      * tmpstr;
  void* tempptr;
  char type;

  /* Initialize the datasize to the size of the key + NULL char */
  datasize = keysize + 1;

  /* We first need to determine how large the data segment will be
     so that we can allocate the send buffer.  Loop through each key
     and keep a running count of the size so far. */
  for (int i=1; i < keysize; i++) {
    switch(key[i]) {
      case 'i':
        lastint = va_arg(ap1, int);
        datasize += intsize;
        break;
      case 'f':
        va_arg (ap1, double);
        datasize += floatsize;
        break;
      case 's':
        datasize += (strlen(va_arg(ap1, char*))+1);
        break;
      case 'I':
        va_arg(ap1, int*);
        datasize += (lastint * intsize);
        break;
      case 'F':
        va_arg(ap1, float*);
        datasize += (lastint * floatsize);
        break;
      case 'p':
        va_arg(ap1, void*);
        datasize += addrsize;
        break;
    }
  }

  /* We have the size of the data segment, so allocate the send buffer.
     Keep in mind that we need additional space for the header. */
  buffer = new char[headersize + datasize];
  index = buffer;

  /* Copy the size of the data segment to the first header int. */
  memcpy(index, &datasize, intsize);
  index += intsize;

  /* Copy the operation type to the second header int. */
  memcpy(index, &op, intsize);
  index += intsize;

  /* Copy the key size to the beginning of the data segment. */
  memcpy(index, key, keysize+1);
  index += (keysize+1);

  /* Now loop through the elements to store the data into the data segment. */
  for (int j=1; j < keysize; j++) {
    type = key[j];
    switch(type) {
      case 'i':
        lastint = va_arg(ap2, int);
        memcpy(index, &lastint, intsize);
        index += intsize;
        break;
      case 'f':
        f = (float) va_arg(ap2, double);
        memcpy(index, &f, floatsize);
        index += floatsize;
        break;
      case 's':
        tmpstr = va_arg(ap2, char*);
        memcpy(index, tmpstr, strlen(tmpstr)+1);
        index += (strlen(tmpstr)+1);
        break;
      case 'I':
        memcpy(index, va_arg(ap2, int*), intsize*lastint);
        index += (intsize*lastint);
        break;
      case 'F':
        memcpy(index, va_arg(ap2, float*), floatsize*lastint);
        index += (floatsize*lastint);
        break;
      case 'p':
        tempptr = va_arg(ap2, void*);
        memcpy(index, &tempptr, addrsize);
        index += addrsize;
        break;
    }
  }

  if (key[0] == 'r') {		//remote call
    /* The send buffer has been filled, deliver it to the manager. */
    WriteData(SD, buffer, (datasize + headersize));

    /* free up the send buffer */
    delete [] buffer;
    return NULL;
  }
  if (key[0] == 'l') {		//local call (handled directly by mgr)
    Tuple* t;
    char* DataBuffer = new char[datasize];
    memcpy(DataBuffer, buffer+headersize, datasize);
    t = BuildTuple(DataBuffer, op, datasize);
    delete [] buffer;
    return t;
  }

  Error("Invalid first key (must be l or r).\n");
  return NULL;
}

/* DecodePacketSrvr() is used by the manager to decode an incoming
   packet sent by EncodePacketClnt() from out(), in(), rd(), or tmexec().
   It analyzes the header buffer to establish the size of the
   pending data segment and operation type. It then allocates the data
   buffer based on this size and retrieves this data buffer from the
   socket.  Finally, it calls BuildTuple with this data buffer to convert
   it into a UCTuple, and returns this UCTuple. */

Tuple* DecodePacketSrvr(int client) {
  int size,	//stores the size of the data segment
      NBytes,
      op;	//stores the operation type
  char* DataBuffer;

  /* Establish the size of the data segment from the header buffer. */
  memcpy(&size, HeaderBuffer, intsize);

  /* Establish the operation type from the header buffer. */
  memcpy(&op, HeaderBuffer + intsize, intsize);
 
  /* We now know the size of the data buffer, so allocate it and
     retrieve it. */
  DataBuffer = new char [size];
  
  NBytes = ReadData(client, DataBuffer, size);
  if (NBytes != size) Error("DecodePacketSrvr failed: sizes do not match\n");

  /* Build a new UCTuple from this data buffer and return a pointer to it. */
  return BuildTuple(DataBuffer, op, size);
}

/* EncodePacketSrvr() is used by the manager to package return data for
   delivery to a worker, for in() and rd() calls.  It takes as input the
   original UCTuple delivered to the manager by the worker (possibly
   containing "p" keys) and the matched UCTuple retrieved from the UCTuple
   space which contains all the requested information that will be returned
   to the worker.  It then builds a packet containing only the data
   requested by the worker in "p" keys (the data to write to the worker).
   This packet consists of header and data segments.  The header segment
   consists of two ints concatenated together and is placed at the head
   of the packet to be sent.  The header is organized as follows:

   <size of data segment><number of elements in data segment>

   The data segment is itself segmented into a preamble describing 
   the organization of the data elements, and the data elements
   themselves. This preamble is organized as follows:
   
   <offset to 1st element address><size of 1st element data>
   <offset to 2nd element address><size of 2nd element data>
   .
   .
   .

   The offset specified here is referenced from after this preamble.  The
   offset points to a location containing the address to write the data to.
   This is the address sent to us in "p" fields from the incoming UCTuple.
   Immediately following this address is the data.  The size field above
   represents the size of this data.  So, the data segment following the
   preamble is organized as follows:

   <address of 1st data element><data for 1st data element>
   <address of 2nd data element><data for 2nd data element>
   .
   .
   .
    
   The header and data segments are concatenated together and delivered to
   the manager as a single packet. */
   
void EncodePacketSrvr(int client, Tuple* InTuple, Tuple* OutTuple) {
  int count = 0,		//tracks the number of "p" fields seen thus far
      datasize = 0,	//running sum of the datasize thus far
      offset;		//tracks the offset into data buffer for data

  char* data_index,	//index into data buffer to which data is written
      * info_index;	//offset into preamble to which data info is written
  char* buffer;		//buffer for packet to send
 
  /* First, loop through the fields of the incoming tuple to grab
     a count of the "p" fields (which will correspond to the number of
     elements to send) and the datasize so that we can allocate the
     buffer. */
  for (int i=1; i < InTuple->keysize; i++)
    if (InTuple->key[i] == 'p') {
      count++;
      datasize += OutTuple->info[i][1]; //get size info for data
    }

  /* At this point, datasize only contains the size of the raw data, but
     the data segment also contains addresses for each element and preamble
     info, so add these sizes to the datasize. */
  datasize += ((count * addrsize) + (count * (intsize * 2)));

  /* We have the size of our data buffer, so allocate our send buffer.
     Don't forget about the header. */ 
  buffer = new char[datasize + headersize];
  info_index = buffer;

  /* Copy the size of the data segment to the 1st header int. */
  memcpy(info_index, &datasize, intsize);
  info_index += intsize;

  /* Copy the number of elements to the 2nd header int. */
  memcpy(info_index, &count, intsize);
  info_index += intsize;

  /* Set our data_index to index into the data segment */
  data_index = info_index + (count * (intsize * 2));
  offset = count * (intsize * 2);

  /* Now, loop through the incoming UCTuple key to find "p" fields
     and copy the relevant data for those fields into the data segment
     preamble and body. */
  for (int j=1; j < InTuple->keysize; j++)
    if (InTuple->key[j] == 'p') {
      /* copy offset of this element to preamble */
      memcpy (info_index, &offset, intsize);
      info_index += intsize;

      /* copy size of this element to preamble */
      memcpy (info_index, &OutTuple->info[j][1], intsize);
      info_index += intsize;

      /* copy the address for this element to which to write the data to */
      memcpy (data_index, InTuple->data + InTuple->info[j][2], addrsize);
      data_index += addrsize;

      /* copy the data */
      memcpy (data_index,
              OutTuple->data + OutTuple->info[j][2],
              OutTuple->info[j][1]);
      data_index += OutTuple->info[j][1];

      /* update the offset for the next element */
      offset += (addrsize + OutTuple->info[j][1]);
    }

  if (InTuple->op == IN) delete OutTuple;
  delete InTuple;

  if (InTuple->key[0] == 'r')
    /* Buffer built, send it to the worker. */
    WriteData(client, buffer, (datasize + headersize));

  delete [] buffer;
}

/* DecodePacketClnt() is called by the worker to decode the packet
   constructed by the manager using the function EncodePacketSrvr().
   This will break the packet up by analyzing the data preamble to 
   determine the addresses to which to write requested data. It takes
   a header buffer as input, which contains the size of the pending
   data buffer and the number of elements in the data buffer. */
void DecodePacketClnt() {
  int size,		//stores the size of the data buffer
      NBytes,
      count,	//stores the number of elements in the data buffer
      offset,	//keeps track of the offset for the current element
      datasize;	//keeps track of the data size for the current element
  void* addr;	//stores the address to which to write the current element
  char* DataBuffer;	//data buffer
  char* index;	//index into the data buffer

  /* Retrieve the size of the data segment */
  memcpy(&size, HeaderBuffer, intsize);

  /* Retrieve the number of elements */
  memcpy(&count, HeaderBuffer + intsize, intsize);

  /* We now know the size of the pending data buffer, so allocate it, and
     retrieve it. */
  DataBuffer = new char [size];
  
  NBytes = ReadData(SD, DataBuffer, size);
  if (NBytes != size) Error("DecodePacketClnt failed: sizes do not match\n");

  /* Init the offset into the data buffer. */
  index = DataBuffer;

  /* Loop through data buffer, write element data to the associated address. */
  for (int i=0; i < count; i++) {
    /* Determine offset for this elements address */
    memcpy (&offset, index, intsize);
    index += intsize;

    /* Determine the size of the data block for this element. */
    memcpy (&datasize, index, intsize);
    index += intsize;

    /* Retrieve the address to which to write the data block. */
    memcpy (&addr, DataBuffer + offset, addrsize);

    /* Write the data block to the specified address. */
    memcpy (addr, DataBuffer + offset + addrsize, datasize);
  }
}