/*
 * (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 <osmocom/core/fsm.h>
#include <osmocom/core/logging.h>
#include <osmocom/core/tdef.h>
#include <osmocom/core/utils.h>

#include <osmocom/pfcp/pfcp_msg.h>
#include <osmocom/pfcp/pfcp_endpoint.h>

#include <osmocom/upf/upf.h>
#include <osmocom/upf/up_peer.h>
#include <osmocom/upf/up_endpoint.h>
#include <osmocom/upf/up_session.h>

enum up_peer_fsm_state {
	UP_PEER_ST_NOT_ASSOCIATED,
	UP_PEER_ST_ASSOCIATED,
	UP_PEER_ST_GRACEFUL_RELEASE,
	UP_PEER_ST_WAIT_USE_COUNT,
};

static const struct value_string up_peer_fsm_event_names[] = {
	OSMO_VALUE_STRING(UP_PEER_EV_RX_ASSOC_SETUP_REQ),
	OSMO_VALUE_STRING(UP_PEER_EV_RX_ASSOC_UPD_REQ),
	OSMO_VALUE_STRING(UP_PEER_EV_RX_ASSOC_REL_REQ),
	OSMO_VALUE_STRING(UP_PEER_EV_RX_SESSION_EST_REQ),
	OSMO_VALUE_STRING(UP_PEER_EV_HEARTBEAT_FAILURE),
	OSMO_VALUE_STRING(UP_PEER_EV_USE_COUNT_ZERO),
	OSMO_VALUE_STRING(UP_PEER_EV_SESSION_TERM),
	{}
};

static struct osmo_fsm up_peer_fsm;

static const struct osmo_tdef_state_timeout up_peer_fsm_timeouts[32] = {
	[UP_PEER_ST_GRACEFUL_RELEASE] = { .T = -21 },
};

/* Transition to a state, using the T timer defined in up_peer_fsm_timeouts.
 * Assumes local variable fi exists. */
#define up_peer_fsm_state_chg(state) \
	osmo_tdef_fsm_inst_state_chg(fi, state, \
				     up_peer_fsm_timeouts, \
				     osmo_pfcp_tdefs, \
				     5)

static int up_peer_use_cb(struct osmo_use_count_entry *e, int32_t old_use_count, const char *file, int line)
{
	struct up_peer *peer = e->use_count->talloc_object;
	int32_t total;
	int level;

	if (!e->use)
		return -EINVAL;

	total = osmo_use_count_total(&peer->use_count);

	if (total == 0
	    || (total == 1 && old_use_count == 0 && e->count == 1))
		level = LOGL_INFO;
	else
		level = LOGL_DEBUG;

	LOGPFSMSLSRC(peer->fi, DREF, level, file, line,
		     "%s %s: now used by %s\n",
		     (e->count - old_use_count) > 0 ? "+" : "-", e->use,
		     osmo_use_count_to_str_c(OTC_SELECT, &peer->use_count));

	if (e->count < 0)
		return -ERANGE;

	if (total == 0)
		osmo_fsm_inst_dispatch(peer->fi, UP_PEER_EV_USE_COUNT_ZERO, NULL);
	return 0;
}

char *up_peer_remote_addr_str(struct up_peer *peer)
{
	struct osmo_sockaddr remote_addr = peer->remote_addr;

	/* Zero the port, it is not interesting information. The port for PFCP is defined fixed, and there is no use
	 * printing it in the logs */
	osmo_sockaddr_set_port(&remote_addr.u.sa, 0);

	return osmo_sockaddr_to_str_c(OTC_SELECT, &remote_addr);
}

static void up_peer_update_id(struct up_peer *peer)
{
	osmo_fsm_inst_update_id_f_sanitize(peer->fi, '-', "%s", up_peer_remote_addr_str(peer));
	LOGPFSML(peer->fi, LOGL_DEBUG, "Updated id\n");
}

static struct up_peer *up_peer_add(struct up_endpoint *up_endpoint, const struct osmo_sockaddr *remote_addr)
{
	struct up_peer *peer;

	struct osmo_fsm_inst *fi = osmo_fsm_inst_alloc(&up_peer_fsm, up_endpoint, NULL, LOGL_DEBUG, NULL);
	OSMO_ASSERT(fi);

	peer = talloc(fi, struct up_peer);
	OSMO_ASSERT(peer);
	fi->priv = peer;

