/*
 * (C) 2021-2022 by sysmocom - s.f.m.c. GmbH <info@sysmocom.de>
 * All Rights Reserved.
 *
 * Author: Neels Janosch Hofmeyr <nhofmeyr@sysmocom.de>
 *
 * SPDX-License-Identifier: GPL-2.0+
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

#include <errno.h>
#include <netinet/in.h>

#include <osmocom/core/utils.h>
#include <osmocom/core/talloc.h>

#include <osmocom/pfcp/pfcp_endpoint.h>

#include <osmocom/upfloadgen/pfcp_tool.h>
#include <osmocom/upfloadgen/gtp_flood.h>

struct g_pfcp_tool *g_pfcp_tool = NULL;

struct osmo_tdef_group g_pfcp_tool_tdef_groups[] = {
	{ .name = "pfcp", .tdefs = osmo_pfcp_tdefs, .desc = "PFCP" },
	{}
};

void g_pfcp_tool_alloc(void *ctx)
{
	OSMO_ASSERT(g_pfcp_tool == NULL);
	g_pfcp_tool = talloc_zero(ctx, struct g_pfcp_tool);

	*g_pfcp_tool = (struct g_pfcp_tool){
		.vty_cfg = {
			.local_ip = talloc_strdup(g_pfcp_tool, "0.0.0.0"),
			.local_port = OSMO_PFCP_PORT,
		},
		.next_teid_state = 23,
		.gtp.flood = {
			.cfg = {
				.num_rx_workers = 1,
				.num_tx_workers = 1,
				.queue_size = 4096,
				.slew_us = 0,
			},
			.flows_per_session = 1,
			.packets_per_flow = 0,
		},
	};

	INIT_LLIST_HEAD(&g_pfcp_tool->peers);
	INIT_LLIST_HEAD(&g_pfcp_tool->gtp.gtp_local_addrs);
	INIT_LLIST_HEAD(&g_pfcp_tool->gtp.gtp_core_addrs);
}

struct pfcp_tool_peer *pfcp_tool_peer_find(const struct osmo_sockaddr *remote_addr)
{
	struct pfcp_tool_peer *peer;
	llist_for_each_entry(peer, &g_pfcp_tool->peers, entry) {
		if (osmo_sockaddr_cmp(&peer->remote_addr, remote_addr) == 0)
			return peer;
	}
	return NULL;
}

struct pfcp_tool_peer *pfcp_tool_peer_find_or_create(const struct osmo_sockaddr *remote_addr)
{
	struct pfcp_tool_peer *peer = pfcp_tool_peer_find(remote_addr);
	if (peer)
		return peer;

	peer = talloc_zero(g_pfcp_tool, struct pfcp_tool_peer);
	peer->remote_addr = *remote_addr;
	peer->next_seid_state = 0x1234567;
	INIT_LLIST_HEAD(&peer->sessions);
	llist_add(&peer->entry, &g_pfcp_tool->peers);
	return peer;
}

struct pfcp_tool_session *pfcp_tool_session_find(struct pfcp_tool_peer *peer, uint64_t cp_seid)
{
	struct pfcp_tool_session *session;
	llist_for_each_entry(session, &peer->sessions, entry) {
		if (session->cp_seid == cp_seid)
			return session;
	}
	return NULL;
}

struct pfcp_tool_session *pfcp_tool_session_find_or_create(struct pfcp_tool_peer *peer, uint64_t cp_seid,
							   enum up_gtp_action_kind kind)
{
	struct pfcp_tool_session *session = pfcp_tool_session_find(peer, cp_seid);
	if (session)
		return session;

	session = talloc(peer, struct pfcp_tool_session);
	*session = (struct pfcp_tool_session){
		.peer = peer,
		.cp_seid = cp_seid,
		.kind = kind,
	};
	llist_add(&session->entry, &peer->sessions);
	return session;
}

static void rx_assoc_setup_resp(struct osmo_pfcp_endpoint *ep, struct osmo_pfcp_msg *m)
{
	if (m->ies.assoc_setup_resp.up_function_features_present)
		OSMO_LOG_PFCP_MSG(m, LOGL_NOTICE, "Associated. UP Peer features: %s\n",
				  osmo_pfcp_bits_to_str_c(OTC_SELECT,
							  m->ies.assoc_setup_resp.up_function_features.bits,
							  osmo_pfcp_up_feature_strs));

	if (m->ies.assoc_setup_resp.cp_function_features_present)
		OSMO_LOG_PFCP_MSG(m, LOGL_NOTICE, "Associated. CP Peer features: %s\n",
				  osmo_pfcp_bits_to_str_c(OTC_SELECT,
							  m->ies.assoc_setup_resp.cp_function_features.bits,
							  osmo_pfcp_cp_feature_strs));
}

static void rx_session_est_resp(struct osmo_pfcp_endpoint *ep, struct osmo_pfcp_msg *m)
{
	struct pfcp_tool_peer *peer;
	struct pfcp_tool_session *session;
	enum osmo_pfcp_cause *cause = osmo_pfcp_msg_cause(m);
	if (!cause) {
		OSMO_LOG_PFCP_MSG(m, LOGL_ERROR, "Session Establishment Response should contain a Cause\n");
		return;
	}
	if (*cause != OSMO_PFCP_CAUSE_REQUEST_ACCEPTED) {
		OSMO_LOG_PFCP_MSG(m, LOGL_ERROR, "Peer responds that Session Establishment failed\n");
		return;
	}
	if (!m->h.seid_present) {
		OSMO_LOG_PFCP_MSG(m, LOGL_ERROR, "Session Establishment Response should contain a SEID\n");
		return;
	}
	if (!m->ies.session_est_resp.up_f_seid_present) {
		OSMO_LOG_PFCP_MSG(m, LOGL_ERROR, "Session Establishment Response without UP F-SEID\n");
		return;
	}
	peer = pfcp_tool_peer_find(&m->remote_addr);
	if (!peer)
		return;
	session = pfcp_tool_session_find(peer, m->h.seid);
	if (!session)
		return;
	session->up_f_seid = m->ies.session_est_resp.up_f_seid;
}

void pfcp_tool_rx_msg(struct osmo_pfcp_endpoint *ep, struct osmo_pfcp_msg *m, struct osmo_pfcp_msg *req)
{
	switch (m->h.message_type) {
	case OSMO_PFCP_MSGT_ASSOC_SETUP_RESP:
		rx_assoc_setup_resp(ep, m);
		break;
	case OSMO_PFCP_MSGT_SESSION_EST_RESP:
		rx_session_est_resp(ep, m);
		break;
	default: break;
	}
}

static void copy_msg(struct osmo_pfcp_msg *dst, const struct osmo_pfcp_msg *m)
{
	*dst = *m;
	dst->encoded = NULL;
	dst->ctx.peer_use_token = NULL;
	dst->ctx.session_use_token = NULL;
	dst->ctx.resp_cb = NULL;
}

int peer_tx(struct pfcp_tool_peer *peer, struct osmo_pfcp_msg *m)
{
	int rc;
	if (m->is_response)
		copy_msg(&peer->last_resp, m);
	else
		copy_msg(&peer->last_req, m);
	rc = osmo_pfcp_endpoint_tx(g_pfcp_tool->ep, m);
	if (rc)
		LOGP(DLPFCP, LOGL_ERROR, "Failed to transmit PFCP: %s\n", strerror(-rc));
	return rc;
}

uint64_t peer_new_seid(struct pfcp_tool_peer *peer)
{
	return peer->next_seid_state++;
}

uint32_t pfcp_tool_new_teid(void)
{
	return g_pfcp_tool->next_teid_state++;
}

int pfcp_tool_next_ue_addr(struct osmo_sockaddr *dst)
{
	range_next(&g_pfcp_tool->ue.ip_range);
	if (g_pfcp_tool->ue.ip_range.wrapped) {
		LOGP(DLGLOBAL, LOGL_ERROR, "insufficient UE IP addresses, wrapped back to first\n");
		g_pfcp_tool->ue.ip_range.wrapped = false;
	}
	range_val_get_addr(dst, &g_pfcp_tool->ue.ip_range.next);
	return 0;
}

struct udp_port *pfcp_tool_have_udp_port_by_osa(struct llist_head *list, const struct osmo_sockaddr *_osa,
						uint16_t fallback_port)
{
	struct udp_port *port;
	/* copy osa and have a non-const pointer */
	struct osmo_sockaddr osa_mutable = *_osa;
	struct osmo_sockaddr *osa = &osa_mutable;

	if (!osmo_sockaddr_port(&osa->u.sa))
		osmo_sockaddr_set_port(&osa->u.sa, fallback_port);
	OSMO_ASSERT(osmo_sockaddr_port(&osa->u.sa));

	llist_for_each_entry(port, list, entry) {
		if (osmo_sockaddr_cmp(&port->osa, osa) == 0)
			return port;
	}

	port = talloc_zero(g_pfcp_tool, struct udp_port);
	port->osa = *osa;
	port->ofd.fd = -1;

	llist_add_tail(&port->entry, list);
	return port;
}

