#include "ehs/io/socket/ICMP_LNX.h"
#include "ehs/Serializer.h"
#include "ehs/Log.h"

#include <arpa/inet.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <ifaddrs.h>
#include <unistd.h>

namespace ehs
{
	ICMP::~ICMP()
	{
		if (close(hdl) == -1)
			EHS_LOG_INT(LogType::ERR, 0, "Failed to close socket.");
	}

	ICMP::ICMP()
		: hdl(EHS_INVALID_SOCKET), src{}
	{
	}

	ICMP::ICMP(const IP version)
		: BaseICMP(version), src{}
	{
		if (version == IP::V6)
			hdl = socket(AF_INET6, SOCK_RAW, IPPROTO_ICMPV6);
		else
			hdl = socket(AF_INET, SOCK_RAW, IPPROTO_ICMP);

		if (hdl < 0)
		{
			EHS_LOG_INT(LogType::ERR, 0, "Failed to create ICMP socket with error #" + Str_8::FromNum(errno) + ".");
			return;
		}

		EHS_LOG_SUCCESS();
	}

	ICMP::ICMP(ICMP &&icmp) noexcept
		: BaseICMP((BaseICMP &&)icmp), hdl(icmp.hdl), src(icmp.src)
	{
		icmp.hdl = EHS_INVALID_SOCKET;
		icmp.src = {};
	}

	ICMP::ICMP(const ICMP &icmp)
		: BaseICMP(icmp), hdl(icmp.hdl), src{}
	{
	}

	ICMP & ICMP::operator=(ICMP &&icmp) noexcept
	{
		if (this == &icmp)
			return *this;

		BaseICMP::operator=((BaseICMP &&)icmp);

		hdl = icmp.hdl;
		src = icmp.src;

		icmp.hdl = EHS_INVALID_SOCKET;
		icmp.src = {};

		return *this;
	}

	ICMP & ICMP::operator=(const ICMP &icmp)
	{
		if (this == &icmp)
			return *this;

		BaseICMP::operator=(icmp);

		hdl = icmp.hdl;
		src = {};

		return *this;
	}

	void ICMP::Release()
	{
		if (close(hdl) == -1)
		{
			EHS_LOG_INT(LogType::ERR, 0, "Failed to close socket.");

			return;
		}

		hdl = EHS_INVALID_SOCKET;

		EHS_LOG_SUCCESS();
	}

	void ICMP::SetReceiveTimeout(const UInt_64 timeout)
	{
		timeval result = {};
		result.tv_sec = (long)timeout;
		result.tv_usec = 0;

		if (setsockopt(hdl, SOL_SOCKET, SO_RCVTIMEO, &result, sizeof(result)) < 0)
		{
			EHS_LOG_INT(LogType::WARN, 0, "Failed to set receive timeout with error #" + Str_8::FromNum(errno) + ".");

			return;
		}

		EHS_LOG_SUCCESS();
	}

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

	bool ICMP::IsLinkLocal(const in6_addr &addr)
	{
		return addr.s6_addr[0] == 0xfe && (addr.s6_addr[1] & 0xc0) == 0x80;
	}

	sockaddr_in6 ICMP::RetrieveSrcAddress()
	{
		ifaddrs *ifaddr;
		sockaddr_in6 addr = {};

		if (getifaddrs(&ifaddr) == -1)
		{
			EHS_LOG_INT(LogType::ERR, 0, "Failed to retrieve public address with error #" + Str_8::FromNum(errno) + ".");

			return addr;
		}

		for (ifaddrs *ifa = ifaddr; ifa; ifa = ifa->ifa_next)
		{
			if (ifa->ifa_addr == nullptr || ifa->ifa_addr->sa_family != AF_INET6)
				continue;

			addr = *(sockaddr_in6 *)ifa->ifa_addr;
			if (!addr.sin6_addr.s6_addr32[0] || IsLinkLocal(addr.sin6_addr))
				continue;

			break;
		}

		freeifaddrs(ifaddr);

		EHS_LOG_SUCCESS();

		return addr;
	}

	UInt_32 ICMP::CalculatePseudoHeaderChecksum(const PseudoICMPv6_Header &header)
	{
		UInt_32 checksum = 0;

		for (UInt_8 i = 0; i < 16; ++i)
			checksum += header.src.sin6_addr.s6_addr[i];

		for (UInt_8 i = 0; i < 16; ++i)
			checksum += header.dst.sin6_addr.s6_addr[i];

		checksum += 58;

		checksum += htons(header.length);

		checksum = (checksum >> 16) + (checksum & 0xFFFF);
		checksum += (checksum >> 16);

		return checksum;
	}

	UInt_16 ICMP::ComputeChecksumV6(UInt_16 *buffer, Size length, const sockaddr_in6 &dst)
	{
		UInt_32 checksum = 0;

		if (!src.sin6_addr.s6_addr32[0])
		{
			src = RetrieveSrcAddress();
			if (!src.sin6_addr.s6_addr32[0])
			{
				EHS_LOG_INT(LogType::ERR, 0, "Could not retrieve a suitable global address.");
				return checksum;
			}
		}

		checksum += CalculatePseudoHeaderChecksum({src, dst, (UInt_32)length});

		while (length > 1)
		{
			checksum += *buffer++;
			length -= sizeof(UInt_16);
		}

		if (length == 1)
			checksum += *(UInt_8 *)buffer;

		// Carry over any overflow from the 16-bit result
		checksum = (checksum >> 16) + (checksum & 0xFFFF);
		checksum += (checksum >> 16);

		// Return the 16-bit complement
		return ~checksum;
	}