	*peer = (struct up_peer) {
		.fi = fi,
		.up_endpoint = up_endpoint,
		.remote_addr = *remote_addr,
		.heartbeat_fi = NULL /* FIXME */,
		.use_count = {
			.talloc_object = peer,
			.use_cb = up_peer_use_cb,
		},
	};
	osmo_use_count_make_static_entries(&peer->use_count, peer->use_count_buf, ARRAY_SIZE(peer->use_count_buf));
	hash_init(peer->sessions_by_up_seid);
	hash_init(peer->sessions_by_cp_seid);

	osmo_pfcp_bits_set(peer->local_up_features.bits, OSMO_PFCP_UP_FEAT_BUNDL, true);
	osmo_pfcp_bits_set(peer->local_up_features.bits, OSMO_PFCP_UP_FEAT_RTTL, true);
	osmo_pfcp_bits_set(peer->local_up_features.bits, OSMO_PFCP_UP_FEAT_FTUP, true);

	up_peer_update_id(peer);

	llist_add(&peer->entry, &up_endpoint->peers);
	return peer;
}

struct up_peer *up_peer_find(struct up_endpoint *up_endpoint, const struct osmo_sockaddr *remote_addr)
{
	struct up_peer *peer;
	llist_for_each_entry(peer, &up_endpoint->peers, entry) {
		if (osmo_sockaddr_cmp(&peer->remote_addr, remote_addr))
			continue;
		return peer;
	}
	return NULL;
}

struct up_peer *up_peer_find_or_add(struct up_endpoint *up_endpoint, const struct osmo_sockaddr *remote_addr)
{
	struct up_peer *peer = up_peer_find(up_endpoint, remote_addr);
	if (peer)
		return peer;
	return up_peer_add(up_endpoint, remote_addr);
}

static int up_peer_fsm_timer_cb(struct osmo_fsm_inst *fi)
{
	//struct up_peer *peer = fi->priv;
	/* Return 1 to terminate FSM instance, 0 to keep running */
	return 1;
}

void up_peer_set_msg_ctx(struct up_peer *peer, struct osmo_pfcp_msg *m)
{
	OSMO_ASSERT(!m->ctx.peer_fi);

	m->ctx.peer_fi = peer->fi;
	m->ctx.peer_use_count = &peer->use_count;
	m->ctx.peer_use_token = (m->rx ? UP_USE_MSG_RX : UP_USE_MSG_TX);
	OSMO_ASSERT(osmo_use_count_get_put(m->ctx.peer_use_count, m->ctx.peer_use_token, 1) == 0);
}

struct osmo_pfcp_msg *up_peer_init_tx(struct up_peer *peer, struct osmo_pfcp_msg *in_reply_to,
				      enum osmo_pfcp_message_type message_type)
{
	struct osmo_pfcp_msg *tx;
	if (in_reply_to)
		tx = osmo_pfcp_msg_alloc_tx_resp(OTC_SELECT, in_reply_to, message_type);
	else
		tx = osmo_pfcp_msg_alloc_tx_req(OTC_SELECT, &peer->remote_addr, message_type);
	up_peer_set_msg_ctx(peer, tx);
	return tx;
}

static int up_peer_tx_assoc_setup_resp(struct up_peer *peer, struct osmo_pfcp_msg *m, enum osmo_pfcp_cause cause)
{
	struct osmo_pfcp_msg *resp;

	resp = up_peer_init_tx(peer, m, OSMO_PFCP_MSGT_ASSOC_SETUP_RESP);

	resp->ies.assoc_setup_resp = (struct osmo_pfcp_msg_assoc_setup_resp) {
		.cause = cause,
		.recovery_time_stamp = osmo_pfcp_endpoint_get_recovery_timestamp(g_upf->pfcp.ep->pfcp_ep),
		.up_function_features_present = true,
		.up_function_features = peer->local_up_features,
	};

	if (osmo_pfcp_endpoint_tx(peer->up_endpoint->pfcp_ep, resp)) {
		OSMO_LOG_PFCP_MSG(m, LOGL_ERROR, "Error sending response to this message,"
				  " cannot associate with peer\n");
		return -EIO;
	}
	return 0;
}

