/////////////////////////////////////////////////////////////////
// Sockets.cpp - C++ wrapper for C socket api                  //
// ver 4.0                                                     //
// Jim Fawcett, CSE687 - Object Oriented Design, Spring 2015   //
/////////////////////////////////////////////////////////////////

#include "Sockets.h"
#include <iostream>
#include <sstream>
#include <thread>
#include <memory>
#include <functional>

#include "../ApplicationHelpers/AppHelpers.h"
#include "../WindowsHelpers/WindowsHelpers.h"

using namespace ApplicationHelpers;
using namespace WindowsHelpers;

const bool always = true;

/////////////////////////////////////////////////////////////////////////////
// SocketSystem class members

//----< constructor starts up sockets by loading winsock lib >---------------

SocketSystem::SocketSystem()
{
  int iResult = WSAStartup(MAKEWORD(2, 2), &wsaData);
  if (iResult != 0) {
    Verbose::show("WSAStartup failed with error = " + toString(iResult), always);
  }
}
//-----< destructor frees winsock lib >--------------------------------------

SocketSystem::~SocketSystem()
{
  WSACleanup();
}

/////////////////////////////////////////////////////////////////////////////
// Socket class members

//----< constructor sets TCP protocol and Stream mode >----------------------

Socket::Socket()
{
  ZeroMemory(&hints, sizeof(hints));
  hints.ai_family = AF_UNSPEC;
  hints.ai_socktype = SOCK_STREAM;
  hints.ai_protocol = IPPROTO_TCP;
}
//----< promotes Win32 socket to Socket >------------------------------------

Socket::Socket(::SOCKET sock) : socket_(sock)
{
  ZeroMemory(&hints, sizeof(hints));
  hints.ai_family = AF_UNSPEC;
  hints.ai_socktype = SOCK_STREAM;
  hints.ai_protocol = IPPROTO_TCP;
}
//----< transfer socket ownership with move constructor >--------------------

Socket::Socket(Socket&& s)
{
  socket_ = s.socket_;
  s.socket_ = INVALID_SOCKET;
  ZeroMemory(&hints, sizeof(hints));
  hints.ai_family = AF_UNSPEC;
  hints.ai_socktype = SOCK_STREAM;
  hints.ai_protocol = IPPROTO_TCP;
}
//----< transfer socket ownership with move assignment >---------------------

Socket& Socket::operator=(Socket&& s)
{
  if (this == &s) return *this;
  socket_ = s.socket_;
  s.socket_ = INVALID_SOCKET;
  return *this;
  ZeroMemory(&hints, sizeof(hints));
  hints.ai_family = AF_UNSPEC;
  hints.ai_socktype = SOCK_STREAM;
  hints.ai_protocol = IPPROTO_TCP;
}
//----< close connection >---------------------------------------------------

void Socket::close()
{
  ::closesocket(socket_);
}
//----< tells receiver there will be no more sends from this socket >--------

bool Socket::shutDownSend()
{
  ::shutdown(socket_, SD_SEND);
  if (socket_ != INVALID_SOCKET)
    return true;
  return false;
}

//----< tells receiver this socket won't call receive anymore >--------------

bool Socket::shutDownRecv()
{
  ::shutdown(socket_, SD_RECEIVE);
  if (socket_ != INVALID_SOCKET)
    return true;
  return false;
}
//----< tells receiver there will be no more sends or recvs >----------------

bool Socket::shutDown()
{
  ::shutdown(socket_, SD_BOTH);
  if (socket_ != INVALID_SOCKET)
    return true;
  return false;

}
//----< destructor closes socket handle >------------------------------------

Socket::~Socket() {
  close();
}
//----< builds a std::vector<byte> from a std::string >----------------------
/*
 *   Makes a Block with last byte being a null terminator
 */
Socket::Block makeBlock(const std::string& str)
{
  Socket::Block blk;
  for (auto ch : str)
    blk.push_back(ch);
  blk.push_back('\0');
  return blk;
}
//----< extracts std::string from Block >------------------------------------

std::string blockToString(const Socket::Block& blk)
{
  std::string ret;
  for (auto ch : blk)
    ret += ch;
  return ret;
}
//----< sends a block >------------------------------------------------------
/*
 *  Doesn't return until entire Block has been sent
 *
 * ToDo:
 * - I'm thinking about removing the Block transfers or changing
 *   them to use another termination strategy.
 */
