// SPDX-License-Identifier: BSD-3-Clause OR Apache-2.0

#include "aap2/aap2_client.h"
#include "aap2/aap2_util.h"

#include "cla/posix/cla_tcp_util.h"

#include "platform/hal_io.h"
#include "platform/posix/socket_util.h"

#include "ud3tn/common.h"
#include "ud3tn/bundle.h"
#include "ud3tn/eid.h"

#include <pb_decode.h>
#include <pb_encode.h>

#include <errno.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <unistd.h>

struct aap2_client {
	int socket_fd;

	pb_istream_t pb_istream;

	bool is_configured;
	bool is_subscriber;

	char *ud3tn_node_id;
	char *registered_eid;

	uint32_t keepalive_timeout_s;
};

void aap2_client_close(struct aap2_client *client)
{
	if (!client)
		return;

	shutdown(client->socket_fd, SHUT_RDWR);
	close(client->socket_fd);
	// NOTE: free(NULL) is a NOP, so we can safely free() everything even if unconfigured.
	free(client->ud3tn_node_id);
	free(client->registered_eid);
	free(client);
}

static struct aap2_client *aap2_client_create_from_sockfd(int socket_fd, uint64_t timeout_ms)
{
	ASSERT(socket_fd >= 0);

	// Struct allocation and initialization

	struct aap2_client *const client = malloc(sizeof(struct aap2_client));

	if (!client) {
		LOG_ERRNO("AAP2Client", "malloc()", errno);
		return NULL;
	}

	client->socket_fd = socket_fd;
	client->pb_istream = (pb_istream_t) {
		&aap2_pb_recv_callback,
		(void *)(intptr_t)socket_fd,
		SIZE_MAX,
		NULL,
	};
	client->is_configured = false;
	client->is_subscriber = false;
	// Set EIDs to NULL so we are able to determine whether they are set.
	client->ud3tn_node_id = NULL;
	client->registered_eid = NULL;

	// Marker byte and WELCOME message processing

	uint8_t marker_byte;
	aap2_AAPMessage welcome_msg;

	if (timeout_ms != 0 && timeout_ms <= INT32_MAX) {
		switch (poll_recv_timeout(client->socket_fd, timeout_ms)) {
		case 0:
			// timeout
			LOG_DEBUG("AAP2Client: Timed out.");
			aap2_client_close(client);
			return NULL;
		case 1:
			break;
		default:
			// general error
			LOG_WARN("AAP2Client: Error while polling on socket for receiving data!");
			aap2_client_close(client);
			return NULL;
		}
	}

	// Check first byte received from the server, which must be 0x2f.
	if (tcp_recv_all(socket_fd, &marker_byte, 1) != 1 || marker_byte != 0x2f) {
		LOG_WARN("AAP2Client: Did not receive marker byte.");
		aap2_client_close(client);
		return NULL;
	}

	if (aap2_recv_message(&client->pb_istream, aap2_AAPMessage_fields, &welcome_msg)) {
		LOG_WARN("AAP2Client: Did not receive Welcome message.");
		aap2_client_close(client);
		return NULL;
	}

	if (welcome_msg.which_msg != aap2_AAPMessage_welcome_tag) {
		LOG_WARN("AAP2Client: Invalid AAP2 message received (expected Welcome).");
		pb_release(aap2_AAPMessage_fields, &welcome_msg);
		aap2_client_close(client);
		return NULL;
	}

	// Take ownership of node_id transmitted in Welcome message (which is just a normal
	// malloc-ed null-terminated string).
	client->ud3tn_node_id = welcome_msg.msg.welcome.node_id;
	welcome_msg.msg.welcome.node_id = NULL;

	// Release the rest of the Welcome message.
	pb_release(aap2_AAPMessage_fields, &welcome_msg);

	LOGF_DEBUG(
		"AAP2Client: Got Welcome message back: Connected to AAP2 agent with node ID: %s",
		client->ud3tn_node_id
	);