	UInt_64 ICMP::SendV6(const Str_8 &address, ICMP_Header header, const Byte *data, UInt_64 size)
	{
		if (!IsValid())
		{
			EHS_LOG_INT(LogType::WARN, 0, "Socket is not initialized.");

			return 0;
		}

		header.checksum = 0;

		Serializer<UInt_64> payload(Endianness::LE);
		payload.Write(header);
		payload.Resize(payload.Size() + size);

		Util::Copy(&payload[payload.GetOffset()], data, size);

		payload.SetOffset(payload.GetOffset() + size);

		sockaddr_in6 dst = {};
		dst.sin6_family = AF_INET6;
		inet_pton(AF_INET6, address, &(dst.sin6_addr));

		header.checksum = ComputeChecksumV6((UInt_16 *)&payload[0], payload.Size(), dst);

		payload.SetOffset(0);
		payload.Write(header);
		payload.SetOffset(payload.Size());

		SInt_64 sent = sendto(hdl, payload, payload.Size(), 0, (sockaddr *)&dst, sizeof(sockaddr_in6));
		if (sent < 0)
		{
			EHS_LOG_INT(LogType::ERR, 0, "Failed to send packet with error #" + Str_8::FromNum(errno) + ".");

			return 0;
		}

		EHS_LOG_SUCCESS();

		return sent;
	}

	UInt_64 ICMP::SendV4(const Str_8 &address, ICMP_Header header, const Byte *data, UInt_64 size)
	{
		if (!IsValid())
		{
			EHS_LOG_INT(LogType::WARN, 0, "Socket is not initialized.");

			return 0;
		}

		header.checksum = 0;

		Serializer<UInt_64> payload(Endianness::LE);
		payload.Write(header);
		payload.Resize(payload.Size() + size);

		Util::Copy(&payload[payload.GetOffset()], data, size);

		payload.SetOffset(payload.GetOffset() + size);

		header.checksum = ComputeChecksumV4((UInt_16 *)&payload[0], payload.Size());

		payload.SetOffset(0);
		payload.Write(header);
		payload.SetOffset(payload.Size());

		sockaddr_in dst = {};
		dst.sin_family = AF_INET;
		inet_pton(AF_INET, address, &(dst.sin_addr));

		SInt_64 sent = sendto(hdl, payload, payload.Size(), 0, (sockaddr *)&dst, sizeof(sockaddr_in));
		if (sent < 0)
		{
			EHS_LOG_INT(LogType::ERR, 0, "Failed to send packet with error #" + Str_8::FromNum(errno) + ".");

			return 0;
		}

		EHS_LOG_SUCCESS();

		return sent;
	}

	UInt_64 ICMP::ReceiveV6(Str_8 &address, ICMP_Header &header, Serializer<UInt_64> &data) const
	{
		if (!IsValid())
		{
			EHS_LOG_INT(LogType::WARN, 0, "Socket is not initialized.");

			return 0;
		}

		Serializer<UInt_64> payload(Endianness::LE);
		payload.Resize(1500);

		sockaddr_in6 remote = {};
		socklen_t from_len = sizeof(remote);

		SInt_64 recv = recvfrom(hdl, payload, 1500, 0, (sockaddr *)&remote, &from_len);
		if (recv < 0)
		{
			int code = errno;
			if (code == EAGAIN)
				EHS_LOG_SUCCESS();
			else
				EHS_LOG_INT(LogType::ERR, 0, "Failed to receive packet with error #" + Str_8::FromNum(code) + ".");

			return 0;
		}

		payload.Resize(recv);

		char tmpAddr[INET6_ADDRSTRLEN];

		if (!inet_ntop(remote.sin6_family, &remote.sin6_addr, tmpAddr, INET6_ADDRSTRLEN))
		{
			EHS_LOG_INT(LogType::ERR, 1, "Failed to convert IPv6 address with error #" + Str_8::FromNum(errno) + ".");

			return recv;
		}

		address = tmpAddr;
		header = payload.Read<ICMP_Header>();
		data = Serializer<UInt_64>(payload.GetEndianness(), &payload[payload.GetOffset()], payload.Size() - payload.GetOffset());

		EHS_LOG_SUCCESS();

		return recv;
	}

	UInt_64 ICMP::ReceiveV4(Str_8 &address, ICMP_Header &header, Serializer<UInt_64> &data) const
	{
		if (!IsValid())
		{
			EHS_LOG_INT(LogType::WARN, 0, "Socket is not initialized.");

			return 0;
		}

		Serializer<UInt_64> payload(Endianness::LE);
		payload.Resize(1500);

		sockaddr_in remote = {};
		socklen_t from_len = sizeof(remote);

		SInt_64 recv = recvfrom(hdl, payload, 1500, 0, (sockaddr *)&remote, &from_len);
		if (recv < 0)
		{
			int code = errno;
			if (code == EAGAIN)
				EHS_LOG_SUCCESS();
			else
				EHS_LOG_INT(LogType::ERR, 0, "Failed to receive packet with error #" + Str_8::FromNum(code) + ".");

			return 0;
		}

		payload.Resize(recv);

		char tmpAddr[INET_ADDRSTRLEN];

		if (!inet_ntop(remote.sin_family, &remote.sin_addr, tmpAddr, INET_ADDRSTRLEN))
		{
			EHS_LOG_INT(LogType::ERR, 1, "Failed to convert IPv4 address with error #" + Str_8::FromNum(errno) + ".");

			return recv;
		}

		address = tmpAddr;
		iphdr ipHeader = payload.Read<iphdr>();
		header = payload.Read<ICMP_Header>();
		data = Serializer<UInt_64>(payload.GetEndianness(), &payload[payload.GetOffset()], payload.Size() - payload.GetOffset());

		EHS_LOG_SUCCESS();

		return recv;
	}
}