#include <algorithm>
#include "ClientChannel.h"
#ifdef _WIN32
#  include <WinSock2.h>
#else
#  include <sys/socket.h>
#  include <netinet/in.h>
#  include <errno.h>
#  include <arpa/inet.h>
#  include <fcntl.h>
#  include <unistd.h>
#endif


struct FileDescriptorSet
{
  FileDescriptorSet()
  {
    FD_ZERO(&m_set);
  }
  FileDescriptorSet(int a_socket)
    : FileDescriptorSet()
  {
    AddSocket(a_socket);
  }
  void AddSocket(int a_socket)
  {
    FD_SET(a_socket, &m_set);
  }
  operator fd_set*() { return &m_set; }
  fd_set m_set;
};


ClientChannel::ClientChannel(int a_serverSocket)
  : m_serverSocket(a_serverSocket)
{
#ifdef _WIN32
  DWORD nonBlocking = 1;
  if (ioctlsocket(m_serverSocket, FIONBIO, &nonBlocking) != 0)
  {
    m_serverSocket = InvalidSocket;
  }
#else
  int flags = fcntl(m_serverSocket, F_GETFL, 0) | O_NONBLOCK;
  if (fcntl(m_serverSocket, F_SETFL, flags) != 0)
  {
    m_serverSocket = InvalidSocket;
  }
#endif
}


int ClientChannel::Select(bool a_read, bool a_write, bool a_except, std::chrono::milliseconds a_timeout)
{
  uint64_t timeOutMicroSeconds = a_timeout.count() * 1000LL;
  struct timeval tv = { static_cast<long>(timeOutMicroSeconds / 1000000LL), static_cast<int32_t>(timeOutMicroSeconds % 1000000LL) };
  fd_set rfds, wfds, xfds; 
  FD_ZERO(&rfds);
  FD_ZERO(&wfds);
  FD_ZERO(&xfds);
  if (a_read)   { FD_SET(m_serverSocket, &rfds); }
  if (a_write)  { FD_SET(m_serverSocket, &wfds); }
  if (a_except) { FD_SET(m_serverSocket, &xfds); }
  return FunctionWrapper(-1, ::select, m_serverSocket + 1, &rfds, &wfds, &xfds, &tv);
}


// Warning: it will be closed if you try to read <= 0 bytes
int ClientChannel::Read(std::vector<uint8_t>& a_recvBuffer, int a_readSize)
{
  return FunctionWrapper(0, ::recv, m_serverSocket, reinterpret_cast<char*>(a_recvBuffer.data()), std::min<int>(static_cast<int>(a_recvBuffer.size()), a_readSize), 0);
}


int ClientChannel::Write(const std::vector<uint8_t>& a_sendBuffer, int a_writeSize)
{
  return FunctionWrapper(0, ::send, m_serverSocket, reinterpret_cast<const char*>(a_sendBuffer.data()), std::min<int>(static_cast<int>(a_sendBuffer.size()), a_writeSize), 0);
}


ClientChannel::Result ClientChannel::SelectAndRead(std::vector<uint8_t>& a_recvBuffer, std::chrono::milliseconds a_timeout)
{
  size_t capacity = a_recvBuffer.capacity();
  size_t spaceAvailable = capacity - a_recvBuffer.size();
  if (!spaceAvailable)
  {
    return InsufficientCapacity; // No space to receive data
  }

  uint64_t timeOutMicroSeconds = a_timeout.count() * 1000LL;
  struct timeval tv = { static_cast<long>(timeOutMicroSeconds / 1000000LL), static_cast<int32_t>(timeOutMicroSeconds % 1000000LL) };

  FileDescriptorSet fdReadSet(m_serverSocket);
  if (FunctionWrapper(-1, ::select, m_serverSocket + 1, fdReadSet, nullptr, nullptr, &tv) <= 0)
  {
    return DataUnavailable; // No data ready to read
  }

  a_recvBuffer.resize(capacity);
  int bytesRead = FunctionWrapper(0, ::recv, m_serverSocket, reinterpret_cast<char*>(a_recvBuffer.data() + a_recvBuffer.size() - spaceAvailable), spaceAvailable, 0);
  a_recvBuffer.resize(capacity - spaceAvailable + std::max<int>(bytesRead, 0));
  return (bytesRead <= 0) ? Error : Success;
}


ClientChannel::Result ClientChannel::SelectAndWrite(const std::vector<uint8_t>& a_sendBuffer, int& a_offset, std::chrono::milliseconds a_timeout)
{
  size_t size = a_sendBuffer.size();
  int dataAvailable = size - a_offset;
  if (dataAvailable <= 0)
  {
    return DataUnavailable; // No data to send
  }

  uint64_t timeOutMicroSeconds = a_timeout.count() * 1000LL;
  struct timeval tv = { static_cast<long>(timeOutMicroSeconds / 1000000LL), static_cast<int32_t>(timeOutMicroSeconds % 1000000LL) };

  FileDescriptorSet fdWriteSet(m_serverSocket);
  if (FunctionWrapper(-1, ::select, m_serverSocket + 1, nullptr, fdWriteSet, nullptr, &tv) <= 0)
  {
    return InsufficientCapacity; // No space in buffer to send
  }

  int bytesWritten = FunctionWrapper(0, ::send, m_serverSocket, reinterpret_cast<const char*>(a_sendBuffer.data() + a_offset), dataAvailable, 0);
  a_offset += std::max<int>(bytesWritten, 0);
  return (bytesWritten <= 0) ? Error : Success;
}


void ClientChannel::Close()
{
#ifdef _WIN32
  FunctionWrapper(INT_MAX, ::closesocket, m_serverSocket);
#else
  FunctionWrapper(INT_MAX, ::close, m_serverSocket);
#endif
}


bool ClientChannel::IsConnected()
{
  return m_serverSocket != InvalidSocket;
}