#include <errno.h>
#include <kernel/types.h>
#include <kernel/string.h>
#include <kernel/printf.h>
#include <kernel/syscall.h>
#include <kernel/vfs.h>
#include <kernel/hashmap.h>
#include <kernel/mod/net.h>
#include <kernel/net/netif.h>
#include <kernel/net/eth.h>

#include <sys/socket.h>

struct arp_header {
	uint16_t arp_htype;
	uint16_t arp_ptype;
	uint8_t  arp_hlen;
	uint8_t  arp_plen;
	uint16_t arp_oper;
	union {
		struct {
			uint8_t  arp_sha[6];
			uint32_t arp_spa;
			uint8_t arp_tha[6];
			uint32_t arp_tpa;
		} __attribute__((packed)) arp_eth_ipv4;
	} arp_data;
} __attribute__((packed));

static void ip_ntoa(const uint32_t src_addr, char * out) {
	snprintf(out, 16, "%d.%d.%d.%d",
		(src_addr & 0xFF000000) >> 24,
		(src_addr & 0xFF0000) >> 16,
		(src_addr & 0xFF00) >> 8,
		(src_addr & 0xFF));
}

spin_lock_t net_arp_cache_lock = {0};
hashmap_t * net_arp_cache = NULL;

struct ArpCacheEntry {
	uint8_t hwaddr[6];
	uint16_t flags;
	struct EthernetDevice * iface;
};

void net_arp_cache_add(struct EthernetDevice * iface, uint32_t addr, uint8_t * hwaddr, uint16_t flags) {
	spin_lock(net_arp_cache_lock);
	if (!net_arp_cache) net_arp_cache = hashmap_create_int(10);
	struct ArpCacheEntry * entry = hashmap_get(net_arp_cache, (void*)(uintptr_t)addr);
	if (!entry) entry = malloc(sizeof(struct ArpCacheEntry));
	memcpy(entry->hwaddr, hwaddr, 6);
	entry->flags = flags;
	entry->iface = iface;
	hashmap_set(net_arp_cache, (void*)(uintptr_t)addr, entry);
	spin_unlock(net_arp_cache_lock);
}

struct ArpCacheEntry * net_arp_cache_get(uint32_t addr) {
	if (!net_arp_cache) return NULL;
	spin_lock(net_arp_cache_lock);
	struct ArpCacheEntry * out = hashmap_get(net_arp_cache, (void*)(uintptr_t)addr);
	spin_unlock(net_arp_cache_lock);
	return out;
}

void net_arp_handle(struct arp_header * packet, fs_node_t * nic) {
	printf("net: arp: hardware %d protocol %d operation %d hlen %d plen %d\n",
		ntohs(packet->arp_htype), ntohs(packet->arp_ptype), ntohs(packet->arp_oper),
		packet->arp_hlen, packet->arp_plen);
	struct EthernetDevice * eth_dev = nic->device;

	if (ntohs(packet->arp_htype) == 1 && ntohs(packet->arp_ptype) == ETHERNET_TYPE_IPV4) {
		/* Ethernet, IPv4 */
		if (packet->arp_data.arp_eth_ipv4.arp_spa) {
			net_arp_cache_add(eth_dev, packet->arp_data.arp_eth_ipv4.arp_spa, packet->arp_data.arp_eth_ipv4.arp_sha, 0);
		}
		if (ntohs(packet->arp_oper) == 1) {
			char spa[17];
			ip_ntoa(ntohl(packet->arp_data.arp_eth_ipv4.arp_spa), spa);
			char tpa[17];
			ip_ntoa(ntohl(packet->arp_data.arp_eth_ipv4.arp_tpa), tpa);
			printf("net: arp: " MAC_FORMAT " (%s) wants to know who %s is\n",
				FORMAT_MAC(packet->arp_data.arp_eth_ipv4.arp_sha),
				spa, tpa);
			if (eth_dev->ipv4_addr &&  packet->arp_data.arp_eth_ipv4.arp_tpa == eth_dev->ipv4_addr) {
				printf("net: arp: that's us, we should reply...\n");

				struct arp_header response = {0};
				response.arp_htype = htons(1);
				response.arp_ptype = htons(ETHERNET_TYPE_IPV4);
				response.arp_hlen = 6;
				response.arp_plen = 4;
				response.arp_oper = htons(2);
				memcpy(response.arp_data.arp_eth_ipv4.arp_sha, eth_dev->mac, 6);
				memcpy(response.arp_data.arp_eth_ipv4.arp_tha, packet->arp_data.arp_eth_ipv4.arp_sha, 6);
				response.arp_data.arp_eth_ipv4.arp_spa = eth_dev->ipv4_addr;
				response.arp_data.arp_eth_ipv4.arp_tpa = packet->arp_data.arp_eth_ipv4.arp_spa;
				net_eth_send(eth_dev, sizeof(struct arp_header), &response, ETHERNET_TYPE_ARP, packet->arp_data.arp_eth_ipv4.arp_sha);
			}
		} else if (ntohs(packet->arp_oper) == 2) {
			char spa[17];
			ip_ntoa(ntohl(packet->arp_data.arp_eth_ipv4.arp_spa), spa);
			printf("net: arp: " MAC_FORMAT " says they are %s\n",
				FORMAT_MAC(packet->arp_data.arp_eth_ipv4.arp_sha),
				spa);
		}
	}
}