bool Socket::send(const Block& blk)
{
  size_t bytesSent, bytesRemaining = blk.size();
  const byte* pBuf = &(*blk.begin());
  while (bytesRemaining > 0)
  {
    bytesSent = ::send(socket_, pBuf, bytesRemaining, 0);
    if (bytesSent == INVALID_SOCKET)
      return false;
    bytesRemaining -= bytesSent;
    pBuf += bytesSent;
  }
  return true;
}
//----< sends a null terminated string >-------------------------------------
/*
 *  Doesn't return until entire string has been sent
 */
bool Socket::sendString(const std::string& str)
{
  size_t bytesSent, bytesRemaining = str.size();
  const byte* pBuf = &(*str.begin());
  while (bytesRemaining > 0)
  {
    bytesSent = ::send(socket_, pBuf, bytesRemaining, 0);
    if (bytesSent == INVALID_SOCKET)
      return false;
    bytesRemaining -= bytesSent;
    pBuf += bytesSent;
  }
  ::send(socket_, '\0', 1, 0);
  return true;
}
//----< recv a Block = std::vector<byte> >-----------------------------------
/*
 * Doesn't return until a null byte is received
 *
 * ToDo:
 * - I'm thinking about removing the Block transfers or changing
 *   them to use another termination strategy.
 */
Socket::Block Socket::recv()
{
  static const int buflen = 1;
  char buffer[1];
  Block blk;
  while (true)
  {
    iResult = ::recv(socket_, buffer, buflen, 0);
    if (iResult == 0)
      break;
    blk.push_back(buffer[0]);
    if (buffer[0] == '\0')
      break;
  }
  return blk;
}
//----< receives null-terminated string >------------------------------------
/*
 * ToDo:
 * - needs reads of one byte to be replaced by bulk reads into a
 *   stream buffer to improve efficiency.
 * - That will require building a circular buffer.
 */
std::string Socket::recvString()
{
  static const int buflen = 1;
  char buffer[1];
  std::string str;
  while (true)
  {
    iResult = ::recv(socket_, buffer, buflen, 0);
    if (iResult == 0)
      break;
    str += buffer[0];
    if (buffer[0] == '\0')
      break;
  }
  return str;
}
//----< returns bytes available in recv buffer >-----------------------------

size_t Socket::bytesWaiting()
{
  unsigned long int ret;
  ::ioctlsocket(socket_, FIONREAD, &ret);
  return (size_t)ret;
}
//----< waits for server data, checking every timeToCheck millisec >---------

bool Socket::waitForData(size_t timeToWait, size_t timeToCheck)
{
  size_t MaxCount = timeToWait / timeToCheck;
  static size_t count = 0;
  while (bytesWaiting() == 0)
  {
    if (++count < MaxCount)
      ::Sleep(timeToCheck);
    else
      return false;
  }
  return true;
}
/////////////////////////////////////////////////////////////////////////////
// SocketConnector class members

//----< request to connect to ip and port >----------------------------------

