#include "UdpSocket.h"

#ifdef PLATFORM_POSIX
#include <sys/ioctl.h>
#endif

// constructor
UdpSocket::UdpSocket(uint16_t port) : _port(port)
{
#ifdef PLATFORM_WIN32
	ONCELOCK {
		WSADATA wsadata;
		WSAStartup(MAKEWORD(2, 2), &wsadata);
	}
#endif

	// not connected
	_serverSocket = -1;

	// not in a multicast group
	_mcast = false;
}

// destructor
UdpSocket::~UdpSocket()
{
	// close if still connected
	if(_serverSocket != -1)
		Stop();
}

// send data
bool UdpSocket::SendTo(IPV4Address host, uint16_t port, String const &data)
{
	SOCKET_TYPE sock;
	struct sockaddr_in server;
	
	sock = socket(AF_INET, SOCK_DGRAM, 0);
	if(sock < 0)
		return false;
	
	memset(&server, 0, sizeof(struct sockaddr_in));
	server.sin_family = AF_INET;
	server.sin_addr.s_addr = host;
	server.sin_port = htons(_port);
	int n = sendto(sock, ~data, data.GetCount(), 0, (struct sockaddr *)&server, sizeof(struct sockaddr_in));
#ifdef PLATFORM_WIN32
	closesocket(sock);
#else
	close(sock);
#endif

	return n == data.GetCount();
}

// start listening for connections
bool UdpSocket::Listen()
{
	struct sockaddr_in server;

	// create the socket
	_serverSocket = socket(AF_INET, SOCK_DGRAM, 0);
	
	// check for error
	if(_serverSocket < 0)
	{
		_serverSocket = -1;
		return false;
	}
	
	// allow multiple sockets to use the same PORT number
	// (needed for multicast to work...)
	int yes = 1;
	if (setsockopt(_serverSocket, SOL_SOCKET, SO_REUSEADDR, (const char *)&yes, sizeof(yes)) < 0)
	{
#ifdef PLATFORM_WIN32
		closesocket(_serverSocket);
#else
		close(_serverSocket);
#endif
		_serverSocket = -1;
		return false;
	}

	// put socket in non-blocking mode
#ifdef PLATFORM_POSIX
	int flags = fcntl(_serverSocket, F_GETFL, 0);
	if(flags == -1)
		flags = 0;
	fcntl(_serverSocket, F_SETFL, flags | O_NONBLOCK);
#else
	u_long iMode=1;
	ioctlsocket(_serverSocket, FIONBIO, &iMode);
#endif
	
	// bind to port
	memset(&server, 0, sizeof(server));
	server.sin_family = AF_INET;
	server.sin_addr.s_addr = INADDR_ANY;
	server.sin_port = htons(_port);
	if(bind(_serverSocket, (struct sockaddr *)&server, sizeof(server)) < 0)
	{
#ifdef PLATFORM_WIN32
		closesocket(_serverSocket);
#else
		close(_serverSocket);
#endif
		_serverSocket = -1;
		return false;
	}
	return true;
}

// check if data is available from server
int UdpSocket::Available(void)
{
	if(_serverSocket == -1)
		return 0;
	
#ifdef PLATFORM_POSIX
	int count;
	ioctl(_serverSocket, FIONREAD, &count);
#else
	u_long count;
	ioctlsocket(_serverSocket, FIONREAD, &count);
#endif
	return count;
}

// read data from socket - return 0 bytes if none is available
// on host and port the corresponding client's ones
int UdpSocket::Receive(IPV4Address &host, uint16_t &port, String &data)
{
	if(_serverSocket == -1)
	{
		data.Clear();
		return 0;
	}

	int avail = Available();
	if(!avail)
		return 0;
	
	StringBuffer buf(avail);
	struct sockaddr_in from;
	socklen_t fromlen = sizeof(struct sockaddr_in);
	int n = recvfrom(_serverSocket, ~buf, avail, 0,(struct sockaddr *)&from, &fromlen);
	if(n < 0)
	{
		data.Clear();
		return 0;
	}
	host = from.sin_addr.s_addr;
	port = ntohs(from.sin_port);
	data = buf;
	return n;
}

// stop listening for connections
bool UdpSocket::Stop(void)
{
	// if not connected, just do nothing
	if(_serverSocket == -1)
	{
		_mcast = false;
		return true;
	}
	
	// if we're into an IGMP group, leave it
	if(_mcast)
		IGMPLeave();
	
#ifdef PLATFORM_WIN32
		closesocket(_serverSocket);
#else
		close(_serverSocket);
#endif
	_serverSocket = -1;
	return true;
}

// join multicast group
bool UdpSocket::IGMPJoin(IPV4Address addr)
{
	// if already joined a multicast group, leave it before
	if(_mcast)
		IGMPLeave();
	
	struct ip_mreq mreq;
	
	// if still not listening, start listening now
	if(_serverSocket < 0)
		if(!Listen())
			return false;
	
	// use setsockopt() to request that the kernel join a multicast group */
	mreq.imr_multiaddr.s_addr = addr;
	mreq.imr_interface.s_addr = htonl(INADDR_ANY);
	if (setsockopt(_serverSocket, IPPROTO_IP, IP_ADD_MEMBERSHIP, (const char *)&mreq, sizeof(mreq)) < 0)
		return false;
	
	// remember we're in a group
	_mcast = true;
	_mcastAddr = addr;
	
	return true;
}

// leave multicast group
bool UdpSocket::IGMPLeave(void)
{
	struct ip_mreq mreq;

	if(!_mcast)
		return true;
	
	// mark that we're no more in a group, even if leave fails
	_mcast = false;

	if(_serverSocket < 0)
		return false;

	// leave the group
	mreq.imr_multiaddr.s_addr = _mcastAddr;
	mreq.imr_interface.s_addr = htonl(INADDR_ANY);
	if (setsockopt(_serverSocket, IPPROTO_IP, IP_DROP_MEMBERSHIP, (const char *)&mreq, sizeof(mreq)) < 0)
		return false;
	return true;
}