struct udp_port *pfcp_tool_have_udp_port_by_str(struct llist_head *list, const struct osmo_sockaddr_str *addr,
						uint16_t fallback_port)
{
	struct osmo_sockaddr osa;
	if (osmo_sockaddr_str_to_osa(addr, &osa))
		return NULL;
	return pfcp_tool_have_udp_port_by_osa(list, &osa, fallback_port);
}

struct udp_port *pfcp_tool_have_local_udp_port_by_osa(const struct osmo_sockaddr *osa, uint16_t fallback_port)
{
	struct udp_port *port = pfcp_tool_have_udp_port_by_osa(&g_pfcp_tool->gtp.gtp_local_addrs, osa, fallback_port);

	/* already bound? */
	if (port->ofd.fd >= 0)
		return port;

	/* create and bind socket */
	int rc;
	rc = osmo_sock_init_osa(SOCK_DGRAM, IPPROTO_UDP, &port->osa, NULL, OSMO_SOCK_F_BIND);
	if (rc < 0) {
		LOGP(DLGLOBAL, LOGL_ERROR, "Failed to bind socket on UDP %s\n",
		     osmo_sockaddr_to_str_c(OTC_SELECT, &port->osa));
		exit(1);
	}
	port->ofd.fd = rc;

#if 0
	int a_lot;
	a_lot = 1024 * (1024*1024);
	rc = setsockopt(port->ofd.fd, SOL_SOCKET, SO_SNDBUF, &a_lot, sizeof(a_lot));
	a_lot = 1024 * (1024*1024);
	rc = setsockopt(port->ofd.fd, SOL_SOCKET, SO_RCVBUF, &a_lot, sizeof(a_lot));
#endif

