/*]  Copyright (c) 2009-2010, Charles McGarvey  [**************************
**]  All rights reserved.
*
* vi:ts=4 sw=4 tw=75
*
* Distributable under the terms and conditions of the 2-clause BSD license;
* see the file COPYING for a complete text of the license.
*
**************************************************************************/

#ifndef _MOOF_SOCKET_HH_
#define _MOOF_SOCKET_HH_

#include <cstring>
#include <sstream>
#include <string>
#include <vector>

#if defined(_WIN32)
#include <winsock2.h>
#include <ws2tcpip.h>
#include <wspiapi.h>
#else
#include <arpa/inet.h>
#include <netdb.h>
#include <sys/socket.h>
#include <sys/types.h>
#if HAVE_FCNTL_H
#include <fcntl.h>
#else
#include <sys/ioctl.h>
#endif
#endif

#include <Moof/Log.hh>
#include <Moof/Packet.hh>
#include <Moof/Thread.hh>


#ifndef SO_NONBLOCK
#define SO_NONBLOCK 1024
#endif


namespace Mf {


class SocketAddress
{
public:

	SocketAddress() :
		mSize(0),
		mType(0)
	{
		mAddress.sa.sa_family = AF_UNSPEC;
		mAddress.v4.sin_port = 0;
	}

	SocketAddress(const std::string& service, const std::string& name,
			int type = SOCK_STREAM, int family = AF_UNSPEC)
	{
		init(service, name, type, family);
	}

	SocketAddress(const std::string& service,
			int type = SOCK_STREAM, int family = AF_UNSPEC)
	{
		init(service, type, family);
	}

	SocketAddress(const struct addrinfo* addr, const std::string& name)
	{
		mType = addr->ai_socktype;
		memcpy(&mAddress.sa, addr->ai_addr, addr->ai_addrlen);
		mName = name;
		mSize = addr->ai_addrlen;
	}

	SocketAddress(const struct sockaddr* addr, size_t size,
			int type = SOCK_STREAM)
	{
		mType = type;
		memcpy(&mAddress.sa, addr, size);
		mSize = size;
		setNameFromAddress();
	}


	static SocketAddress broadcast(const std::string& service)
	{
		std::istringstream stream(service);
		unsigned short port;
		stream >> port;

		struct sockaddr_in addr;
		addr.sin_family = AF_INET;
		addr.sin_port = htons(port);
		addr.sin_addr.s_addr = INADDR_BROADCAST;
		memset(&addr.sin_zero, 0, sizeof(addr.sin_zero));
		return SocketAddress((sockaddr*)&addr, sizeof(addr), SOCK_DGRAM);
	}


	void init(const std::string& service, const std::string& name = "",
			int type = SOCK_STREAM, int family = AF_UNSPEC)
	{
		ASSERT(type == SOCK_STREAM || type == SOCK_DGRAM);
		ASSERT(family == AF_INET || family == AF_INET6 || family == AF_UNSPEC);

		struct addrinfo hints;
		memset(&hints, 0, sizeof(hints));
		hints.ai_family = family;
		hints.ai_socktype = type;
		hints.ai_flags = AI_PASSIVE;

		struct addrinfo* addr;
		int status = getaddrinfo(name.length() > 0 ? name.c_str() : 0,
				service.c_str(), &hints, &addr);
		if (status == 0)
		{
			mType = addr->ai_socktype;
			memcpy(&mAddress.sa, addr->ai_addr, addr->ai_addrlen);
			mSize = addr->ai_addrlen;

			if (name != "") mName = name;
			else            setNameFromAddress();

			freeaddrinfo(addr);
		}
		else
		{
			Mf::logWarning(gai_strerror(status));
			mType = 0;
			mSize = 0;
			mAddress.sa.sa_family = AF_UNSPEC;
			mAddress.v4.sin_port = 0;
		}
	}

	void init(const std::string& service,
			int type = SOCK_STREAM, int family = AF_UNSPEC)
	{
		init(service, "", type, family);
	}


	const std::string& name() const
	{
		return mName;
	}

	void setName(const std::string& name)
	{
		mName = name;
	}

	unsigned short port() const
	{
		return ntohs(mAddress.v4.sin_port);
	}

	int type() const
	{
		return mType;
	}

	int family() const
	{
		return mAddress.sa.sa_family;
	}


	const struct sockaddr* address() const
	{
		return mSize != 0 ? &mAddress.sa : 0;
	}

	size_t size() const
	{
		return mSize;
	}


	static int resolve(const std::string& service, const std::string& name,
			int type, int family, std::vector<SocketAddress>& resolved)
	{
		ASSERT(type == SOCK_STREAM || type == SOCK_DGRAM);
		ASSERT(family == AF_INET || family == AF_INET6 || family == AF_UNSPEC);

		resolved.clear();

		struct addrinfo hints;
		memset(&hints, 0, sizeof(hints));
		hints.ai_family = family;
		hints.ai_socktype = type;
		hints.ai_flags = AI_PASSIVE;

		struct addrinfo* list;
		int status = getaddrinfo(name.length() > 0 ? name.c_str() : 0,
				service.length() > 0 ? service.c_str() : 0, &hints, &list);
		if (status == 0)
		{
			for (struct addrinfo* addr = list;
					addr != 0; addr = addr->ai_next)
			{
				resolved.push_back(SocketAddress(addr, name));
			}

			freeaddrinfo(list);
		}
		else
		{
			Mf::logWarning(gai_strerror(status));
			return -1;
		}

		return 0;
	}


private:

