#include <assert.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netdb.h>
#include <unistd.h>
#include <fcntl.h>
#include <poll.h>
#include <errno.h>

#include "MiscSupport.h"
#include "ZLibMalloc.h"
#include "SimpleLogFile.h"
#include "ThreadMonitor.h"

#include "TalkWithServer.h"


/////////////////////////////////////////////////////////////////////
// TcpIpConnection
/////////////////////////////////////////////////////////////////////

TcpIpConnection::TcpIpConnection(std::string const &name) :
  _name(name),
  _handle(-1),
  _status(NotYetOpen),
  _writeSucceeded(false)
{
}

TcpIpConnection::~TcpIpConnection()
{
  close();
}

void TcpIpConnection::connect(std::string address, std::string port,
			      bool asyncConnect)
{
  assert(_status == NotYetOpen);
  // Protocol #6 is TCP.  This comes from /etc/protocols.
  _handle = socket(AF_INET, SOCK_STREAM, 6);
  if (_handle == -1)
  { // An error occurred.  This should never happen.  All we've done so
    // far is update the address family and protocol.  We haven't tried
    // to connect to or even name the remote server yet.
    std::string errorMessage = errorString();
    TclList msg;
    msg<<FLF<<"socket()"<<_name<<errorMessage;
    sendToLogFile(msg);
    _status = Closed;
    return;
  }
  addrinfo *addresses;
  int aiResult = getaddrinfo(address.c_str(), port.c_str(), NULL, &addresses);
  if (aiResult != 0)
  { // An error occurred.  Unable to parse the address.
    TclList msg;
    msg<<FLF<<_name<<"getaddrinfo()"<<gai_strerror(aiResult)<<address<<port;
    sendToLogFile(msg);
    ::close(_handle);
    _handle = -1;
    _status = Closed;
    return;
  }
  if (asyncConnect)
    // connect() is asynchronous, as is everything else going forward.
    if (!setNonBlocking())
    {
      freeaddrinfo(addresses);
      return;
    }
  {
    ThreadMonitor::SetState tm("::connect");
    tm.increment("::connect");
    int connectResult =
      ::connect(_handle, addresses->ai_addr, addresses->ai_addrlen);
    freeaddrinfo(addresses);
    if ((connectResult != 0)
	&& !(asyncConnect && (errno == EINPROGRESS)))
    { // An error occurred.
      // EINPROGRESS is not an error.  That's what we expect when we try to
      // connect asynchronously.
      TclList msg;
      msg<<FLF<<_name<<"::connect()"<<errorString()<<address<<port;
      sendToLogFile(msg);
      ::close(_handle);
      _handle = -1;
      _status = Closed;
      return;
    }
  }
  if (!asyncConnect)
  { // connect() was blocking but everything going forward is asynchronous.
    if (!setNonBlocking())
      return;
  }
  // We're calling this open.  We're not sure if we really connected or not.
  // If you want to distinguish between connected then disconnected vs never
  // connected check writeSucceeded().
  _status = Open;
}

bool TcpIpConnection::setNonBlocking()
{
  int fctlResult = fcntl(_handle, F_SETFL, O_NONBLOCK);
  if (fctlResult)
  { // An error occurred.
    TclList msg;
    msg<<FLF<<_name<<"fcntl()"<<errorString();
    sendToLogFile(msg);
    ::close(_handle);
    _handle = -1;
    _status = Closed;
    return false;
  }
  return true;
}

bool TcpIpConnection::closed() const
{
  return _status == Closed;
}

void TcpIpConnection::close()
{
  if (_status == Open)
    ::close(_handle);
  _status = Closed;
}

bool TcpIpConnection::readyToSend()
{
  if (_status != Open)
    // Handle is not ready.
    return false;
  if (!toSend.empty())
    // Some data in the buffer already, so presumably write already told us to
    // pause.  Be sure to call readyToSend() before adding to the buffer, and
    // save the result rather than calling it multiple times.
    return false;
  // Our local buffer is empty, so we can't say for certain if we can write
  // anything but we will try.
  return true;
}

bool TcpIpConnection::wantsRead()
{
  return _status == Open;
}

bool TcpIpConnection::wantsWrite()
{ // Listening for write is the preferred way to wait for an asynchronous
  // connect() to finish.  That works espeically well since we always try
  // to write immediately.  So we should wake up as soon as the 
  return (_status == Open) && !toSend.empty();
}

int TcpIpConnection::getHandle()
{
  return _handle;
}