bool SocketConnecter::connect(const std::string& ip, size_t port)
{
  //size_t uport = htonl(port);
  size_t uport = port;
  std::string sPort = toString(uport);

  // Resolve the server address and port
  const char* pTemp = ip.c_str();
  iResult = getaddrinfo(pTemp, sPort.c_str(), &hints, &result);  // was DEFAULT_PORT
  if (iResult != 0) {
    Verbose::show("\n  -- getaddrinfo failed with error: " + toString(iResult), always);
    return false;
  }

  // Attempt to connect to an address until one succeeds
  for (ptr = result; ptr != NULL; ptr = ptr->ai_next) {

    char ipstr[INET6_ADDRSTRLEN];
    void *addr;
    char *ipver;

    // get pointer to address - different fields in IPv4 and IPv6:

    if (ptr->ai_family == AF_INET) { // IPv4
      struct sockaddr_in *ipv4 = (struct sockaddr_in *)ptr->ai_addr;
      addr = &(ipv4->sin_addr);
      ipver = "IPv4";
    }
    else { // IPv6
      struct sockaddr_in6 *ipv6 = (struct sockaddr_in6 *)ptr->ai_addr;
      addr = &(ipv6->sin6_addr);
      ipver = "IPv6";
    }

    // convert the IP to a string and print it:
    inet_ntop(ptr->ai_family, addr, ipstr, sizeof ipstr);
    printf("  %s: %s\n", ipver, ipstr);

    // Create a SOCKET for connecting to server
    socket_ = socket(ptr->ai_family, ptr->ai_socktype,
      ptr->ai_protocol);
    if (socket_ == INVALID_SOCKET) {
      int error = WSAGetLastError();
      Verbose::show("\n  -- socket failed with error: " + toString<int>(error), always);
      return false;
    }

    iResult = ::connect(socket_, ptr->ai_addr, (int)ptr->ai_addrlen);
    if (iResult == SOCKET_ERROR) {
      closesocket(socket_);
      socket_ = INVALID_SOCKET;
      continue;
    }
    break;
  }

  freeaddrinfo(result);

  if (socket_ == INVALID_SOCKET) {
    int error = WSAGetLastError();
    Verbose::show("\n  -- unable to connect to server, error = " + toString(error));
    return false;
  }
  return true;
}
/////////////////////////////////////////////////////////////////////////////
// SocketListener class members

//----< constructs SocketListener, specifying type of protocol to use >------

SocketListener::SocketListener(size_t port, IpVer ipv) : Socket(ipv), port_(port)
{
  if (ipv == Socket::IP6)
    hints.ai_family = AF_INET6;       // use this if you want an IP6 address
  else
    hints.ai_family = AF_INET;        // this gives IP4 address
}
//----< binds SocketListener to a network adddress on local machine >--------

bool SocketListener::bind()
{
  // Resolve the server address and port

  //size_t uport = ::htonl(port_);
  size_t uport = port_;
  std::string sPort = toString(uport);
  iResult = getaddrinfo(NULL, sPort.c_str(), &hints, &result);
  if (iResult != 0) {
    Verbose::show("\n  -- getaddrinfo failed with error: " + toString(iResult), always);
    return false;
  }

  // Iterate through all results and bind to first available

  for (auto pResult = result; pResult != NULL; pResult = pResult->ai_next)
  {
    // Create a SOCKET for connecting to server

    socket_ = socket(pResult->ai_family, pResult->ai_socktype, pResult->ai_protocol);
    if (socket_ == INVALID_SOCKET) {
      int error = WSAGetLastError();
      Verbose::show("\n  -- socket failed with error: " + toString(error), always);
      continue;
    }
    Verbose::show("server created ListenSocket");

    // Setup the TCP listening socket

    iResult = ::bind(socket_, pResult->ai_addr, (int)pResult->ai_addrlen);
    if (iResult == SOCKET_ERROR) {
      int error = WSAGetLastError();
      Verbose::show("\n  -- bind failed with error: " + toString(error), always);
      closesocket(socket_);
      continue;
    }
    else
    {
      break;
    }
  }
  return true;
}
//----< put SocketListener in listen mode, doesn't block >-------------------

bool SocketListener::listen()
{
  iResult = ::listen(socket_, SOMAXCONN);
  if (iResult == SOCKET_ERROR) {
    int error = WSAGetLastError();
    Verbose::show("\n  -- listen failed with error: " + toString(error), always);
    closesocket(socket_);
    return false;
  }
  Verbose::show("server TCP listening socket setup complete");
  return true;
}
//----< accepts incoming requrests to connect - blocking call >--------------

Socket SocketListener::accept()
{
  ::SOCKET sock = ::accept(socket_, NULL, NULL);
  Socket clientSocket = sock;    // uses Socket(::SOCKET) promotion ctor
  if (!clientSocket.validState()) {
    int error = WSAGetLastError();
    Verbose::show("\n  -- server accept failed with error: " + toString(error), always);
    return clientSocket;
  }
  return clientSocket;
}
//----< request SocketListener to stop accepting connections >---------------