	LOGP(DLGLOBAL, LOGL_NOTICE, "bound UDP %s fd=%d\n",
	     osmo_sockaddr_to_str_c(OTC_SELECT, &port->osa), rc);

	/* Set Don't Fragment (DF) bit on IP packets transmitted by socket: */
	int val = IP_PMTUDISC_DO;
	rc = setsockopt(port->ofd.fd, IPPROTO_IP, IP_MTU_DISCOVER, &val, sizeof(val));
	if (rc == -1)
		LOGP(DLGLOBAL, LOGL_ERROR, "setsockopt(IPPROTO_IP, IP_DONTFRAG) failed errno=%d", errno);
	return port;
}

struct udp_port *pfcp_tool_have_local_udp_port_by_str(const struct osmo_sockaddr_str *addr, uint16_t fallback_port)
{
	struct osmo_sockaddr osa;
	if (osmo_sockaddr_str_to_osa(addr, &osa))
		return NULL;
	return pfcp_tool_have_local_udp_port_by_osa(&osa, fallback_port);
}

struct udp_port *pfcp_tool_have_core_udp_port_by_osa(const struct osmo_sockaddr *osa, uint16_t fallback_port)
{
	return pfcp_tool_have_udp_port_by_osa(&g_pfcp_tool->gtp.gtp_core_addrs, osa, fallback_port);
}

struct udp_port *pfcp_tool_have_core_udp_port_by_str(const struct osmo_sockaddr_str *addr, uint16_t fallback_port)
{
	struct osmo_sockaddr osa;
	if (osmo_sockaddr_str_to_osa(addr, &osa))
		return NULL;
	return pfcp_tool_have_core_udp_port_by_osa(&osa, fallback_port);
}

