#pragma once

#include "Utilities.hpp"
#include <netinet/in.h>
#include <arpa/inet.h>
#include <string>
#include <cstring>
#include <variant> // Use std::variant

namespace reactor
{

class InetAddress
{
private:
	std::variant<sockaddr_in, sockaddr_in6> addr_;

public:
	/*
	 * Constructs an address.
	 */
	explicit InetAddress(uint16_t port = 0, bool ipv6 = false, bool loopback = false)
	{
		if (ipv6)
		{
			sockaddr_in6 addr6;
			memset(&addr6, 0, sizeof(addr6));
			addr6.sin6_family = AF_INET6;
			addr6.sin6_addr = loopback ? in6addr_loopback : in6addr_any;
			addr6.sin6_port = htons(port);
			addr_ = addr6;
		}
		else
		{
			sockaddr_in addr4;
			memset(&addr4, 0, sizeof(addr4));
			addr4.sin_family = AF_INET;
			addr4.sin_addr.s_addr = htonl(loopback ? INADDR_LOOPBACK : INADDR_ANY);
			addr4.sin_port = htons(port);
			addr_ = addr4;
		}
		LOG_TRACE << "InetAddress created: " << toIpPort();
	}

	/*
	 * Constructs an address from an IP and port.
	 */
	InetAddress(const std::string &ip, uint16_t port)
	{
		if (ip.find(':') != std::string::npos)
		{
			sockaddr_in6 addr6;
			memset(&addr6, 0, sizeof(addr6));
			addr6.sin6_family = AF_INET6;
			addr6.sin6_port = htons(port);
			if (inet_pton(AF_INET6, ip.c_str(), &addr6.sin6_addr) <= 0)
			{
				LOG_ERROR << "Invalid IPv6 address: " << ip;
			}
			addr_ = addr6;
		}
		else
		{
			sockaddr_in addr4;
			memset(&addr4, 0, sizeof(addr4));
			addr4.sin_family = AF_INET;
			addr4.sin_port = htons(port);
			if (inet_pton(AF_INET, ip.c_str(), &addr4.sin_addr) <= 0)
			{
				LOG_ERROR << "Invalid IPv4 address: " << ip;
			}
			addr_ = addr4;
		}
		LOG_TRACE << "InetAddress created from ip:port: " << toIpPort();
	}

	/*
	 * Constructs an address from a sockaddr_in struct.
	 */
	explicit InetAddress(const sockaddr_in &addr) : addr_(addr)
	{
		LOG_TRACE << "InetAddress created from sockaddr_in: " << toIpPort();
	}

	/*
	 * Constructs an address from a sockaddr_in6 struct.
	 */
	explicit InetAddress(const sockaddr_in6 &addr) : addr_(addr)
	{
		LOG_TRACE << "InetAddress created from sockaddr_in6: " << toIpPort();
	}

	const sockaddr *getSockAddr() const
	{
		// std::visit gets the pointer from the active variant member
		return std::visit(
			[](const auto &addr)
			{
				return reinterpret_cast<const sockaddr *>(&addr);
			},
			addr_);
	}

	socklen_t getSockLen() const
	{
		return std::visit([](const auto &addr)
		{ return sizeof(addr); }, addr_);
	}
	bool isIpV6() const { return std::holds_alternative<sockaddr_in6>(addr_); }
	uint16_t port() const
	{
		// Use if constexpr to access members with different names
		return std::visit(
			[](const auto &addr)
			{
				if constexpr (std::is_same_v<std::decay_t<decltype(addr)>, sockaddr_in>)
				{
					return ntohs(addr.sin_port);
				}
				else
				{
					return ntohs(addr.sin6_port);
				}
			},
			addr_);
	}

	std::string toIp() const
	{
		char buf[INET6_ADDRSTRLEN];
		std::visit(
			[&buf](const auto &addr)
			{
				if constexpr (std::is_same_v<std::decay_t<decltype(addr)>, sockaddr_in>)
				{
					inet_ntop(AF_INET, &addr.sin_addr, buf, sizeof(buf));
				}
				else
				{
					inet_ntop(AF_INET6, &addr.sin6_addr, buf, sizeof(buf));
				}
			},
			addr_);
		return std::string(buf);
	}

	std::string toIpPort() const
	{
		return isIpV6() ? "[" + toIp() + "]:" + std::to_string(port())
						: toIp() + ":" + std::to_string(port());
	}

	bool operator==(const InetAddress &other) const
	{
		if (addr_.index() != other.addr_.index()) return false;
		return std::visit(
			[&other](const auto &self_addr)
			{
				const auto &other_addr = std::get<std::decay_t<decltype(self_addr)>>(other.addr_);
				return memcmp(&self_addr, &other_addr, sizeof(self_addr)) == 0;
			},
			addr_);
	}

	bool operator!=(const InetAddress &other) const
	{
		return !(*this == other);
	}

	bool operator<(const InetAddress &other) const
	{
		if (isIpV6() != other.isIpV6())
		{
			return !isIpV6();
		}

		return std::visit(
			[&other](const auto &self_addr)
			{
				const auto &other_addr = std::get<std::decay_t<decltype(self_addr)>>(other.addr_);
				return memcmp(&self_addr, &other_addr, sizeof(self_addr)) < 0;
			},
			addr_);
	}

	std::string familyToString() const
	{
		return isIpV6() ? "IPv6" : "IPv4";
	}

	static bool resolve(const std::string &hostname, InetAddress &result)
	{
		// Simple resolution - in a real implementation you'd use getaddrinfo
		if (hostname == "localhost")
		{
			result = InetAddress(0, false, true);
			return true;
		}

		// Try to parse as IP address directly
		InetAddress addr(hostname, 0);
		if (addr.toIp() != "0.0.0.0" && addr.toIp() != "::")
		{
			result = addr;
			return true;
		}

		LOG_WARN << "Could not resolve hostname: " << hostname;
		return false;
	}

	// Make variant accessible for hashing
	const std::variant<sockaddr_in, sockaddr_in6> &getVariant() const
	{
		return addr_;
	}
};

} // namespace reactor

namespace std
{
template <>
struct hash<reactor::InetAddress>
{
	size_t operator()(const reactor::InetAddress &addr) const
	{
		// Hash the raw bytes of the underlying struct for better performance
		// than converting to a string first.
		return std::visit(
			[](const auto &a)
			{
				size_t seed = 0;
				std::string_view bytes(reinterpret_cast<const char *>(&a), sizeof(a));
				reactor::hashCombine(seed, std::hash<std::string_view>{}(bytes));
				return seed;
			},
			addr.getVariant());
	}
};
} // namespace std