void SocketListener::stop()
{
  stop_.exchange(true);
}
//----< SocketListener start function runs listener on its own thread >------
/*
 *  - Accepts Callable Object that defines the operations
 *    to handle client requests.
 *  - You will find an example Callable Object, ClientProc,
 *    used in the test stub below
*/
template<typename CallObj>
bool SocketListener::start(CallObj& co)
{
  if (!bind())
  {
    return false;
  }

  if (!listen())
  {
    return false;
  }

  std::thread ListenThread(
    [&]()
  {
    Verbose::show("server waiting for connection");

    while (true)
    {
      if (stop_.load())
        break;

      // Accept a client socket - blocking call

      Socket clientSocket = accept();    // uses move ctor
      if (!clientSocket.validState()) {
        continue;
      }
      Verbose::show("server accepted connection");

      co.takeSocket(std::move(clientSocket));  // transfer ownership of client socket to co
      std::thread clientThread(std::ref(co));  // start a thread to handle client request
      clientThread.detach();                   // detach because we won't need it again
    }
  }
  );
  ListenThread.detach();  // SocketListener has been stopped so detach thread
  return true;
}

/////////////////////////////////////////////////////////////////////////////
// Server's Client Processing handler class
// - must be callable object
// - passed to SocketListener::start(CallObject& co)
/*
 *  ToDo:
 *  - operator() currently uses Win32 socket operations
 *  - will change this to use Socket class operations.
 */
class ClientProc
{
public:
  ClientProc() : socket_(INVALID_SOCKET) {}

  void takeSocket(Socket s) { socket_ = std::move(s); }

  void operator()()  // this is function that defines thread operations
  {
    int iResult, iSendResult;
    ::SOCKET ClientSocket = socket_;
    static const int recvbuflen = DEFAULT_BUFLEN;
    char recvbuf[recvbuflen];

    do {
      iResult = recv(ClientSocket, recvbuf, recvbuflen, 0);
      if (iResult > 0) {
        Verbose::show("bytes recvd at server: " + toString(iResult));

        if (iResult < recvbuflen)
          recvbuf[iResult] = '\0';
        else
          recvbuf[recvbuflen - 1] = '\0';
        std::string msg = std::string(recvbuf);
        std::cout << "\n  server recvd: " << msg;

        // Echo the buffer back to the sender

        iSendResult = send(ClientSocket, recvbuf, iResult + 1, 0);
        if (iSendResult == SOCKET_ERROR) {
          int error = WSAGetLastError();
          Verbose::show("\n  -- server send failed with error: " + toString(error), always);
          closesocket(ClientSocket);
          return;
        }
        Verbose::show("bytes sent by server: " + toString(iSendResult), always);
      }
      else if (iResult == 0)
        Verbose::show("server connection closing");
      else  {
        int error = WSAGetLastError();
        Verbose::show("\n  -- server recv failed with error: " + toString(error), always);
        closesocket(ClientSocket);
        return;
      }
    } while (iResult > 0);
  }
private:
  Socket socket_;
};

struct Cosmetic
{
  ~Cosmetic()
  {
    std::cout << "\n  press key to exit: ";
    std::cin.get();
    std::cout << "\n\n";
  }
} aGlobalForCosmeticAction;

#ifdef TEST_SOCKETS

int main()
{
  title("Testing Sockets", '=');

  Verbose v(true);
  SocketSystem ss;
  SocketConnecter si;
  SocketListener sl(8080, Socket::IP6);
  ClientProc cp;
  sl.start(cp);

  while (!si.connect("localhost", 8080))
  {
    std::cout << "\n  client waiting to connect";
    ::Sleep(100);
  }
  std::string text = "Hello World!";
  for (size_t i = 0; i < 5; ++i)
  {
    // Test send(Block blk)

    si.send(makeBlock(text));
    
    // Test sendString(text)

    //si.sendString(text);
    std::cout << "\n  client sending: " << text;
  }

  std::cout << "\n\n  client calling send shutdown\n";
  si.shutDownSend();

  while (true)
  {
   if (si.waitForData(1000, 100))
    {
      // Test recv Block

      Socket::Block blk = si.recv();
      if (blk.size() == 0)
        break;
      std::string temp = blockToString(blk);
      if (temp.size() > 0 && temp[0] != '\0')
        std::cout << "\n  client recvd: " << temp;

      // Test recvString()

      //std::string str = si.recvString();
      //if (str.size() == 0)
      //  break;
      //std::cout << "\n  " << str;
   }
    else
    {
      std::cout << "\n  client waiting for server time-out";
      break;
    }
  }
}

#endif