static std::string s_TcpIpConnection_wakeUp = "TcpIpConnection::wakeUp()";
//static std::string s_EAGAIN_read = "TcpIpConnection <- EAGAIN on read()";
static std::string s_EINTR_read = "TcpIpConnection <- EINTR on read()";
static std::string s_read_0_no_write = "TcpIpConnection read 0 no write";
static std::string s_read_0 = "TcpIpConnection read 0";
static std::string s_read_1 = "TcpIpConnection read 1";
static std::string s_read = "TcpIpConnection read ";
static std::string s_EAGAIN_write = "TcpIpConnection <- EAGAIN on write()";
static std::string s_EINTR_write = "TcpIpConnection <- EINTR on write()";
static std::string s_wrote_all = "TcpIpConnection wrote all";
static std::string s_wrote_partial = "TcpIpConnection wrote partial";

void TcpIpConnection::wakeUp()
{
  ThreadMonitor::SetState tm(s_TcpIpConnection_wakeUp);
  tm.increment(s_TcpIpConnection_wakeUp);
  //std::cout<<"wakeUp()"<<std::endl;
  if (_status != Open)
    return;
  int successfulReadCount = 0;
  while (true)
  {
    const int READ_BUFFER_SIZE = 10000;
    char readBuffer[READ_BUFFER_SIZE];
    ssize_t result = read(_handle, readBuffer, READ_BUFFER_SIZE);
    //std::cout<<"read() -> "<<result<<std::endl;
    if (result == -1)
    { // Some type of error.
      if (errno == EAGAIN)
	// Try again soon, but not immediately.
	// Allow writing, if needed.
	// No need to report anything.  There will be a summary report below
	// using successfulReadCount.
	break;
      else if (errno == EINTR)
      { // Try again soon, but not immediately.
        // Allow writing, if needed.
        // I've never seen this happen, but it could.
        tm.increment(s_EINTR_read);
        break;
      }
      else
      { // Retrying does not make sense.  These are rare if they
        // happen at all.
        TclList msg;
        msg<<FLF<<_name<<"read()"<<errorString();
        sendToLogFile(msg);
        close();
        return;
      }
    }
    if (result == 0)
    { // End of file.  Stop and report it.
      close();
      return;
    }
    // Store it in the buffer.
    received.append(readBuffer, result);
    successfulReadCount++;

    if (successfulReadCount >= 25)
      // We don't want to be doing this for too long.  Just take a break and
      // we'll come back soon.
      break;
    
    // Try to read more.
  }
  switch (successfulReadCount)
  {
  case 0:
    if (toSend.empty())
      // This is bad.  We were woken up, but there is no work to do.  It might
      // be unavoidable sometimes, but it shouldn't happen a lot.
      tm.increment(s_read_0_no_write);
    else
      // There was nothing to read, but there was something to write.  This
      // happens a lot.  For simplicity we have just one wakeup call shared
      // by the readable and writable events.
      tm.increment(s_read_0);
    break;
  case 1:
    // This should be a common case.  We woke up.  We read something.  It all
    // fit in our buffer.
    tm.increment(s_read_1);
    break;
  default:
    // This is not an error.  I suspect it's not common, but we won't know
    // until we try this.  The amount of data available for reading was bigger
    // than our buffer, so we had to do more than one read.
    tm.increment(s_read + ntoa(successfulReadCount));
    break;
  }
  if (!toSend.empty())
  {
    ssize_t result = write(_handle,toSend.data(), toSend.size());
    //std::cout<<"write() -> "<<result<<std::endl;
    if (result == -1)
    { // Some type of error.
      if ((errno == EAGAIN) || (errno == EINTR))
      { // Try again soon, but not immediately.
        tm.increment((errno == EAGAIN)?s_EAGAIN_write:s_EINTR_write);
      }
      else
      { // Retrying does not make sense.
        // Reporting something to the log might make sense.
        TclList msg;
        msg<<FLF<<_name<<"write()"<<errorString();
        sendToLogFile(msg);
        close();
        return;
      }
    }
    else
    { // Remove whatever we sent from our buffer
      toSend.erase(0, result);
      tm.increment(toSend.empty()?s_wrote_all:s_wrote_partial);	
      // Record this.  This means that we must have gotten a confirmation to
      // our connect request from the remote server.
      _writeSucceeded = true;
    }
  }
}


/////////////////////////////////////////////////////////////////////
// ZLibOutputStream
/////////////////////////////////////////////////////////////////////