static int up_peer_tx_assoc_rel_resp(struct up_peer *peer, struct osmo_pfcp_msg *m, enum osmo_pfcp_cause cause)
{
	struct osmo_pfcp_msg *resp;

	resp = up_peer_init_tx(peer, m, OSMO_PFCP_MSGT_ASSOC_RELEASE_RESP);

	resp->ies.assoc_release_resp = (struct osmo_pfcp_msg_assoc_release_resp) {
		.cause = cause,
	};

	if (osmo_pfcp_endpoint_tx(peer->up_endpoint->pfcp_ep, resp)) {
		OSMO_LOG_PFCP_MSG(m, LOGL_ERROR, "Error sending response to this message\n");
		return -EIO;
	}
	return 0;
}

static void up_peer_clear_sessions(struct up_peer *peer)
{
	struct up_session *session;
	int bkt;
	struct hlist_node *tmp;
	int count = 0;
	hash_for_each_safe(peer->sessions_by_up_seid, bkt, tmp, session, node_by_up_seid) {
		count += up_session_discard(session);
	}
	if (count)
		LOGPFSML(peer->fi, LOGL_NOTICE, "terminated %d sessions\n", count);
}

static void up_peer_rx_assoc_setup_req(struct up_peer *peer, struct osmo_pfcp_msg *m)
{
	struct osmo_fsm_inst *fi = peer->fi;
	enum osmo_pfcp_cause cause = OSMO_PFCP_CAUSE_REQUEST_ACCEPTED;

	if (fi->state == UP_PEER_ST_ASSOCIATED) {
		/* Retransmissions of the ACK response happen in pfcp_endpoint.c. So if we get this, it is a genuine
		 * duplicate association setup request. We could reject it. But why. Just "replace" with the new
		 * association. Continue. */
		/* If the peer has restarted, it has forgotten about all sessions. */
		if (peer->remote_recovery_timestamp != m->ies.assoc_setup_req.recovery_time_stamp) {
			LOGPFSML(fi, LOGL_NOTICE, "another Association Setup Request, with different Recovery Timestamp."
				 " Clearing sessions, sending ACK.\n");
			up_peer_clear_sessions(peer);
		} else {
			LOGPFSML(fi, LOGL_NOTICE, "another Association Setup Request, with same Recovery Timestamp."
				 " Keeping sessions, sending ACK.\n");
		}
	} else if (up_peer_fsm_state_chg(UP_PEER_ST_ASSOCIATED)) {
		/* Not allowed to transition to ST_ASSOCIATED */
		cause = OSMO_PFCP_CAUSE_REQUEST_REJECTED;
	} else {
		/* Successfully transitioned to ST_ASSOCIATED */
		peer->remote_recovery_timestamp = m->ies.assoc_setup_req.recovery_time_stamp;
		peer->remote_node_id = m->ies.assoc_setup_req.node_id;
		if (m->ies.assoc_setup_req.cp_function_features_present)
			peer->peer_cp_features = m->ies.assoc_setup_req.cp_function_features;
	}

	if (up_peer_tx_assoc_setup_resp(peer, m, cause)
	    || cause != OSMO_PFCP_CAUSE_REQUEST_ACCEPTED)
		up_peer_fsm_state_chg(UP_PEER_ST_WAIT_USE_COUNT);

	LOGPFSML(fi, LOGL_NOTICE, "Peer associated, Node-Id=%s. Local UP features: [%s]; Peer CP features: [%s]\n",
		 osmo_pfcp_ie_node_id_to_str_c(OTC_SELECT, &peer->remote_node_id),
		 osmo_pfcp_bits_to_str_c(OTC_SELECT, peer->local_up_features.bits, osmo_pfcp_up_feature_strs),
		 osmo_pfcp_bits_to_str_c(OTC_SELECT, peer->peer_cp_features.bits, osmo_pfcp_cp_feature_strs));
}

static void up_peer_rx_assoc_rel_req(struct up_peer *peer, struct osmo_pfcp_msg *m)
{
	struct osmo_fsm_inst *fi = peer->fi;
	up_peer_tx_assoc_rel_resp(peer, m, OSMO_PFCP_CAUSE_REQUEST_ACCEPTED);
	up_peer_fsm_state_chg(UP_PEER_ST_WAIT_USE_COUNT);
}

