#include "Socket.h"

#include "WinSock2.h"
#include "WS2tcpip.h"
#include "Windows.h"

#include <exception>
#include <cassert>


static bool wsa_initialized = false;


Socket::Socket()
{
	// Move out to a separate function
	if (!wsa_initialized)
	{
		WSAData data;
		if (WSAStartup(MAKEWORD(2, 2), &data))
			throw std::exception("Failed to initialize WSA");
		wsa_initialized = true;
	}

	_sock = socket(
		AF_INET,
		SOCK_STREAM,
		0 // IPPROTO_TCP
	);

	if (_sock == INVALID_SOCKET)
		throw std::exception("Failed to create socket");
}

Socket::~Socket()
{
	if (!_rc.last())
		return;
	// Flush comunication before closing
	shutdown(_sock, SD_SEND);
	char buffer[512];
	while (read(buffer, sizeof(buffer))) {}
	shutdown(_sock, SD_RECEIVE);

	closesocket(_sock);
}

void Socket::bind(const char *address, unsigned short port)
{
	sockaddr_in addr = { 0 };

	addr.sin_family = AF_INET;
	addr.sin_port = htons(port);
	if (address)
		inet_pton(AF_INET, address, &addr.sin_addr);
	else
		addr.sin_addr.s_addr = htonl(INADDR_ANY);

	if (::bind(_sock, (sockaddr*)&addr, sizeof(addr)) == SOCKET_ERROR)
		throw std::exception("Failed to bind socket");
}

void Socket::listen(int backlog)
{
	if (::listen(_sock, backlog) == SOCKET_ERROR)
		throw std::exception("Failed to listen on socket");
}

Socket Socket::accept()
{
	SOCKET sock = ::accept(_sock, nullptr, nullptr);
	if (sock == INVALID_SOCKET)
		throw std::exception("Failed to accept socket");
	Socket socket(sock);
	return std::move(socket);
}

int Socket::read(void* buffer, int size)
{
	int ret = ::recv(_sock, (char*)buffer, size, 0);
	return max(0, ret); // Clamp to 0, as socket is always synchronious
}

int Socket::write(const char *string)
{
	return write(string, (int)strlen(string));
}

int Socket::write(const void *data, int size)
{
	return ::send(_sock, (const char *)data, size, 0);
}

void Socket::end()
{
	::shutdown(_sock, SD_SEND);
}


SocketStreamBuf::~SocketStreamBuf()
{
	sync();
}

int SocketStreamBuf::underflow()
{
	// System should be buffering received data
	if (!_sock.read(&_c, 1))
		return EOF;
	setg(&_c, &_c, &_c+1);
	return _c;
}

int SocketStreamBuf::overflow(int c)
{
	if (_buf_size >= _buf_size_max)
		sync();

	// This check is not required, but it muffles a warning in VisualStudio
	if (_buf_size < _buf_size_max)
		_buf[_buf_size++] = c;
	else
		assert(0);

	return c;
}

int SocketStreamBuf::sync()
{
	_sock.write(_buf, _buf_size);
	_buf_size = 0;
	return 0;
}