void ZLibOutputStream::compress(std::string const &source, bool flush)
{
  if (!_valid)
    return;
  const int outputBufferSize = 10000;
  char outputBuffer[outputBufferSize];
  int flushType = flush?Z_SYNC_FLUSH:Z_NO_FLUSH;
  //int realOutputCount = 0;
  bool outputBufferFull = false;
  // Why isn't this a constant pointer?  This cast must be right.
  _zlibBuffer.next_in = (Bytef *)source.data();
  _zlibBuffer.avail_in = source.size();
  _zlibBuffer.total_in = 0;
  while (_zlibBuffer.avail_in || outputBufferFull)
  {
    _zlibBuffer.next_out = (Bytef *)outputBuffer;
    _zlibBuffer.avail_out = outputBufferSize;
    _zlibBuffer.total_out = 0;
    int result = deflate(&_zlibBuffer, flushType);
    if (result != Z_OK)
    {
      _valid = false;
      return;
    }
    outputBufferFull = !_zlibBuffer.avail_out;
    if (_zlibBuffer.total_out)
    {
      output.append(outputBuffer, _zlibBuffer.total_out);
      //realOutputCount += _zlibBuffer.total_out;
      _zlibBuffer.total_out = 0;
    }
  }
}

ZLibOutputStream::ZLibOutputStream(int level)
{
  _zlibBuffer.zalloc = zlibMalloc;
  _zlibBuffer.zfree = zlibFree;
  _zlibBuffer.opaque = Z_NULL;
  int result = deflateInit(&_zlibBuffer, level);
  _valid = result == Z_OK;
}

ZLibOutputStream::~ZLibOutputStream()
{
  deflateEnd(&_zlibBuffer);
}

void ZLibOutputStream::add(void const *bytes, int count)
{
  add(std::string((char *)bytes, count));
}

void ZLibOutputStream::add(std::string bytes)
{
  if (bytes.empty())
    return;
  if (!_last.empty())
    compress(_last, false);
  _last = bytes;
}

void ZLibOutputStream::flush()
{
  if (_last.empty())
    return;
  compress(_last, true);
  _last.clear();
}


/////////////////////////////////////////////////////////////////////
// ZLibInputStream
/////////////////////////////////////////////////////////////////////

ZLibInputStream::ZLibInputStream()
{
  _zlibBuffer.zalloc = zlibMalloc;
  _zlibBuffer.zfree = zlibFree;
  _zlibBuffer.opaque = Z_NULL;
  int result = inflateInit(&_zlibBuffer);
  _valid = result == Z_OK;
}

ZLibInputStream::~ZLibInputStream()
{
  inflateEnd(&_zlibBuffer);
}
void ZLibInputStream::add(std::string const &bytes)
{
  add(bytes.data(), bytes.size());
}

void ZLibInputStream::add(void const *bytes, int count)
{
  if (!_valid)
    return;
  _zlibBuffer.next_in = (Bytef *)bytes;
  _zlibBuffer.avail_in = count;
  bool outputBufferFull = false;
  while ((_zlibBuffer.avail_in || outputBufferFull) && _valid)
  {
    const int bufferLength = 1024;
    char buffer[bufferLength];
    _zlibBuffer.next_out = (Bytef *)buffer;
    _zlibBuffer.avail_out = bufferLength;
    _zlibBuffer.total_out = 0;
    int result = inflate(&_zlibBuffer, Z_SYNC_FLUSH);
    // Note:  Z_BUF_ERROR is unavoidable in some cases.  This means that you
    // provided zlib with no new input, and it had no output for you.  This
    // happens when we are finished reading the input to zlib, the last
    // output from zlib perfectly filled the output buffer, and zlib has
    // no more data for us.
    if (!((result == Z_OK) || (outputBufferFull && (result == Z_BUF_ERROR))))
    {
      _valid = false;
      output.clear();
    }
    if (_zlibBuffer.total_out)
      output.append(buffer, _zlibBuffer.total_out);
    outputBufferFull = !_zlibBuffer.avail_out;
  }
}


/////////////////////////////////////////////////////////////////////
// TalkWithServer
/////////////////////////////////////////////////////////////////////


TalkWithServer::CancelId TalkWithServer::_lastCancelId = 0;

TalkWithServer::CancelId TalkWithServer::getNextCancelId()
{ // For simplicity make this atomic.  A __thread variable would probably
  // have worked, too.  Each client probably stays in one thread.  But it
  // was just as easy to do it this way and have fewer assumptions.
  return __sync_add_and_fetch(&_lastCancelId, 1);
}

TalkWithServer::TalkWithServer(std::string const &name) :
  _name(name),
  _tcpIpConnection(name),
  _status(NotYetOpen)
{
}

