#include "ehs/io/socket/TCP_W32.h"
#include "ehs/io/socket/DNS.h"

#include <WinSock2.h>
#include <WS2tcpip.h>

namespace ehs
{
	TCP::~TCP()
	{
		if (hdl == EHS_INVALID_SOCKET)
			return;

		Int_32 code;

		if (connection)
		{
			code = shutdown(hdl, SD_SEND);
			if (code == SOCKET_ERROR)
				EHS_LOG_INT(LogType::ERR, 0, "Failed to shutdown socket with error #" + Str_8::FromNum(GetLastError()) + ".");
		}

		code = closesocket(hdl);
		if (code == SOCKET_ERROR)
			EHS_LOG_INT(LogType::ERR, 1, "Failed to close socket with error #" + Str_8::FromNum(GetLastError()) + ".");

		if (!connection && WSACleanup() == SOCKET_ERROR)
			EHS_LOG_INT(LogType::ERR, 2, "Failed to shutdown WSA with error #" + Str_8::FromNum(WSAGetLastError()) + ".");
	}

	TCP::TCP()
		: hdl(EHS_INVALID_SOCKET)
	{
	}

	TCP::TCP(const IP IP)
		: BaseTCP(IP), hdl(EHS_INVALID_SOCKET)
	{
		TCP::Initialize();
	}

	TCP::TCP(TCP&& tcp) noexcept
		: BaseTCP(std::move(tcp)), hdl(tcp.hdl)
	{
		tcp.hdl = EHS_INVALID_SOCKET;
	}

	TCP::TCP(const TCP& tcp)
		: BaseTCP(tcp), hdl(EHS_INVALID_SOCKET)
	{
		TCP::Initialize();
	}

	TCP& TCP::operator=(TCP&& tcp) noexcept
	{
		if (this == &tcp)
			return *this;

		BaseTCP::operator=(tcp);

		hdl = tcp.hdl;

		tcp.hdl = EHS_INVALID_SOCKET;

		return *this;
	}

	TCP& TCP::operator=(const TCP& tcp)
	{
		if (this == &tcp)
			return *this;

		BaseTCP::operator=(tcp);

		hdl = EHS_INVALID_SOCKET;

		TCP::Initialize();

		return *this;
	}

	void TCP::Initialize()
	{
		WSADATA data = {};

		int code = WSAStartup(MAKEWORD(2, 2), &data);
		if (code)
		{
			EHS_LOG_INT(LogType::ERR, 0, "WSAStartup failed with the error #" + Str_8::FromNum(code) + ".");
			return;
		}

		if (ip == IP::V6)
			hdl = socket(AF_INET6, SOCK_STREAM, IPPROTO_TCP);
		else if (ip == IP::V4)
			hdl = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
		else
			return;

		if (hdl == EHS_INVALID_SOCKET)
		{
			UInt_32 code = WSAGetLastError();

			EHS_LOG_INT(LogType::ERR, 1, "Failed to create socket with error #" + Str_8::FromNum(code) + ".");

			if (WSACleanup() == SOCKET_ERROR)
				EHS_LOG_INT(LogType::ERR, 2, "Failed to shutdown WSA with error #" + Str_8::FromNum(WSAGetLastError()) + ".");
		}
	}

	void TCP::Release()
	{
		if (!IsValid())
			return;

		Int_32 code;

		if (connection)
		{
			code = shutdown(hdl, SD_SEND);
			if (code == SOCKET_ERROR)
				EHS_LOG_INT(LogType::ERR, 0, "Failed to shutdown socket with error #" + Str_8::FromNum(GetLastError()) + ".");
		}

		code = closesocket(hdl);
		if (code == SOCKET_ERROR)
			EHS_LOG_INT(LogType::ERR, 1, "Failed to close socket with error #" + Str_8::FromNum(GetLastError()) + ".");

		hdl = EHS_INVALID_SOCKET;

		if (!connection && WSACleanup() == SOCKET_ERROR)
			EHS_LOG_INT(LogType::ERR, 2, "Failed to shutdown WSA with error #" + Str_8::FromNum(WSAGetLastError()) + ".");

		connection = false;
		bound = false;
		listening = false;
		connected = false;
	}

