#include "stdafx.h"
#include <WBANetwork/common.h>
#include <WBANetwork/WBANetwork.h>
#include <WBANetwork/Net/Socket.h>
#include <WBANetwork/util/AsyncResult.h>
#include <mstcpip.h>

using namespace WBANetwork;

Socket::Socket()
:	m_Socket( INVALID_SOCKET ),
	m_ProtocolType( Protocol_None )
{
	memset(m_szIpAddress, 0, sizeof(m_szIpAddress));
}

Socket::~Socket()
{
	Close();
}

DWORD Socket::SetKeepAlive(u_long onoff, u_long KeepaliveTime, u_long Keepaliveinterval)
{
	DWORD dwRet = 0;

	if(INVALID_SOCKET != m_Socket)
	{		
		tcp_keepalive kl;
		kl.onoff = onoff;
		kl.keepalivetime = KeepaliveTime;
		kl.keepaliveinterval = Keepaliveinterval;
		::WSAIoctl(m_Socket, SIO_KEEPALIVE_VALS, &kl , sizeof(tcp_keepalive), 0, 0, &dwRet, NULL, NULL); 
	}

	return dwRet;
}

bool Socket::Create( ProtocolType type, bool overlapped )
{
	if( INVALID_SOCKET != m_Socket || 
		Protocol_None == type )
	{
		//_ASSERTE(!"Socket::Create");
		return false;
	}

	m_ProtocolType = type;
	switch( type )
	{
	case  Protocol_TCP:
		{
			if( true == overlapped )
			{
				m_Socket  = ::WSASocket( AF_INET, 
										SOCK_STREAM, 
										IPPROTO_TCP, 
										0, 
										0, 
										WSA_FLAG_OVERLAPPED );
			}
			else
			{
				m_Socket = ::socket( AF_INET, SOCK_STREAM, IPPROTO_TCP );
			}			

			if( INVALID_SOCKET == m_Socket )
			{
				//_ASSERTE(!"Socket::Create");
				return false;
			}

			::memset( &m_SockAddr, 0, sizeof( m_SockAddr ) );
		}
		break;

	case  Protocol_UDP:
		{
			if( true == overlapped )
			{
				m_Socket  = ::WSASocket( AF_INET, 
										SOCK_DGRAM,
										IPPROTO_UDP, 
										0,
										0,
										WSA_FLAG_OVERLAPPED );
			}
			else
			{
				m_Socket = ::socket( AF_INET, SOCK_DGRAM, IPPROTO_UDP );
			}

			if( INVALID_SOCKET == m_Socket )
			{
				//_ASSERTE(!"Socket::Create");
				return false;
			}

			::memset( &m_SockAddr, 0, sizeof( m_SockAddr ) );
		}
		break;
	}

	return true;
}

bool Socket::IsClose()
{
	if(INVALID_SOCKET == m_Socket) return true;
	return false;
}

void Socket::Close()
{
	if( INVALID_SOCKET == m_Socket )
	{
		return;
	}

	::closesocket( m_Socket );

	m_Socket		= INVALID_SOCKET;
	m_ProtocolType	= Protocol_None;
}

bool Socket::Attach( ProtocolType type, SOCKET handle, SocketAddr* addr )
{
	if( Protocol_None != type )
	{
		m_Socket		= handle;
		m_ProtocolType	= type;
	}

	if( 0 != addr )
	{
		if ( FAILED(StringCbPrintf(m_szIpAddress, sizeof(m_szIpAddress),
									_T("%d.%d.%d.%d"),
									addr->sin_addr.S_un.S_un_b.s_b1,
									addr->sin_addr.S_un.S_un_b.s_b2,
									addr->sin_addr.S_un.S_un_b.s_b3,
									addr->sin_addr.S_un.S_un_b.s_b4 ) ))
		{
			//_ASSERTE(!"Socket::Attach StringCbPrintf");
		}

		::memcpy( &m_SockAddr, addr, sizeof( SocketAddr ) );
	}

	return true;
}