static void up_peer_rx_session_est_req(struct up_peer *peer, struct osmo_pfcp_msg *m)
{
	enum osmo_pfcp_cause cause = OSMO_PFCP_CAUSE_REQUEST_ACCEPTED;
	struct osmo_pfcp_msg *resp;
	struct up_session *session = up_session_find_or_add(peer, &m->ies.session_est_req.cp_f_seid);

	if (!session) {
		cause = OSMO_PFCP_CAUSE_NO_RESOURCES_AVAILABLE;
		goto nack_response;
	}

	up_session_set_msg_ctx(session, m);

	if (osmo_fsm_inst_dispatch(session->fi, UP_SESSION_EV_RX_SESSION_EST_REQ, m)) {
		cause = OSMO_PFCP_CAUSE_REQUEST_REJECTED;
		goto nack_response;
	}
	return;

nack_response:
	resp = up_peer_init_tx(peer, m, OSMO_PFCP_MSGT_SESSION_EST_RESP);
	resp->h.seid = m->ies.session_est_req.cp_f_seid.seid;
	resp->h.seid_present = true;
	resp->ies.session_est_resp = (struct osmo_pfcp_msg_session_est_resp){
		.cause = cause,
	};
	osmo_pfcp_endpoint_tx(peer->up_endpoint->pfcp_ep, resp);
	if (session)
		up_session_discard(session);
}

static void up_peer_not_associated_action(struct osmo_fsm_inst *fi, uint32_t event, void *data)
{
	struct up_peer *peer = fi->priv;

	switch (event) {

	case UP_PEER_EV_RX_ASSOC_SETUP_REQ:
		up_peer_rx_assoc_setup_req(peer, data);
		break;

	case UP_PEER_EV_USE_COUNT_ZERO:
		/* Not associated and no pending messages. discard peer. */
		up_peer_fsm_state_chg(UP_PEER_ST_WAIT_USE_COUNT);
		return;

	default:
		OSMO_ASSERT(false);
	}
}

static void up_peer_associated_action(struct osmo_fsm_inst *fi, uint32_t event, void *data)
{
	struct up_peer *peer = fi->priv;

	switch (event) {

	case UP_PEER_EV_RX_ASSOC_SETUP_REQ:
		up_peer_rx_assoc_setup_req(peer, data);
		break;

	case UP_PEER_EV_RX_ASSOC_UPD_REQ:
		// FIXME
		break;

	case UP_PEER_EV_RX_SESSION_EST_REQ:
		up_peer_rx_session_est_req(peer, data);
		break;

	case UP_PEER_EV_HEARTBEAT_FAILURE:
		// FIXME
		break;

	case UP_PEER_EV_USE_COUNT_ZERO:
		/* Stay associated. */
		return;

	default:
		OSMO_ASSERT(false);
	}
}

static void up_peer_associated_onleave(struct osmo_fsm_inst *fi, uint32_t next_state)
{
	struct up_peer *peer = fi->priv;
	if (next_state != UP_PEER_ST_ASSOCIATED)
		LOGPFSML(fi, LOGL_NOTICE, "Peer %s released\n",
			 osmo_pfcp_ie_node_id_to_str_c(OTC_SELECT, &peer->remote_node_id));
}

static void up_peer_graceful_release_onenter(struct osmo_fsm_inst *fi, uint32_t prev_state)
{
	//struct up_peer *peer = fi->priv;
	// FIXME
}

static void up_peer_graceful_release_action(struct osmo_fsm_inst *fi, uint32_t event, void *data)
{
	struct up_peer *peer = fi->priv;

	switch (event) {

	case UP_PEER_EV_HEARTBEAT_FAILURE:
		up_peer_fsm_state_chg(UP_PEER_ST_WAIT_USE_COUNT);
		break;

	case UP_PEER_EV_USE_COUNT_ZERO:
		/* When there are still sessions, stay around. */
		if (!hash_empty(peer->sessions_by_up_seid))
			return;
		up_peer_fsm_state_chg(UP_PEER_ST_WAIT_USE_COUNT);
		return;

	default:
		OSMO_ASSERT(false);
	}
}

static void up_peer_wait_use_count_onenter(struct osmo_fsm_inst *fi, uint32_t prev_state)
{
	struct up_peer *peer = fi->priv;
	up_peer_clear_sessions(peer);
	if (!osmo_use_count_total(&peer->use_count))
		osmo_fsm_inst_term(fi, OSMO_FSM_TERM_REGULAR, NULL);
}

static void up_peer_wait_use_count_action(struct osmo_fsm_inst *fi, uint32_t event, void *data)
{
	struct up_peer *peer = fi->priv;
	switch (event) {

	case UP_PEER_EV_USE_COUNT_ZERO:
		osmo_fsm_inst_term(fi, OSMO_FSM_TERM_REGULAR, NULL);
		return;

	case UP_PEER_EV_RX_ASSOC_SETUP_REQ:
		up_peer_rx_assoc_setup_req(peer, data);
		break;

	default:
		OSMO_ASSERT(false);
	}
}