	void TCP::Bind(const Str_8& address, UInt_16 port)
	{
		if (!IsValid() || bound || connection)
			return;

		if (ip == IP::V6)
			Bind_v6(address, port);
		else if (ip == IP::V4)
			Bind_v4(address, port);

		this->localAddr = address;
		this->localPort = port;

		bound = true;
	}

	void TCP::Listen()
	{
		if (connection || !IsValid() || !bound || listening)
			return;

		int code = listen(hdl, SOMAXCONN);
		if (code == -1)
		{
			EHS_LOG_INT(LogType::ERR, 0, "Failed to listen with error #" + Str_8::FromNum(WSAGetLastError()) + ".");

			return;
		}

		listening = true;
	}

	TCP* TCP::Accept()
	{
		if (connection || !IsValid() || !bound || !listening)
			return nullptr;

		sockaddr_in6 remote = {};
		UInt_32 addrLen = sizeof(sockaddr_in6);

		TCP* client = new TCP();
		client->ip = ip;
		client->localAddr = localAddr;
		client->localPort = localPort;
		client->connection = true;
		client->hdl = accept(hdl, (sockaddr*)&remote, (int*)&addrLen);

		if (client->hdl == EHS_INVALID_SOCKET)
		{
			Int_32 code = WSAGetLastError();

			if (code != WSAEWOULDBLOCK)
				EHS_LOG_INT(LogType::ERR, 0, "Failed to accept client with error #" + Str_8::FromNum(code) + ".");

			delete client;

			return nullptr;
		}

		if (addrLen == sizeof(sockaddr_in6))
		{
			char tmpAddr[INET6_ADDRSTRLEN];

			if (!inet_ntop(remote.sin6_family, &remote.sin6_addr, tmpAddr, INET6_ADDRSTRLEN))
			{
				Int_32 code = WSAGetLastError();

				EHS_LOG_INT(LogType::ERR, 1, "Failed to convert IPv6 address with error #" + Str_8::FromNum(code) + ".");

				delete client;

				return nullptr;
			}

			client->remoteAddr = tmpAddr;
			client->remotePort = ntohs(remote.sin6_port);
		}
		else if (addrLen == sizeof(sockaddr_in))
		{
			char tmpAddr[INET_ADDRSTRLEN];

			if (!inet_ntop(((sockaddr_in*)&remote)->sin_family, &((sockaddr_in*)&remote)->sin_addr, tmpAddr, INET_ADDRSTRLEN))
			{
				Int_32 code = WSAGetLastError();

				EHS_LOG_INT(LogType::ERR, 1, "Failed to convert IPv4 address with error #" + Str_8::FromNum(code) + ".");

				delete client;

				return nullptr;
			}

			client->remoteAddr = tmpAddr;
			client->remotePort = ntohs(((sockaddr_in*)&remote)->sin_port);
		}

		return client;
	}

	void TCP::Connect(const Str_8& address, const UInt_16 port)
	{
		if (connection || !IsValid() || listening)
			return;

		remoteHostName = address;
		remotePort = port;

		if (ip == IP::V6)
		{
			if (IsIPv6Only())
				remoteAddr = DNS::Resolve(IP::V6, address);
			else
				remoteAddr = DNS::Resolve(address);

			Connect_v6(remoteAddr, port);
		}
		else if (ip == IP::V4)
		{
			remoteAddr = DNS::Resolve(IP::V4, address);

			Connect_v4(remoteAddr, port);
		}

		connected = true;
	}

	UInt_64 TCP::Send(const Byte *const buffer, const UInt_32 size)
	{
		if (!IsValid())
		{
			EHS_LOG_INT(LogType::ERR, 0, "Attempted to send while socket is not initialized.");
			return 0;
		}

		if ((!connection && !connected))
		{
			EHS_LOG_INT(LogType::ERR, 1, "Attempted to send while socket is not connected or a connection.");
			return 0;
		}

		SInt_64 sent = (SInt_64)send(hdl, (char*)buffer, (int)size, 0);
		if (sent == SOCKET_ERROR)
		{
			int err = WSAGetLastError();
			if (err == WSAECONNRESET)
			{
				Release();
				EHS_LOG_INT(LogType::INFO, 0, "Connection dropped.");
			}
			else
				EHS_LOG_INT(LogType::ERR, 1, "Failed to send with error #" + Str_8::FromNum(err) + ".");

			return 0;
		}

		return (UInt_64)sent;
	}

