// SPDX-License-Identifier: BSD-3-Clause
// Copyright (c) 2025 Anthony Harivel

#include "control_output.h"
#include "dhcp.h"
#include "graph.h"
#include "ip4_datapath.h"
#include "l4.h"
#include "mbuf.h"
#include "rxtx.h"

#include <gr_macro.h>

#include <rte_byteorder.h>
#include <rte_mbuf.h>
#include <rte_udp.h>

struct dhcp_input_trace_data {
	uint32_t xid;
	uint8_t msg_type;
	uint8_t op;
};

enum {
	CONTROL = 0,
	NO_HEADROOM,
	EDGE_COUNT,
};

static uint16_t
dhcp_input_process(struct rte_graph *graph, struct rte_node *node, void **objs, uint16_t nb_objs) {
	const struct iface *iface;
	struct rte_mbuf *mbuf;
	rte_edge_t edge;

	for (uint16_t i = 0; i < nb_objs; i++) {
		mbuf = objs[i];
		iface = mbuf_data(mbuf)->iface;

		if (gr_mbuf_is_traced(mbuf)) {
			struct dhcp_input_trace_data *t;
			const struct rte_udp_hdr *udp;
			const struct dhcp_packet *dhcp;
			uint16_t pkt_len;

			t = gr_mbuf_trace_add(mbuf, node, sizeof(*t));
			pkt_len = rte_pktmbuf_data_len(mbuf);

			// Parse minimal DHCP header for trace
			if (pkt_len >= sizeof(*udp) + sizeof(*dhcp)) {
				udp = rte_pktmbuf_mtod(mbuf, const struct rte_udp_hdr *);
				dhcp = PAYLOAD(udp);

				t->xid = rte_be_to_cpu_32(dhcp->xid);
				t->op = dhcp->op;

				// Extract message type from options (option 53)
				t->msg_type = 0;
				const uint8_t *options = dhcp->options;
				uint16_t options_len = pkt_len - sizeof(*udp) - sizeof(*dhcp);
				uint16_t pos = 0;

				while (pos < options_len && pos < 64) { // Limit search
					uint8_t opt = options[pos++];
					if (opt == DHCP_OPT_END)
						break;
					if (opt == DHCP_OPT_PAD)
						continue;
					if (pos >= options_len)
						break;
					uint8_t len = options[pos++];
					if (opt == DHCP_OPT_MESSAGE_TYPE && len == 1
					    && pos < options_len) {
						t->msg_type = options[pos];
						break;
					}
					pos += len;
				}
			}
		}

		if (dhcp_enabled(iface->id)) {
			control_output_set_cb(mbuf, dhcp_input_cb, 0);
		} else {
			struct ip_local_mbuf_data *d = ip_local_mbuf_data(mbuf);

			// reconstruct the IP and Ethernet headers before sending to control plane
			// interface
			struct rte_ipv4_hdr *ip = gr_mbuf_prepend(mbuf, ip);
			if (ip == NULL) {
				edge = NO_HEADROOM;
				goto next;
			}
			ip_set_fields(ip, d);

			struct rte_ether_hdr *eth = gr_mbuf_prepend(mbuf, eth);
			if (eth == NULL) {
				edge = NO_HEADROOM;
				goto next;
			}
			memset(&eth->src_addr, 0, sizeof(eth->src_addr));
			iface_get_eth_addr(iface, &eth->dst_addr);
			eth->ether_type = RTE_BE16(RTE_ETHER_TYPE_IPV4);

			// No VLAN header is written to the TAP. Clear the residual
			// mbuf metadata so traces reflect the actual frame bytes.
			iface_mbuf_data(mbuf)->vlan_id = 0;

			control_output_set_cb(mbuf, iface_cp_tx, 0);
		}

		edge = CONTROL;
next:
		rte_node_enqueue_x1(graph, node, edge, mbuf);
	}

	return nb_objs;
}

static const char *dhcp_msg_type_str(uint8_t type) {
	switch (type) {
	case DHCP_DISCOVER:
		return "DISCOVER";
	case DHCP_OFFER:
		return "OFFER";
	case DHCP_REQUEST:
		return "REQUEST";
	case DHCP_DECLINE:
		return "DECLINE";
	case DHCP_ACK:
		return "ACK";
	case DHCP_NAK:
		return "NAK";
	case DHCP_RELEASE:
		return "RELEASE";
	case DHCP_INFORM:
		return "INFORM";
	default:
		return "UNKNOWN";
	}
}

static int dhcp_input_trace_format(char *buf, size_t len, const void *data, size_t /*data_len*/) {
	const struct dhcp_input_trace_data *t = data;
	return snprintf(
		buf,
		len,
		"%s xid=0x%08x %s",
		t->op == BOOTREQUEST	   ? "REQUEST" :
			t->op == BOOTREPLY ? "REPLY" :
					     "?",
		t->xid,
		dhcp_msg_type_str(t->msg_type)
	);
}

static void dhcp_input_register_port(void) {
	l4_input_register_port(IPPROTO_UDP, RTE_BE16(68), "dhcp_input");
}

static struct rte_node_register node = {
	.name = "dhcp_input",
	.process = dhcp_input_process,
	.nb_edges = EDGE_COUNT,
	.next_nodes = {
		[CONTROL] = "control_output",
		[NO_HEADROOM] = "error_no_headroom",
	},
};

static struct gr_node_info info = {
	.node = &node,
	.type = GR_NODE_T_L4 | GR_NODE_T_CONTROL,
	.trace_format = dhcp_input_trace_format,
	.register_callback = dhcp_input_register_port,
};

GR_NODE_REGISTER(info);