	if (!eidstr_is_valid(client->ud3tn_node_id)) {
		LOGF_WARN(
			"AAP2Client: Provided node ID is not a valid EID: %s",
			client->ud3tn_node_id
		);
		aap2_client_close(client);
		return NULL;
	}

	return client;
}

struct aap2_client *aap2_client_open_socket(const char *sock_path, uint64_t timeout_ms)
{
	ASSERT(sock_path != NULL);
	LOGF_DEBUG("AAP2Client: Connecting to socket path: %s", sock_path);

	const int socket_fd = socket(AF_UNIX, SOCK_STREAM, 0);

	if (socket_fd == -1) {
		LOG_ERRNO("AAP2Client", "socket(AF_UNIX)", errno);
		return NULL;
	}

	struct sockaddr_un addr = {
		.sun_family = AF_UNIX,
	};
	const int path_printed_chars = snprintf(
		addr.sun_path,
		sizeof(addr.sun_path),
		"%s",
		sock_path
	);

	if (path_printed_chars <= 0 || (unsigned int)path_printed_chars > sizeof(addr.sun_path)) {
		LOGF_WARN(
			"AAP2Client: Invalid socket path; len = %d, maxlen = %zu",
			path_printed_chars,
			sizeof(addr.sun_path)
		);
		return NULL;
	}

	if (connect(socket_fd, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
		LOG_ERRNO_WARN("AAP2Client", "connect(unix_domain_socket)", errno);
		close(socket_fd);
		return NULL;
	}

	return aap2_client_create_from_sockfd(socket_fd, timeout_ms);
}

struct aap2_client *aap2_client_open_tcp(const char *node, const char *service, uint64_t timeout_ms)
{
	LOGF_DEBUG("AAP2Client: Connecting using TCP to %s (service = %s)", node, service);

	const int socket_fd = create_tcp_socket(node, service, true, NULL);

	if (socket_fd == -1) {
		LOG_WARN("AAP2Client: Failed to open TCP socket.");
		return NULL;
	}