	UInt_64 TCP::Receive(Byte* const buffer, const UInt_32 size)
	{
		if (!IsValid())
		{
			EHS_LOG_INT(LogType::ERR, 0, "Attempted to receive while socket is not initialized.");
			return 0;
		}

		if ((!connection && !connected))
		{
			EHS_LOG_INT(LogType::ERR, 1, "Attempted to receive while socket is not connected or a connection.");
			return 0;
		}

		SInt_64 received = (SInt_64)recv(hdl, (char*)buffer, (int)size, 0);
		if (received == SOCKET_ERROR)
		{
			int err = WSAGetLastError();
			if (err == WSAECONNRESET)
			{
				Release();
				EHS_LOG_INT(LogType::INFO, 0, "Connection dropped.");
			}
			else if (err == WSAECONNABORTED)
			{
				EHS_LOG_INT(LogType::INFO, 1, "Receiving timed-out.");
			}
			else if (err != WSAEWOULDBLOCK)
			{
				EHS_LOG_INT(LogType::ERR, 2, "Failed to receive with error #" + Str_8::FromNum(err) + ".");
			}

			return 0;
		}

		return (UInt_64)received;
	}

	void TCP::SetBlocking(const bool blocking)
	{
		if (!IsValid())
		{
			EHS_LOG_INT(LogType::ERR, 0, "Attempted to toggle blocking while socket is not initialized.");
			return;
		}

		u_long r = (u_long)!blocking;

		int result = ioctlsocket(hdl, FIONBIO, &r);
		if (result != NO_ERROR)
			EHS_LOG_INT(LogType::ERR, 1, "Failed to toggle non-blocking mode with error #" + Str_8::FromNum(result) + ".");
	}

	bool TCP::IsBlocking() const
	{
		u_long r = 0;
		if (ioctlsocket(hdl, FIONREAD, &r) == SOCKET_ERROR)
			EHS_LOG_INT(LogType::ERR, 0, "Failed to retrieve socket info.");

		return (bool)r;
	}

	void TCP::SetIPv6Only(const bool value)
	{
		if (ip != IP::V6)
		{
			EHS_LOG_INT(LogType::WARN, 0, "Cannot set IPv6 only mode while socket is not using IPv6.");
			return;
		}

		if (!IsValid())
		{
			EHS_LOG_INT(LogType::WARN, 1, "Attempted to set IPv6 only mode while socket is not initialized.");
			return;
		}

		const int result = (int)value;
		if (setsockopt(hdl, IPPROTO_IPV6, IPV6_V6ONLY, (const char *)&result, sizeof(int)) == -1)
		{
			EHS_LOG_INT(LogType::ERR, 2, "Failed to set IPv6 only mode with error #" + Str_8::FromNum(errno) + ".");
			return;
		}

		EHS_LOG_SUCCESS();
	}

	bool TCP::IsIPv6Only() const
	{
		if (ip != IP::V6)
			return false;

		if (!IsValid())
		{
			EHS_LOG_INT(LogType::WARN, 1, "Attempted to set IPv6 only mode while socket is not initialized.");
			return false;
		}

		int result;
		socklen_t len = sizeof(int);

		if (getsockopt(hdl, IPPROTO_IPV6, IPV6_V6ONLY, (char *)&result, &len) == -1)
		{
			EHS_LOG_INT(LogType::ERR, 2, "Failed to set IPv6 only mode with error #" + Str_8::FromNum(errno) + ".");
			return false;
		}

		EHS_LOG_SUCCESS();

		return result;
	}

	bool TCP::IsValid() const
	{
		return hdl != EHS_INVALID_SOCKET;
	}