SOCKET Socket::Accept( DWORD ms, SocketAddr* addr )
{
	if( NULL == addr )
	{
		//_ASSERTE(!"Socket::Accept");
		return INVALID_SOCKET;
	}

	struct timeval	timeout;
	fd_set			fds;
	int				len = sizeof( SocketAddr );

	FD_ZERO( &fds );
	FD_SET( m_Socket, &fds );

	timeout.tv_sec  = ms / 1000;
	timeout.tv_usec = ms % 1000;

	if( ::select( 0, &fds, 0, 0, &timeout ) == SOCKET_ERROR )
		return INVALID_SOCKET;

	if( FD_ISSET( m_Socket, &fds ) )
	{
		SOCKET		acceptedSocket = INVALID_SOCKET;

		try
		{
			acceptedSocket = ::accept( m_Socket, ( struct sockaddr* )addr, &len );

			if( 0 != addr )
			{
				if ( FAILED(StringCbPrintf(m_szIpAddress, sizeof(m_szIpAddress),
											_T("%d.%d.%d.%d"),
											addr->sin_addr.S_un.S_un_b.s_b1, 
											addr->sin_addr.S_un.S_un_b.s_b2,
											addr->sin_addr.S_un.S_un_b.s_b3, 
											addr->sin_addr.S_un.S_un_b.s_b4 ) ))
				{
					//_ASSERTE(!"Socket::Accept StringCbPrintf");
				}
			}
			else
			{
				throw ;
			}
		}
		catch( ... )
		{
			return INVALID_SOCKET;
		}

		return acceptedSocket;
	}

	return INVALID_SOCKET;
}

bool Socket::Connect( char* ipAddress, unsigned short portNo )
{
	if( NULL == ipAddress ||
		INVALID_SOCKET == m_Socket )
	{
		//_ASSERTE(!"Socket::Connect");
		return false;
	}

	fd_set	fds;

	FD_ZERO( &fds );
	FD_SET( m_Socket, &fds );

	::memset( &m_SockAddr, 0, sizeof( SocketAddr ) );
	m_SockAddr.sin_family			= AF_INET;
	m_SockAddr.sin_addr.s_addr		= ::inet_addr( ipAddress );
	m_SockAddr.sin_port				= ::htons( portNo );


	int ret	= ::connect( m_Socket, 
						 ( struct sockaddr* )&m_SockAddr, 
						 sizeof( m_SockAddr ) );

	if( SOCKET_ERROR == ret )
	{
		return false;
	}

	if ( FAILED(StringCbPrintf(m_szIpAddress, sizeof(m_szIpAddress),
								_T("%d.%d.%d.%d"),
								m_SockAddr.sin_addr.S_un.S_un_b.s_b1,
								m_SockAddr.sin_addr.S_un.S_un_b.s_b2,
								m_SockAddr.sin_addr.S_un.S_un_b.s_b3,
								m_SockAddr.sin_addr.S_un.S_un_b.s_b4 ) ))
	{
		//_ASSERTE(!"Socket::Connect StringCbPrintf");
	}

	return true;
}

bool Socket::Bind( char* ipAddress, unsigned short portNo )
{
	if( NULL == ipAddress ||
		INVALID_SOCKET == m_Socket )
	{
		//_ASSERTE(!"Socket::Bind");
		return false;
	}

	::memset( &m_SockAddr, 0, sizeof( m_SockAddr ) );
	m_SockAddr.sin_family			= AF_INET;

	if( 0 == strcmp( ipAddress, "127.0.0.1" ) )
	{
		m_SockAddr.sin_addr.s_addr	= htonl( INADDR_ANY );
	}
	else
	{
		m_SockAddr.sin_addr.s_addr	= ::inet_addr( ipAddress );
	}

	m_SockAddr.sin_port				= ::htons( portNo );


	int ret = ::bind( m_Socket, 
					  ( const sockaddr* )&m_SockAddr,
					  sizeof( m_SockAddr ) );

	if ( FAILED(StringCbPrintf(m_szIpAddress, sizeof(m_szIpAddress),
								_T("%d.%d.%d.%d"),
								m_SockAddr.sin_addr.S_un.S_un_b.s_b1,
								m_SockAddr.sin_addr.S_un.S_un_b.s_b2,
								m_SockAddr.sin_addr.S_un.S_un_b.s_b3,
								m_SockAddr.sin_addr.S_un.S_un_b.s_b4 ) ))
	{
		//_ASSERTE(!"Socket::Connect StringCbPrintf");
	}

	return ( SOCKET_ERROR != ret );
}