	return aap2_client_create_from_sockfd(socket_fd, timeout_ms);
}

enum aap2_client_error aap2_client_register(struct aap2_client *client,
					    const struct aap2_client_spec client_spec,
					    uint64_t timeout_ms)
{
	const char *eid = client_spec.eid;
	const char *secret = client_spec.secret;
	bool is_subscriber = client_spec.is_subscriber;

	if (!eid) {
		LOG_WARN("AAP2Client: Invalid argument to register function!");
		return AAP2_CLIENT_EARGUMENT;
	}

	LOGF_DEBUG("AAP2Client: Registering for EID %s (sub = %d)", eid, (int)is_subscriber);

	aap2_AAPMessage reg_msg = aap2_AAPMessage_init_default;
	aap2_AAPResponse response;

	reg_msg.which_msg = aap2_AAPMessage_config_tag;
	reg_msg.msg.config.endpoint_id = strdup(eid);
	if (!reg_msg.msg.config.endpoint_id) {
		pb_release(aap2_AAPMessage_fields, &reg_msg);
		LOG_ERRNO("AAP2Client", "strdup(eid)", errno);
		return AAP2_CLIENT_EFAIL;
	}
	reg_msg.msg.config.is_subscriber = is_subscriber;
	reg_msg.msg.config.auth_type = aap2_AuthType_AUTH_TYPE_DEFAULT;
	if (secret) {
		reg_msg.msg.config.secret = strdup(secret);
		if (!reg_msg.msg.config.secret) {
			pb_release(aap2_AAPMessage_fields, &reg_msg);
			LOG_ERRNO("AAP2Client", "strdup(secret)", errno);
			return AAP2_CLIENT_EFAIL;
		}
	}
	reg_msg.msg.config.keepalive_seconds = client_spec.keepalive_timeout_s;

	if (aap2_send_message(client->socket_fd, aap2_AAPMessage_fields, &reg_msg)) {
		pb_release(aap2_AAPMessage_fields, &reg_msg);
		LOG_WARN("AAP2Client: Failed to send ConnectionConfig message.");
		return AAP2_CLIENT_EFAIL;
	}

	// Take over the duplicated EID for putting into the client struct.
	char *const eid_copy = reg_msg.msg.config.endpoint_id;

	reg_msg.msg.config.endpoint_id = NULL;
	pb_release(aap2_AAPMessage_fields, &reg_msg);

	if (timeout_ms != 0 && timeout_ms <= INT32_MAX) {
		switch (poll_recv_timeout(client->socket_fd, timeout_ms)) {
		case 0:
			// timeout
			LOG_DEBUG("AAP2Client: Timed out.");
			free(eid_copy);
			return AAP2_CLIENT_ETIMEOUT;
		case 1:
			break;
		default:
			// general error
			LOG_WARN("AAP2Client: Error while polling on socket for receiving data!");
			free(eid_copy);
			return AAP2_CLIENT_EFAIL;
		}
	}

	if (aap2_recv_message(&client->pb_istream, aap2_AAPResponse_fields, &response)) {
		LOG_WARN("AAP2Client: Failed to receive ConnectionConfig response.");
		free(eid_copy);
		return AAP2_CLIENT_EFAIL;
	}

	if (response.response_status != aap2_ResponseStatus_RESPONSE_STATUS_SUCCESS) {
		LOGF_WARN(
			"AAP2Client: ConnectionConfig request was not successful; status = %d",
			response.response_status
		);
		free(eid_copy);
		if (response.response_status != aap2_ResponseStatus_RESPONSE_STATUS_UNAUTHORIZED)
			return AAP2_CLIENT_EUNAUTHORIZED;
		return AAP2_CLIENT_EFAIL;
	}

	client->registered_eid = eid_copy;
	client->is_subscriber = is_subscriber;
	client->keepalive_timeout_s = client_spec.keepalive_timeout_s;
	client->is_configured = true;

	LOGF_INFO(
		"AAP2Client: Successfully registered %s with EID: %s",
		is_subscriber ? "bundle source" : "bundle sink",
		eid
	);

	return AAP2_CLIENT_OK;
}

const char *aap2_client_get_node_id(const struct aap2_client *const client)
{
	ASSERT(client->ud3tn_node_id); // Would be a bug: open ensures it is set
	return client->ud3tn_node_id;
}

const char *aap2_client_get_registered_eid(const struct aap2_client *client)
{
	if (!client->is_configured)
		return NULL;
	return client->registered_eid;
}

bool aap2_client_is_configured(const struct aap2_client *const client)
{
	return client->is_configured;
}

bool aap2_client_is_subscriber(const struct aap2_client *const client)
{
	return client->is_subscriber;
}

char *aap2_client_get_eid_for_agent(const char *node_id, const char *agent_id)
{
	if (!node_id || !agent_id)
		return NULL;

	struct eid eid_parsed = eid_from_string(node_id);

	if (eid_is_error(eid_parsed))
		return NULL;

	struct eid agent_eid = eid_get_agent_eid(eid_parsed, agent_id);

	eid_free(eid_parsed);
	if (eid_is_error(agent_eid))
		return NULL;

	char *const result = eid_to_string(agent_eid);

	eid_free(agent_eid);
	return result;
}

enum aap2_client_error aap2_client_send(struct aap2_client *const client,
					const struct aap2_client_bundle_spec spec,
					const uint8_t *const payload, size_t payload_length,
					const uint64_t timeout_ms,
					struct aap2_client_bundle_spec *out_spec)
{
	ASSERT(client);

	if (!client->is_configured || client->is_subscriber) {
		LOG_WARN("AAP2Client: Invalid client configuration, cannot send bundles!");
		return AAP2_CLIENT_EARGUMENT;
	}