	void TCP::Bind_v6(const Str_8& address, UInt_16 port)
	{
		sockaddr_in6 result = {};
		result.sin6_family = AF_INET6;
		result.sin6_port = htons(port);

		if (address.Size())
		{
			int code = inet_pton(AF_INET6, address, &result.sin6_addr);
			if (!code)
			{
				EHS_LOG_INT(LogType::ERR, 0, "The given address, \"" + address + "\" is not valid.");
				return;
			}
			else if (code == -1)
			{
				Int_32 dCode = WSAGetLastError();

				EHS_LOG_INT(LogType::ERR, 1, "Failed to convert address with error #" + Str_8::FromNum(dCode) + ".");
				return;
			}
		}
		else
		{
			result.sin6_addr = in6addr_any;
		}

		int code = bind(hdl, (sockaddr*)&result, sizeof(sockaddr_in6));
		if (code == -1)
		{
			EHS_LOG_INT(LogType::ERR, 2, "Failed to bind socket with error #" + Str_8::FromNum(WSAGetLastError()) + ".");

			return;
		}
	}

	void TCP::Bind_v4(const Str_8& address, UInt_16 port)
	{
		sockaddr_in result = {};
		result.sin_family = AF_INET;
		result.sin_port = htons(port);

		if (address.Size())
		{
			int code = inet_pton(AF_INET, address, &result.sin_addr);
			if (!code)
			{
				EHS_LOG_INT(LogType::ERR, 0, "The given address, \"" + address + "\" is not valid.");
				return;
			}
			else if (code == -1)
			{
				Int_32 dCode = WSAGetLastError();

				EHS_LOG_INT(LogType::ERR, 1, "Failed to convert address with error #" + Str_8::FromNum(dCode) + ".");
				return;
			}
		}
		else
		{
			result.sin_addr.S_un.S_addr = INADDR_ANY;
		}

		int code = bind(hdl, (sockaddr*)&result, sizeof(sockaddr_in));
		if (code == -1)
		{
			EHS_LOG_INT(LogType::ERR, 2, "Failed to bind socket with error #" + Str_8::FromNum(WSAGetLastError()) + ".");

			return;
		}
	}

	void TCP::Connect_v6(const Str_8& address, UInt_16 port)
	{
		sockaddr_in6 result = {};
		result.sin6_family = AF_INET6;
		result.sin6_port = htons(port);

		int code = inet_pton(AF_INET6, address, &result.sin6_addr);
		if (!code)
		{
			EHS_LOG_INT(LogType::ERR, 0, "The given address, \"" + address + "\" is not valid.");
			return;
		}
		else if (code == -1)
		{
			Int_32 dCode = WSAGetLastError();

			EHS_LOG_INT(LogType::ERR, 1, "Failed to convert address with error #" + Str_8::FromNum(dCode) + ".");
			return;
		}

		code = connect(hdl, (sockaddr*)&result, sizeof(sockaddr_in6));
		if (code == SOCKET_ERROR)
		{
			int err = WSAGetLastError();
			if (err == WSAETIMEDOUT)
			{
				EHS_LOG_INT(LogType::INFO, 2, "Connection attempt timed-out.");
			}
			else
			{
				EHS_LOG_INT(LogType::ERR, 3, "Failed to connect with error #" + Str_8::FromNum(err) + ".");
			}

			return;
		}
	}

	void TCP::Connect_v4(const Str_8& address, UInt_16 port)
	{
		sockaddr_in result = {};
		result.sin_family = AF_INET;
		result.sin_port = htons(port);

		int code = inet_pton(AF_INET, address, &result.sin_addr);
		if (!code)
		{
			EHS_LOG_INT(LogType::ERR, 0, "The given address, \"" + address + "\" is not valid.");
			return;
		}
		else if (code == -1)
		{
			Int_32 dCode = WSAGetLastError();

			EHS_LOG_INT(LogType::ERR, 1, "Failed to convert address with error #" + Str_8::FromNum(dCode) + ".");
			return;
		}

		code = connect(hdl, (sockaddr*)&result, sizeof(sockaddr_in));
		if (code == SOCKET_ERROR)
		{
			int err = WSAGetLastError();
			if (err == WSAETIMEDOUT)
			{
				EHS_LOG_INT(LogType::INFO, 2, "Connection attempt timed-out.");
			}
			else
			{
				EHS_LOG_INT(LogType::ERR, 3, "Failed to connect with error #" + Str_8::FromNum(err) + ".");
			}

			return;
		}
	}
}