	void setNameFromAddress()
	{
#if defined(_WIN32)
		// inet_ntop was introduced in Vista
		mName = inet_ntoa(mAddress.v4.sin_addr);
#else
		char name[INET6_ADDRSTRLEN] = {'\0'};
		inet_ntop(mAddress.sa.sa_family, &mAddress.sa, name, sizeof(name));
		mName = name;
#endif
	}


	union
	{
		sockaddr			sa;
		sockaddr_in			v4;
		sockaddr_in6		v6;
		sockaddr_storage	storage;
	}			mAddress;
	size_t		mSize;
	std::string	mName;
	int			mType;
};


class Socket
{
public:

	Socket(const SocketAddress& address) :
		mFd(-1),
		mIsConnected(false),
		mAddress(address)
	{
		mFd = socket(address.family(), address.type(), 0);
	}

	~Socket()
	{
#if defined(_WIN32)
		if (mFd != -1) closesocket(mFd);
#else
		if (mFd != -1) close(mFd);
#endif
	}


	bool isConnected() const
	{
		return mIsConnected;
	}

	const SocketAddress& address() const
	{
		return mAddress;
	}


	int connect()
	{
		int result = ::connect(mFd, mAddress.address(), mAddress.size());
		mIsConnected = result != -1;
		return result;
	}

	int bind()
	{
		return ::bind(mFd, mAddress.address(), mAddress.size());
	}

	int listen(int backlog = SOMAXCONN)
	{
		return ::listen(mFd, backlog > 0 ? backlog : SOMAXCONN);
	}

	Socket accept()
	{
		return Socket(mFd);
	}


	int set(int option, int value = 0)
	{
		if (option == SO_NONBLOCK)
		{
#ifdef HAVE_FCNTL
			int flags = fcntl(mFd, F_GETFL);
			return fcntl(mFd, F_SETFL, (value ? O_NONBLOCK : 0) | flags);
#else
			return ioctl(mFd, FIONBIO, value);
#endif
		}
		return setsockopt(mFd, SOL_SOCKET, option, &value, sizeof(value));
	}

	int set(int option, const std::string& value)
	{
		return setsockopt(mFd, SOL_SOCKET, option,
				value.data(), value.length());
	}

	int get(int option, int& value)
	{
		if (option == SO_NONBLOCK)
		{
#ifdef HAVE_FCNTL
			int flags = fcntl(mFd, F_GETFL);
			return flags & O_NONBLOCK;
#else
			return ioctl(mFd, FIONBIO, &value);
#endif
		}
		socklen_t optlen = sizeof(value);
		return getsockopt(mFd, SOL_SOCKET, option, &value, &optlen);
	}

	int get(int option, std::string& value)
	{
		char str[64] = {'\0'};
		socklen_t optlen = sizeof(str);
		int result = getsockopt(mFd, SOL_SOCKET, option, &str, &optlen);
		value = str;
		return result;
	}


	ssize_t write(const void* bytes, size_t size)
	{
		return send(mFd, bytes, size, 0);
	}
	ssize_t write(const void* bytes, size_t size,
		const SocketAddress& address)
	{
		return sendto(mFd, bytes, size, 0,
			address.address(), address.size());
	}

	ssize_t write(const Packet& packet)
	{
		return write(packet.bytes(), packet.size());
	}

	ssize_t write(const Packet& packet, const SocketAddress& address)
	{
		return write(packet.bytes(), packet.size(), address);
	}


	ssize_t read(void* bytes, size_t size)
	{
		return recv(mFd, bytes, size, 0);
	}

	ssize_t read(void* bytes, size_t size, SocketAddress& address)
	{
		union
		{
			sockaddr			sa;
			sockaddr_storage	storage;
		} addr;
		socklen_t length = sizeof(addr);

		ssize_t result = recvfrom(mFd, bytes, size, 0, &addr.sa, &length);
		if (result != -1)
		{
			address = SocketAddress(&addr.sa, length, mAddress.type());
		}
		return result;
	}

	ssize_t read(Packet& packet)
	{
		char buffer[65536];
		ssize_t result = read(buffer, sizeof(buffer));
		if (result != -1) packet = Packet(buffer, result);
		return result;
	}

	ssize_t read(Packet& packet, SocketAddress& address)
	{
		char buffer[65536];
		ssize_t result = read(buffer, sizeof(buffer), address);
		if (result != -1) packet = Packet(buffer, result);
		return result;
	}


private:

	Socket(int fd)
	{
		// for accepting a socket from fd
		union
		{
			sockaddr			sa;
			sockaddr_storage	storage;
		} addr;
		socklen_t length = sizeof(addr);

		mFd = ::accept(fd, &addr.sa, &length);
		mAddress = SocketAddress(&addr.sa, length);
	}


	int				mFd;
	bool			mIsConnected;
	SocketAddress	mAddress;
};


class ResolverTask : public ThreadedTask
{
public:

	ResolverTask(const std::string& service, const std::string& name,
			int type = SOCK_STREAM, int family = AF_UNSPEC) :
		mIsDone(false)
	{
		mFunction = boost::bind(&ResolverTask::resolve,
				this, service, name, type, family);
	}


	bool isDone() const
	{
		return mIsDone;
	}

	void run()
	{
		if (!mThread) mThread = Mf::detachFunction(mFunction);
	}


	const std::vector<SocketAddress>& addresses() const
	{
		return mAddressList;
	}


private:

	int resolve(const std::string& service, const std::string& name,
			int type, int family)
	{
		int status = SocketAddress::resolve(service, name,
				type, family, mAddressList);
		mIsDone = true;
		return status;
	}


	std::vector<SocketAddress>	mAddressList;
	bool						mIsDone;
	Function					mFunction;
};


} // namespace Mf

#endif // _MOOF_SOCKET_HH_