	LOGF_DEBUG(
		"AAP2Client: Sending bundle (len = %zu) to EID %s",
		payload_length,
		spec.dest_eid
	);

	if (!payload || !eidstr_is_valid(spec.dest_eid) ||
	    (spec.source_eid && !eidstr_is_valid(spec.source_eid)) ||
	    (spec.report_to_eid && !eidstr_is_valid(spec.report_to_eid))) {
		LOG_WARN("AAP2Client: Invalid payload or EID provided!");
		return AAP2_CLIENT_EARGUMENT;
	}

	// NOTE: DO NOT RELEASE THIS, it would double-free the EIDs!
	aap2_AAPMessage send_msg = aap2_AAPMessage_init_default;

	send_msg.which_msg = aap2_AAPMessage_adu_tag;
	send_msg.msg.adu.dst_eid = spec.dest_eid;
	send_msg.msg.adu.src_eid = spec.source_eid;
	send_msg.msg.adu.report_to_eid = spec.report_to_eid;
	send_msg.msg.adu.payload_length = payload_length;
	send_msg.msg.adu.adu_flags_count = aap2_set_adu_flags(
		send_msg.msg.adu.adu_flags,
		spec.flags
	);
	send_msg.msg.adu.creation_timestamp_ms = spec.creation_timestamp_dtn_ms;
	send_msg.msg.adu.sequence_number = spec.sequence_number;

	if (aap2_send_message(client->socket_fd, aap2_AAPMessage_fields, &send_msg) < 0) {
		LOG_WARN("AAP2Client: Cannot send message!");
		return AAP2_CLIENT_EFAIL;
	}

	struct tcp_write_to_socket_param wsp = {
		.socket_fd = client->socket_fd,
		.errno_ = 0,
	};

	tcp_write_to_socket(&wsp, payload, payload_length);
	if (wsp.errno_) {
		LOG_ERRNO_WARN("AAP2Client", "send()", wsp.errno_);
		return AAP2_CLIENT_EFAIL;
	}

	if (timeout_ms != 0 && timeout_ms <= INT32_MAX) {
		switch (poll_recv_timeout(client->socket_fd, timeout_ms)) {
		case 0:
			// timeout
			LOG_DEBUG("AAP2Client: Timed out.");
			return AAP2_CLIENT_ETIMEOUT;
		case 1:
			break;
		default:
			// general error
			LOG_WARN("AAP2Client: Error while polling on socket for receiving data!");
			return AAP2_CLIENT_EFAIL;
		}
	}

	aap2_AAPResponse response;

	if (aap2_recv_message(&client->pb_istream, aap2_AAPResponse_fields, &response)) {
		LOG_WARN("AAP2Client: Error receiving AAP2 response!");
		return AAP2_CLIENT_EFAIL;
	}

	if (response.response_status != aap2_ResponseStatus_RESPONSE_STATUS_SUCCESS) {
		LOGF_WARN(
			"AAP2Client: Response indicates failure sending ADU; status = %d",
			response.response_status
		);
		pb_release(aap2_AAPResponse_fields, &response);
		return AAP2_CLIENT_EFAIL;
	}

	LOG_DEBUG("AAP2Client: Server responded - successfully sent bundle.");

	if (!response.has_bundle_headers) {
		LOG_WARN("AAP2Client: Invalid response, does not contain bundle headers.");
		pb_release(aap2_AAPResponse_fields, &response);
		return AAP2_CLIENT_EFAIL;
	}

	if (out_spec) {
		out_spec->creation_timestamp_dtn_ms = response.bundle_headers.creation_timestamp_ms;
		out_spec->sequence_number = response.bundle_headers.sequence_number;
		// Move the allocated EID strings into out_spec
		out_spec->source_eid = response.bundle_headers.src_eid;
		response.bundle_headers.src_eid = NULL;
		out_spec->dest_eid = response.bundle_headers.dst_eid;
		response.bundle_headers.dst_eid = NULL;
		// Not returned by AAP 2.0 at the moment.
		out_spec->report_to_eid = NULL;
		out_spec->flags = (struct aap2_bundle_adu_flags){ false, false, false };
	}
	pb_release(aap2_AAPResponse_fields, &response);