static bool add_session_to_gtp_flow(struct gtp_flood *gf, struct pfcp_tool_session *session)
{
	struct range *r;
	struct gtp_flood_flow_cfg cfg = {};
	const struct pfcp_tool_gtp_tun *tun_access;
	const struct pfcp_tool_gtp_tun *tun_core = NULL;
	const struct osmo_sockaddr_str *ue_local_addr = NULL;

	switch (session->kind) {
	case UP_GTP_U_TUNEND:
		tun_access = &session->tunend.access;
		ue_local_addr = &session->tunend.core.ue_local_addr;
		break;
	case UP_GTP_U_TUNMAP:
		tun_access = &session->tunmap.access;
		tun_core = &session->tunmap.core;
		break;
	default:
		OSMO_ASSERT(false);
	}

	/* The 'local' and 'remote' naming clashes here. struct pfcp_tool_gtp_tun names its items from the UPF's point
	 * of view: so the GTP port of osmo-upf-load-gen is 'remote'.
	 * Here we are setting up a port to emit GTP from osmo-upf-load-gen, the same port is called 'gtp_local'.
	 */
	cfg.gtp_local = pfcp_tool_have_local_udp_port_by_str(&tun_access->remote.addr, 2152);
	if (!cfg.gtp_local) {
		LOGP(DLGLOBAL, LOGL_ERROR, "Cannot set up GTP flow for session, failed to setup local GTP port\n");
		return false;
	}

	/* The 'local' and 'remote' naming clashes here. struct pfcp_tool_gtp_tun names its items from the UPF's point
	 * of view: so the GTP port of UPF is 'local'.
	 * Here we are reading the GTP port that the UPF has opened to receive GTP traffic; the same port is called
	 * 'gtp_remote' in cfg, which is remote from osmo-upf-load-gen's point of view.
	 */
	if (osmo_sockaddr_str_to_osa(&tun_access->local.addr, &cfg.gtp_remote)) {
		LOGP(DLGLOBAL, LOGL_ERROR, "Cannot set up GTP flow for session, failed to identify UPF's GTP port\n");
		return false;
	}
	if (!osmo_sockaddr_port(&cfg.gtp_remote.u.sa))
		osmo_sockaddr_set_port(&cfg.gtp_remote.u.sa, 2152);
	cfg.gtp_remote_teid = tun_access->local.teid;

	if (ue_local_addr) {
		osmo_sockaddr_str_to_osa(ue_local_addr, &cfg.payload_src);
	} else if (pfcp_tool_next_ue_addr(&cfg.payload_src)) {
		LOGP(DLGLOBAL, LOGL_ERROR, "Cannot set up GTP flow for session, failed to identify UE's IP address\n");
		return false;
	}
	r = &g_pfcp_tool->gtp.payload.source.udp_port_range;
	range_next(r);
	osmo_sockaddr_set_port(&cfg.payload_src.u.sa, range_val_get_int(&r->next));

	r = &g_pfcp_tool->gtp.payload.target.ip_range;
	range_next(r);
	range_val_get_addr(&cfg.payload_dst, &r->next);
	r = &g_pfcp_tool->gtp.payload.target.udp_port_range;
	range_next(r);
	osmo_sockaddr_set_port(&cfg.payload_dst.u.sa, range_val_get_int(&r->next));

	cfg.num_packets = g_pfcp_tool->gtp.flood.packets_per_flow;

	if (g_pfcp_tool->gtp.payload.append_info) {
		cfg.append_payload_info = true;
		cfg.payload_info = (struct gtp_flood_payload_info){};
		if (tun_core)
			cfg.payload_info.return_teid = htonl(tun_core->local.teid);
	}

	for (int i = 0; i < g_pfcp_tool->gtp.flood.flows_per_session; i++)
		gtp_flood_add_flow(gf, &cfg);

	return true;
}

void pfcp_tool_gtp_flood_start(void)
{
	struct gtp_flood *gf;
	struct pfcp_tool_peer *peer;
	struct udp_port *port;

	if (g_pfcp_tool->gtp.flood.state) {
		LOGP(DLGLOBAL, LOGL_ERROR,
		     "another 'gtp flood' is still running; currently we can run only one gtp flood at a time\n");
		return;
	}

	gf = g_pfcp_tool->gtp.flood.state = gtp_flood_alloc(g_pfcp_tool, &g_pfcp_tool->gtp.flood.cfg);
	if (!gf)
		return;

	/* add transmitter flows */
	llist_for_each_entry(peer, &g_pfcp_tool->peers, entry) {
		struct pfcp_tool_session *session;
		llist_for_each_entry(session, &peer->sessions, entry) {
			add_session_to_gtp_flow(gf, session);
		}
	}

	unsigned int rx_flows_added = 0;
	/* add listener flows, one per local UDP port (like a GTP port from a 'gtp local 1.2.3.4' vty command) */
	llist_for_each_entry(port, &g_pfcp_tool->gtp.gtp_local_addrs, entry) {
		struct gtp_flood_flow_cfg cfg = {
			.rx = true,
			.gtp_local = port,
		};
		gtp_flood_add_flow(gf, &cfg);
		rx_flows_added++;
	}
	/* If there's still free Rx workers, use them to spread the load among the Rx flows: */
	do {
		llist_for_each_entry(port, &g_pfcp_tool->gtp.gtp_local_addrs, entry) {
			if (rx_flows_added >= g_pfcp_tool->gtp.flood.cfg.num_rx_workers)
				goto done_filling_rx_workers;
			struct gtp_flood_flow_cfg cfg = {
				.rx = true,
				.gtp_local = port,
			};
			gtp_flood_add_flow(gf, &cfg);
			rx_flows_added++;
		}
	} while (true);
done_filling_rx_workers:

	gtp_flood_start(gf);
}