TalkWithServer::~TalkWithServer()
{
  assert(_validResponses.empty()
         || (_atServer.empty() && _localQueue.empty()));
}

void TalkWithServer::connect(std::string address, std::string port)
{
  assert(_status == NotYetOpen);
  _tcpIpConnection.connect(address, port, false);
  _status = Open;
  _tcpIpConnection.toSend =
    "command=set_output&mode=zlib\r\ncommand=set_input&mode=zlib\r\n";
  checkForClosed();
}

void TalkWithServer::checkForClosed()
{
  if (_tcpIpConnection.closed())
    disconnect();
}

void TalkWithServer::disconnect()
{
  if (_status != Closed)
  {
    _status = Closed;
    _tcpIpConnection.close();
    for (std::map< int, ServerWrapper >::iterator it = _atServer.begin();
         it != _atServer.end();
         it++)
    {
      LocalWrapper localWrapper;
      localWrapper.listener = it->second.listener;
      localWrapper.cancelId = it->second.cancelId;
      localWrapper.success = false;
      localWrapper.streaming = it->second.streaming;
      localWrapper.clientId = it->second.clientId;
      _localQueue.push_back(localWrapper);
    }
    _atServer.clear();
  }
}

bool TalkWithServer::disconnected()
{
  return _status == Closed;
}

void TalkWithServer::cancel(CancelId id)
{
  _validResponses.erase(id);
}

void TalkWithServer::cancelAll()
{
  _validResponses.clear();
}

int TalkWithServer::pendingResponseCount()
{
  return _localQueue.size();
}

void TalkWithServer::doResponses()
{ // It is possible that during this one of these callbacks, new callbacks
  // will be added to the list.  It's is unlikely that a valid response will
  // appear in the list, but it is certaily possible for someone to disconnect
  // the session in a callback, so the list will have new error responses.
  // It is also possible for a request to be canceled while in a callback.
  //
  // Strange.  The comment above makes sense, but the code below uses iterators
  // to walk through the list.  An iterator could be invalidated if the list
  // grows.  I wonder why I made that comment but didn't do anything about it.
  // TODO
  for (std::vector< LocalWrapper >::iterator it = _localQueue.begin();
       it != _localQueue.end();
       it++)
  {
    //TclList msg;
    //msg<<"doResponses"<<it->cancelId<<it->clientId;
    //if (it->success)
    //msg<<it->body;
    //std::cout<<(std::string)msg<<std::endl;
    if (_validResponses.count(it->cancelId))
    { // Not yet canceled.
      //std::cout<<"doResponses NOT CANCELED :)"<<std::endl;
      if (it->success)
        it->listener->onMessage(it->body, it->clientId, it->cancelId);
      else
        it->listener->onAbort(it->clientId, it->cancelId);
      if (!it->streaming)
        // Clean up.  No real functionality here, but without this
        // we might have a memory leak.
        _validResponses.erase(it->cancelId);
    }
  }
  _localQueue.clear();
}

void TalkWithServer::sendMessage(Message const &message)
{
  _outgoing.push_back(message);
}

TalkWithServer::CancelId TalkWithServer::sendMessage(Message message,
                                                     IMessageListener *listener, int clientId,
                                                     bool streaming)
{
  assert(listener);
  int messageId = 1;
  while (_atServer.count(messageId))
    messageId++;
  message["message_id"] = ntoa(messageId);
  ServerWrapper serverWrapper;
  serverWrapper.listener = listener;
  serverWrapper.cancelId = getNextCancelId();
  _validResponses.insert(serverWrapper.cancelId);
  serverWrapper.streaming = streaming;
  serverWrapper.clientId = clientId;
  _atServer[messageId] = serverWrapper;
  sendMessage(message);
  return serverWrapper.cancelId;
}

bool TalkWithServer::wantsRead()
{
  return _tcpIpConnection.wantsRead();
}

bool TalkWithServer::wantsWrite()
{
  return _tcpIpConnection.wantsWrite();
}

int TalkWithServer::getHandle()
{
  return _tcpIpConnection.getHandle();
}

void TalkWithServer::addStringWithSize(std::string toCompress)
{
  addInt32(toCompress.size());
  _zLibOutputStream.add(toCompress);
}

void TalkWithServer::addInt32(int32_t toCompress)
{
  _zLibOutputStream.add(&toCompress, 4);
}