	return AAP2_CLIENT_OK;
}

/**
 * @note BB: this function receives the payload of a bundle. Note that this is because the payload
 * and BundleADU are sent separately.
 */
static uint8_t *receive_payload(pb_istream_t *const istream, const size_t payload_length)
{
	if (payload_length > BUNDLE_MAX_SIZE) {
		LOGF_WARN(
			"AAP2Client: Payload too large: %zu > %zu",
			payload_length,
			BUNDLE_MAX_SIZE
		);
		return NULL;
	}

	// malloc(0) may return NULL on some C libs, but zero-length payloads are fine.
	uint8_t *const payload = payload_length ? malloc(payload_length) : malloc(1);

	if (!payload) {
		LOG_ERROR("AAP2Client: Payload alloc error!");
		return NULL;
	}

	const bool success = payload_length ? pb_read(istream, payload, payload_length) : true;

	if (!success) {
		LOG_WARN("AAP2Client: Payload read error!");
		free(payload);
		return NULL;
	}

	return payload;
}

enum aap2_client_error aap2_client_recv(struct aap2_client *const client,
					struct aap2_client_bundle_spec *const out_spec,
					uint8_t **const out_payload,
					size_t *const out_payload_length,
					uint64_t timeout_ms)
{
	ASSERT(client && out_spec && out_payload && out_payload_length);

	if (!client->is_configured || !client->is_subscriber) {
		LOG_WARN("AAP2Client: Invalid client configuration, cannot receive bundles!");
		return AAP2_CLIENT_EARGUMENT;
	}

	aap2_AAPMessage recv_msg;

	if (timeout_ms != 0 && timeout_ms <= INT32_MAX) {
		switch (poll_recv_timeout(client->socket_fd, timeout_ms)) {
		case 0:
			// timeout
			LOG_DEBUG("AAP2Client: Timed out.");
			return AAP2_CLIENT_ETIMEOUT;
		case 1:
			break;
		default:
			// general error
			LOG_WARN("AAP2Client: Error while polling on socket for receiving data!");
			return AAP2_CLIENT_EFAIL;
		}
	}

	if (aap2_recv_message(&client->pb_istream, aap2_AAPMessage_fields, &recv_msg)) {
		LOG_WARN("AAP2Client: Failed to receive AAP2 message!");
		return AAP2_CLIENT_EFAIL;
	}

	LOGF_DEBUG("AAP2Client: Got a message back. Type is: %d", recv_msg.which_msg);

	aap2_AAPResponse response = aap2_AAPResponse_init_default;

	// If it is not an ADU message (e.g. Keepalive), just ACK that we received it.
	// The caller should retry afterwards if they still want to get an ADU.
	if (recv_msg.which_msg != aap2_AAPMessage_adu_tag) {
		pb_release(aap2_AAPMessage_fields, &recv_msg);
		response.response_status = aap2_ResponseStatus_RESPONSE_STATUS_ACK;
		if (aap2_send_message(client->socket_fd, aap2_AAPResponse_fields, &response) < 0) {
			LOG_WARN("AAP2Client: Cannot send response!");
			return AAP2_CLIENT_EFAIL;
		}
		return AAP2_CLIENT_EAGAIN;
	}

	// Read payload
	uint8_t *const payload_data = receive_payload(
		&client->pb_istream,
		recv_msg.msg.adu.payload_length
	);

	if (!payload_data) {
		pb_release(aap2_AAPMessage_fields, &recv_msg);
		return AAP2_CLIENT_EFAIL;
	}

	*out_payload = payload_data;
	*out_payload_length = recv_msg.msg.adu.payload_length;
	out_spec->source_eid = recv_msg.msg.adu.src_eid;
	recv_msg.msg.adu.src_eid = NULL;
	out_spec->dest_eid = recv_msg.msg.adu.dst_eid;
	recv_msg.msg.adu.dst_eid = NULL;
	out_spec->report_to_eid = recv_msg.msg.adu.report_to_eid;
	recv_msg.msg.adu.report_to_eid = NULL;
	out_spec->creation_timestamp_dtn_ms = recv_msg.msg.adu.creation_timestamp_ms;
	out_spec->sequence_number = recv_msg.msg.adu.sequence_number;
	out_spec->flags = aap2_get_adu_flags(
		recv_msg.msg.adu.adu_flags,
		recv_msg.msg.adu.adu_flags_count
	);
	pb_release(aap2_AAPMessage_fields, &recv_msg);

	response.response_status = aap2_ResponseStatus_RESPONSE_STATUS_SUCCESS;
	if (aap2_send_message(client->socket_fd, aap2_AAPResponse_fields, &response) < 0) {
		LOG_WARN("AAP2Client: Cannot send response!");
		// Do not leak memory if the action failed.
		aap2_client_free_bundle_spec_fields(out_spec);
		return AAP2_CLIENT_EFAIL;
	}

	return AAP2_CLIENT_OK;
}

enum aap2_client_error aap2_client_keepalive(struct aap2_client *client, uint64_t timeout_ms)
{
	ASSERT(client);