#define S(x)    (1 << (x))

static const struct osmo_fsm_state up_peer_fsm_states[] = {
	[UP_PEER_ST_NOT_ASSOCIATED] = {
		.name = "NOT_ASSOCIATED",
		.in_event_mask = 0
			| S(UP_PEER_EV_RX_ASSOC_SETUP_REQ)
			| S(UP_PEER_EV_USE_COUNT_ZERO)
			,
		.out_state_mask = 0
			| S(UP_PEER_ST_ASSOCIATED)
			| S(UP_PEER_ST_WAIT_USE_COUNT)
			,
		.action = up_peer_not_associated_action,
	},
	[UP_PEER_ST_ASSOCIATED] = {
		.name = "ASSOCIATED",
		.in_event_mask = 0
			| S(UP_PEER_EV_RX_ASSOC_SETUP_REQ)
			| S(UP_PEER_EV_RX_ASSOC_UPD_REQ)
			| S(UP_PEER_EV_RX_SESSION_EST_REQ)
			| S(UP_PEER_EV_HEARTBEAT_FAILURE)
			| S(UP_PEER_EV_USE_COUNT_ZERO)
			,
		.out_state_mask = 0
			| S(UP_PEER_ST_ASSOCIATED)
			| S(UP_PEER_ST_GRACEFUL_RELEASE)
			| S(UP_PEER_ST_WAIT_USE_COUNT)
			,
		.action = up_peer_associated_action,
		.onleave = up_peer_associated_onleave,
	},
	[UP_PEER_ST_GRACEFUL_RELEASE] = {
		.name = "GRACEFUL_RELEASE",
		.in_event_mask = 0
			| S(UP_PEER_EV_HEARTBEAT_FAILURE)
			| S(UP_PEER_EV_USE_COUNT_ZERO)
			,
		.out_state_mask = 0
			| S(UP_PEER_ST_WAIT_USE_COUNT)
			,
		.onenter = up_peer_graceful_release_onenter,
		.action = up_peer_graceful_release_action,
	},
	[UP_PEER_ST_WAIT_USE_COUNT] = {
		.name = "WAIT_USE_COUNT",
		.in_event_mask = 0
			| S(UP_PEER_EV_USE_COUNT_ZERO)
			| S(UP_PEER_EV_RX_ASSOC_SETUP_REQ)
			,
		.out_state_mask = 0
			| S(UP_PEER_ST_ASSOCIATED)
			,
		.onenter = up_peer_wait_use_count_onenter,
		.action = up_peer_wait_use_count_action,
	},
};

void up_peer_fsm_cleanup(struct osmo_fsm_inst *fi, enum osmo_fsm_term_cause cause)
{
	struct up_peer *peer = fi->priv;
	LOGPFSML(fi, LOGL_NOTICE, "Peer removed\n");
	llist_del(&peer->entry);
}

static void up_peer_allstate_action(struct osmo_fsm_inst *fi, uint32_t event, void *data)
{
	switch (event) {
	case UP_PEER_EV_SESSION_TERM:
		/* ignore */
		return;
	case UP_PEER_EV_RX_ASSOC_REL_REQ:
		up_peer_rx_assoc_rel_req(fi->priv, data);
		return;
	default:
		OSMO_ASSERT(false);
	}
}

static struct osmo_fsm up_peer_fsm = {
	.name = "up_peer",
	.log_subsys = DPEER,
	.states = up_peer_fsm_states,
	.num_states = ARRAY_SIZE(up_peer_fsm_states),
	.event_names = up_peer_fsm_event_names,
	.timer_cb = up_peer_fsm_timer_cb,
	.cleanup = up_peer_fsm_cleanup,
	.allstate_event_mask = 0
		| S(UP_PEER_EV_RX_ASSOC_REL_REQ)
		| S(UP_PEER_EV_SESSION_TERM)
		,
	.allstate_action = up_peer_allstate_action,
};

static __attribute__((constructor)) void up_peer_fsm_register(void)
{
	OSMO_ASSERT(osmo_fsm_register(&up_peer_fsm) == 0);
}

void up_peer_free(struct up_peer *peer)
{
	osmo_fsm_inst_term(peer->fi, OSMO_FSM_TERM_REGULAR, NULL);
}