#include "arp.h"
#include "ether.h"
#include "memory.h"
#include "string.h"

void init_arp(struct ether_t *ether) {
	ether->arp_table_count = 0;
}

void arp_process_packet(struct ether_t *ether, struct arp_packet_t *packet) {
	if (htons(packet->oper) == 1) {
		// request
		if (htonl(packet->tpa) == ether->ipv4) {
			arp_update_ipv4(ether, packet->sha, htonl(packet->spa));

			struct arp_packet_t arp_packet;

			arp_packet.htype = htons(1); // ethernet
			arp_packet.ptype = htons(0x0800); // ipv4
			arp_packet.hlen = 6;
			arp_packet.plen = 4;
			arp_packet.oper = htons(2); // reply
			arp_packet.sha[0] = ether->mac[0];
			arp_packet.sha[1] = ether->mac[1];
			arp_packet.sha[2] = ether->mac[2];
			arp_packet.sha[3] = ether->mac[3];
			arp_packet.sha[4] = ether->mac[4];
			arp_packet.sha[5] = ether->mac[5];
			arp_packet.spa = htonl(ether->ipv4);
			arp_packet.tha[0] = packet->sha[0];
			arp_packet.tha[1] = packet->sha[1];
			arp_packet.tha[2] = packet->sha[2];
			arp_packet.tha[3] = packet->sha[3];
			arp_packet.tha[4] = packet->sha[4];
			arp_packet.tha[5] = packet->sha[5];
			arp_packet.tpa = packet->spa;
			struct ether_packet_t *ether_packet = (struct ether_packet_t *)malloc(sizeof(struct arp_packet_t) + sizeof(struct ether_packet_t));

			ether_packet->destmac[0] = packet->sha[0];
			ether_packet->destmac[1] = packet->sha[1];
			ether_packet->destmac[2] = packet->sha[2];
			ether_packet->destmac[3] = packet->sha[3];
			ether_packet->destmac[4] = packet->sha[4];
			ether_packet->destmac[5] = packet->sha[5];

			ether_packet->sourcemac[0] = ether->mac[0];
			ether_packet->sourcemac[1] = ether->mac[1];
			ether_packet->sourcemac[2] = ether->mac[2];
			ether_packet->sourcemac[3] = ether->mac[3];
			ether_packet->sourcemac[4] = ether->mac[4];
			ether_packet->sourcemac[5] = ether->mac[5];

			ether_packet->ethertype = htons(0x806);

			memcpy(ether_packet->payload, &arp_packet, sizeof(struct arp_packet_t));

			if (ether_send(ether, (char *)ether_packet, sizeof(struct arp_packet_t) + sizeof(struct ether_packet_t)) == -1) {
				kprintf("Error sending reply.\n");
			}
			
			free(ether_packet);
		} 

	} else if (htons(packet->oper) == 2) {
		// reply
		arp_update_ipv4(ether, packet->sha, htonl(packet->spa));
	}
}

void arp_update_ipv4(struct ether_t *ether, unsigned char *mac, unsigned int ipv4) {
	int i;

	for (i=0;i<ether->arp_table_count;i++) {
		if (ether->arp_table[i]->mac[0] == mac[0] && ether->arp_table[i]->mac[1] == mac[1] && ether->arp_table[i]->mac[2] == mac[2] && ether->arp_table[i]->mac[3] == mac[3] && ether->arp_table[i]->mac[4] == mac[4] && ether->arp_table[i]->mac[5] == mac[5]) {
			ether->arp_table[i]->ipv4 = ipv4;
			return;
		}
	}

	if (ether->arp_table_count == 0) {
		ether->arp_table = (struct arp_table_t **)malloc(sizeof(struct arp_table_t *));
	} else {
		ether->arp_table = (struct arp_table_t **)realloc(ether->arp_table, sizeof(struct arp_table_t *) * (ether->arp_table_count + 1));
	}

	ether->arp_table[ether->arp_table_count] = (struct arp_table_t *)malloc(sizeof(struct arp_table_t));

	ether->arp_table[ether->arp_table_count]->mac[0] = mac[0];
	ether->arp_table[ether->arp_table_count]->mac[1] = mac[1];
	ether->arp_table[ether->arp_table_count]->mac[2] = mac[2];
	ether->arp_table[ether->arp_table_count]->mac[3] = mac[3];
	ether->arp_table[ether->arp_table_count]->mac[4] = mac[4];
	ether->arp_table[ether->arp_table_count]->mac[5] = mac[5];
	ether->arp_table[ether->arp_table_count]->ipv4 = ipv4;

	ether->arp_table_count++;
}

char *arp_req_ipv4(struct ether_t *ether, unsigned int ipv4) {
	int i;

	for (i=0;i<ether->arp_table_count;i++) {
		
		if (ether->arp_table[i]->ipv4 == ipv4) {
			return ether->arp_table[i]->mac;
		}
	}

	// address not in table...

	struct arp_packet_t arp_packet;

	arp_packet.htype = htons(1); // ethernet
	arp_packet.ptype = htons(0x0800); // ipv4
	arp_packet.hlen = 6;
	arp_packet.plen = 4;
	arp_packet.oper = htons(1); // request
	arp_packet.sha[0] = ether->mac[0];
	arp_packet.sha[1] = ether->mac[1];
	arp_packet.sha[2] = ether->mac[2];
	arp_packet.sha[3] = ether->mac[3];
	arp_packet.sha[4] = ether->mac[4];
	arp_packet.sha[5] = ether->mac[5];
	arp_packet.spa = htonl(ether->ipv4);
	arp_packet.tha[0] = 0xff;
	arp_packet.tha[1] = 0xff;
	arp_packet.tha[2] = 0xff;
	arp_packet.tha[3] = 0xff;
	arp_packet.tha[4] = 0xff;
	arp_packet.tha[5] = 0xff;
	arp_packet.tpa = htonl(ipv4);

	struct ether_packet_t *ether_packet = (struct ether_packet_t *)malloc(sizeof(struct arp_packet_t) + sizeof(struct ether_packet_t));

	ether_packet->destmac[0] = 0xff;
	ether_packet->destmac[1] = 0xff;
	ether_packet->destmac[2] = 0xff;
	ether_packet->destmac[3] = 0xff;
	ether_packet->destmac[4] = 0xff;
	ether_packet->destmac[5] = 0xff;

	ether_packet->sourcemac[0] = ether->mac[0];
	ether_packet->sourcemac[1] = ether->mac[1];
	ether_packet->sourcemac[2] = ether->mac[2];
	ether_packet->sourcemac[3] = ether->mac[3];
	ether_packet->sourcemac[4] = ether->mac[4];
	ether_packet->sourcemac[5] = ether->mac[5];

	ether_packet->ethertype = htons(0x806);


	memcpy(ether_packet->payload, &arp_packet, sizeof(struct arp_packet_t));

	ether_send(ether, (char *)ether_packet, sizeof(struct arp_packet_t) + sizeof(struct ether_packet_t));
	
	free(ether_packet);

	return (void *)0;
}