	if (!client->is_configured || client->is_subscriber) {
		LOG_WARN("AAP2Client: Invalid client configuration, cannot send Keepalive!");
		return AAP2_CLIENT_EARGUMENT;
	}

	LOG_DEBUG("AAP2Client: Sending Keepalive message.");

	aap2_AAPMessage send_msg = aap2_AAPMessage_init_default;

	send_msg.which_msg = aap2_AAPMessage_keepalive_tag;

	if (aap2_send_message(client->socket_fd, aap2_AAPMessage_fields, &send_msg) < 0) {
		LOG_WARN("AAP2Client: Cannot send message!");
		return AAP2_CLIENT_EFAIL;
	}

	pb_release(aap2_AAPMessage_fields, &send_msg);

	if (timeout_ms != 0 && timeout_ms <= INT32_MAX) {
		switch (poll_recv_timeout(client->socket_fd, timeout_ms)) {
		case 0:
			// timeout
			LOG_DEBUG("AAP2Client: Timed out.");
			return AAP2_CLIENT_ETIMEOUT;
		case 1:
			break;
		default:
			// general error
			LOG_WARN("AAP2Client: Error while polling on socket for receiving data!");
			return AAP2_CLIENT_EFAIL;
		}
	}

	aap2_AAPResponse response;

	if (aap2_recv_message(&client->pb_istream, aap2_AAPResponse_fields, &response)) {
		LOG_WARN("AAP2Client: Error receiving AAP2 response!");
		return AAP2_CLIENT_EFAIL;
	}

	if (response.response_status != aap2_ResponseStatus_RESPONSE_STATUS_ACK &&
	    response.response_status != aap2_ResponseStatus_RESPONSE_STATUS_SUCCESS) {
		LOGF_WARN(
			"AAP2Client: Invalid response to Keepalive message; status = %d",
			response.response_status
		);
		pb_release(aap2_AAPResponse_fields, &response);
		return AAP2_CLIENT_EFAIL;
	}

	LOG_DEBUG("AAP2Client: Keepalive message acknowledged by server.");
	pb_release(aap2_AAPResponse_fields, &response);

	return AAP2_CLIENT_OK;
}

void aap2_client_free_bundle_spec_fields(struct aap2_client_bundle_spec *spec)
{
	free(spec->dest_eid);
	free(spec->report_to_eid);
	free(spec->source_eid);
}