void TalkWithServer::wakeUp()
{
  if (_tcpIpConnection.readyToSend() && !_outgoing.empty())
  { // Copy outgoing messages from here to the tcp/ip connection.
    // Don't do this before we have to.  Compression will work best if
    // we don't flush the buffer until the last moment.
    for (std::vector< Message >::const_iterator messageIt =
           _outgoing.begin();
         messageIt != _outgoing.end();
         messageIt++)
    {
      for (Message::const_iterator fieldIt = messageIt->begin();
           fieldIt != messageIt->end();
           fieldIt++)
      {
        addStringWithSize(fieldIt->first);
        addStringWithSize(fieldIt->second);
      }
      addInt32(0);
    }
    _outgoing.clear();
    _zLibOutputStream.flush();
    _tcpIpConnection.toSend += _zLibOutputStream.output;
    _zLibOutputStream.output.clear();
  }
  if (!_zLibOutputStream.valid())
  { // This should never happen!  Maybe out of memory?  Or a code error on
    // our part.
    TclList msg;
    msg<<FLF<<_name<<"!_zLibOutputStream.valid()";
    sendToLogFile(msg);
    disconnect();
    return;
  }
  _tcpIpConnection.wakeUp();
  if (!_tcpIpConnection.received.empty())
  {
    //std::cout<<"before receive, _zLibInputStream.size() = "
    //       <<_zLibInputStream.output.size()<<std::endl;
    _zLibInputStream.add(_tcpIpConnection.received);
    //std::cout<<"received "<<_tcpIpConnection.received.size()
    //       <<" bytes of compressed data"<<std::endl
    //       <<"_zLibInputStream.size() = "<<_zLibInputStream.output.size()
    //       <<std::endl;
    _tcpIpConnection.received.clear();
    if (!_zLibInputStream.valid())
    { // This probably means the data got corrupted.  Note that there is no
      // way to end the stream in our protocol; it's made to go on forever.
      // The best thing to do is to close the connection and try again.
      // Note that we abort immediately.  We do not attempt to read any
      // data which might have been received before the bad data.  zlib's
      // ability to detect bad data is very limited.  Often it gives us
      // extra stuff after the last of the good bytes.  Often the extra
      // stuff is a copy of earlier good bytes, so it looks like a good
      // message, but it is not.  The safest thing is to stop immediately.
      TclList msg;
      msg<<FLF<<_name<<"!_zLibInputStream.valid()";
      sendToLogFile(msg);
      disconnect();
      return;
    }
    int consumed = 0;
    int remaining = _zLibInputStream.output.size();
    char const *const start = _zLibInputStream.output.data();
    while (true)
    { // id followed by body size followed by body.
      //std::cout<<"top:  consumed="<<consumed<<
      //  ", remaining="<<remaining<<std::endl;
      if (remaining < 8)
        // Not enough space for a header.
        break;
      const int id = *(int32_t *)(start + consumed);
      //std::cout<<"id="<<id<<std::endl;
      const int size = *(int32_t *)(start + consumed + 4);
      //std::cout<<"size="<<size<<std::endl;
      if (size < 0)
      { // Invalid message.
        TclList msg;
        msg<<FLF<<_name<<"Invalid message header"<<id<<size;
        sendToLogFile(msg);
        disconnect();
        return;
      }
      if (remaining < 8 + size)
        // We have not loaded the entire message yet.
        break;
      if (ServerWrapper *serverWrapper = getProperty(_atServer, id))
      {
        LocalWrapper localWrapper;
        localWrapper.listener = serverWrapper->listener;
        // For simplicity, just copy the cancel id.  We will check for it
        // later.  It might be slightly more effecient to check for it
        // here, too, and delete some messages slightly sooner.
        localWrapper.cancelId = serverWrapper->cancelId;
        localWrapper.body = std::string(start + 8 + consumed, size);
        localWrapper.success = true;
        localWrapper.streaming = serverWrapper->streaming;
        localWrapper.clientId = serverWrapper->clientId;
        _localQueue.push_back(localWrapper);
        if (!serverWrapper->streaming)
          _atServer.erase(id);
        //std::cout<<"cancelId="<<localWrapper.cancelId
        //<<", clientId="<<localWrapper.clientId
        //<<", body="<<tclQuote(localWrapper.body)<<std::endl;
      }
      // Skip past this message when we look for the next one.
      consumed += size + 8;
      remaining -= size + 8;
    }
    // Remove every message we've handled from the buffer.  Remove them all
    // at once, rather than one at a time, so we don't have an O(n*n)
    // algorithm.  Possibly leave a partial message in the buffer.
    _zLibInputStream.output.erase(0, consumed);
  }
  checkForClosed();
}