bool Socket::Listen( int backLogCount )
{
	return ( ::listen( m_Socket, backLogCount ) != SOCKET_ERROR );
}

int Socket::Send( PBYTE buffer, int length, AsyncResult* result )
{
	WSABUF	wsaBuf;
	wsaBuf.buf = (char *)buffer;
	wsaBuf.len = length;

	int ret = ::WSASend( m_Socket, 
						 &wsaBuf, 
						 1,
						 &result->transBytes,
						 0,
						 ( LPOVERLAPPED )result,
						 0 );

	result->error = ERROR_SUCCESS;
	
	if( SOCKET_ERROR == ret )
	{
		result->error = ::WSAGetLastError();
		if( ERROR_IO_PENDING != result->error )
		{
			return 0;
		}
	}

	return result->transBytes;
}

int Socket::Recv( PBYTE buffer, int bufSize, AsyncResult* result )
{
	WSABUF	wsaBuf;
	wsaBuf.buf	= (char *)buffer;
	wsaBuf.len	= bufSize;

	DWORD	flag	= 0;
	int		ret		= ::WSARecv( m_Socket, 
								 &wsaBuf, 
								 1,
								 &result->transBytes, 
								 &flag, 
								 ( LPOVERLAPPED )result,
								 0 );

	result->error	= ERROR_SUCCESS;

	if( SOCKET_ERROR == ret )
	{
		result->error = ::WSAGetLastError();
		if( ERROR_IO_PENDING != result->error )
		{
			return 0;
		}
	}

	return result->transBytes;
}

HANDLE Socket::GetNativeHandle()
{
	HANDLE handle = (m_Socket == INVALID_SOCKET ) 
		? INVALID_HANDLE_VALUE 
		: ( HANDLE )m_Socket;
	
	return handle;
}

void Socket::GetBufferSize( int* sndBufSize, int* rcvBufSize )
{
	if( NULL == sndBufSize ||
		NULL == rcvBufSize )
	{
		//_ASSERTE(!"Socket::GetBufferSize");
		return;
	}

	int len = sizeof( int );

	::getsockopt( m_Socket, SOL_SOCKET, SO_SNDBUF, ( char* )sndBufSize, &len );
	::getsockopt( m_Socket, SOL_SOCKET, SO_RCVBUF, ( char* )rcvBufSize, &len );
}

void Socket::SetReuseAddr()
{
	BOOL reuse = TRUE;

	::setsockopt( m_Socket, 
				  SOL_SOCKET, 
				  SO_REUSEADDR, 
				  ( const char* )&reuse, 
				  sizeof( reuse ) );
}

void Socket::SetIoctlSocket()
{
	unsigned long	argp = 1;

	::ioctlsocket( m_Socket, FIONBIO, &argp );	
}

void Socket::SetLinger( bool flag )
{
	LINGER	opt =
	{
		flag, 0
	};

	::setsockopt( m_Socket, SOL_SOCKET, SO_LINGER, ( char* )&opt, sizeof( opt ) );
}

void Socket::SetKeepAlive( bool flag )
{
	int	argp = flag;

	::setsockopt( m_Socket, SOL_SOCKET, SO_KEEPALIVE, ( char* )&argp, sizeof( argp ) );
}

void Socket::SetBufferSize( int sndBufSize, int rcvBufSize )
{
	::setsockopt( m_Socket, SOL_SOCKET, SO_SNDBUF, ( char* )&sndBufSize, sizeof( int ) );
	::setsockopt( m_Socket, SOL_SOCKET, SO_RCVBUF, ( char* )&rcvBufSize, sizeof( int ) );
}