#include <stdbool.h>
#include <string.h>
#include <errno.h>
#include <unistd.h>
#include <assert.h>
#include <arpa/inet.h>
#include <linux/tcp.h>
#include <net/if.h>

#include <tfe_utils.h>
#include <tfe_tcp_restore.h>

void tfe_tcp_restore_info_dump(const struct tcp_restore_info *info)
{
	char str_client_addr[64] = { 0 };
	char str_server_addr[64] = { 0 };

	const struct tcp_restore_endpoint *client = &info->client;
	const struct tcp_restore_endpoint *server = &info->server;

	assert(client->addr.ss_family == server->addr.ss_family);

	if (client->addr.ss_family == AF_INET)
	{
		struct sockaddr_in *sk_client = (struct sockaddr_in *)&client->addr;
		struct sockaddr_in *sk_server = (struct sockaddr_in *)&server->addr;
		uint16_t port_client = ntohs(sk_client->sin_port);
		uint16_t port_server = ntohs(sk_server->sin_port);

		inet_ntop(AF_INET, &sk_client->sin_addr, str_client_addr, sizeof(str_client_addr));
		inet_ntop(AF_INET, &sk_server->sin_addr, str_server_addr, sizeof(str_client_addr));

		TFE_LOG_DEBUG(g_default_logger, "tcp_restore_info %p: cur_dir=%u, %s:%hu->%s:%hu, seq=%u, ack=%u, "
				"client={ mss=%u, wscale_perm=%u, wscale=%u, ts=%u, sack=%u }, "
				"server={ mss=%u, wscale_perm=%u, wscale=%u, ts=%u, sack=%u }",
				info, info->cur_dir, str_client_addr, port_client, str_server_addr, port_server, info->client.seq, info->client.ack,
				client->mss, (client->wscale_perm ? 1 : 0), client->wscale, (client->timestamp_perm ? 1 : 0), (client->sack_perm ? 1 : 0),
				server->mss, (server->wscale_perm ? 1 : 0), server->wscale, (server->timestamp_perm ? 1 : 0), (server->sack_perm ? 1 : 0));
	}
	else if (client->addr.ss_family == AF_INET6)
	{
		struct sockaddr_in6 *sk_client = (struct sockaddr_in6 *)&client->addr;
		struct sockaddr_in6 *sk_server = (struct sockaddr_in6 *)&server->addr;
		uint16_t port_client = ntohs(sk_client->sin6_port);
		uint16_t port_server = ntohs(sk_server->sin6_port);

		inet_ntop(AF_INET6, &sk_client->sin6_addr, str_client_addr, sizeof(str_client_addr));
		inet_ntop(AF_INET6, &sk_server->sin6_addr, str_server_addr, sizeof(str_client_addr));

		TFE_LOG_DEBUG(g_default_logger, "tcp_restore_info %p: cur_dir=%u, %s:%hu->%s:%hu, seq=%u, ack=%u, "
				"client={ mss=%u, wscale_perm=%u, wscale=%u, ts=%u, sack=%u }, "
				"server={ mss=%u, wscale_perm=%u, wscale=%u, ts=%u, sack=%u }",
				info, info->cur_dir, str_client_addr, port_client, str_server_addr, port_server, info->client.seq, info->client.ack,
				client->mss, (client->wscale_perm ? 1 : 0), client->wscale, (client->timestamp_perm ? 1 : 0), (client->sack_perm ? 1 : 0),
				server->mss, (server->wscale_perm ? 1 : 0), server->wscale, (server->timestamp_perm ? 1 : 0), (server->sack_perm ? 1 : 0));
	}
}

int tfe_tcp_restore_fd_create(const struct tcp_restore_endpoint *endpoint, const struct tcp_restore_endpoint *peer, const char *devname, unsigned int fd_so_mask)
{
	int result = 0;
	int sockopt = 0;
	int sockfd = 0;
	char buffer[IFNAMSIZ] = {0};
	socklen_t buffer_len = sizeof(buffer);
	unsigned int nr_tcp_repair_opts = 0;
	struct tcp_repair_opt tcp_repair_opts[8];
	struct tcp_repair_window tcp_repair_window = { 0 };

	if (endpoint->addr.ss_family == AF_INET)
	{
		sockfd = socket(AF_INET, SOCK_STREAM | SOCK_CLOEXEC, IPPROTO_TCP);
	}
	else if (endpoint->addr.ss_family == AF_INET6)
	{
		sockfd = socket(AF_INET6, SOCK_STREAM | SOCK_CLOEXEC, IPPROTO_TCP);
	}
	else
	{
		errno = EINVAL;
		TFE_LOG_ERROR(g_default_logger, "failed at tcp_restore_fd_create(), %d: %s", errno, strerror(errno));
		goto errout;
	}

	if (sockfd < 0)
	{
		TFE_LOG_ERROR(g_default_logger, "failed at socket(), %d: %s", errno, strerror(errno));
		goto errout;
	}

	sockopt = fd_so_mask;
	result = setsockopt(sockfd, SOL_SOCKET, SO_MARK, (char *)&sockopt, sizeof(sockopt));
	if (result < 0)
	{
		TFE_LOG_ERROR(g_default_logger, "failed at setsockopt(SO_MARK), %d: %s", errno, strerror(errno));
		goto errout;
	}
	
	if (strlen(devname))
	{
		result = setsockopt(sockfd, SOL_SOCKET, SO_BINDTODEVICE, devname, strlen(devname));
		if (result < 0)
		{
			TFE_LOG_ERROR(g_default_logger, "failed at setsockopt(SO_BINDTODEVICE) on %d, %d: %s", devname, errno, strerror(errno));
			goto errout;
		}

		result = getsockopt(sockfd, SOL_SOCKET, SO_BINDTODEVICE, buffer, &buffer_len);
		if (result < 0)
		{
			TFE_LOG_ERROR(g_default_logger, "failed at getsockopt(SO_BINDTODEVICE) on %d, %d: %s", devname, errno, strerror(errno));
			goto errout;
		}
		TFE_LOG_DEBUG(g_default_logger, "sockfd %d successfully bound to %s device, so_mask: %x", sockfd, buffer, fd_so_mask);
	}

	// Setup TCP REPAIR Status
	sockopt = 1;
	result = setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, (char *)&sockopt, sizeof(sockopt));
	if (result < 0)
	{
		TFE_LOG_ERROR(g_default_logger, "failed at setsockopt(SO_REUSEADDR), %d: %s", errno, strerror(errno));
		goto errout;
	}

	sockopt = 1;
	result = setsockopt(sockfd, SOL_IP, IP_TRANSPARENT, (char *)&sockopt, sizeof(sockopt));
	if (result < 0)
	{
		TFE_LOG_ERROR(g_default_logger, "failed at setsockopt(IP_TRANSPARENT), %d: %s", errno, strerror(errno));
		goto errout;
	}

	sockopt = 1;
	result = setsockopt(sockfd, IPPROTO_TCP, TCP_REPAIR, (char *)&sockopt, sizeof(sockopt));
	if (result < 0)
	{
		TFE_LOG_ERROR(g_default_logger, "failed at setsockopt(TCP_REPAIR), %d: %s", errno, strerror(errno));
		goto errout;
	}

	// Setup SEQ/ACK and TCP options
	sockopt = TCP_SEND_QUEUE;
	result = setsockopt(sockfd, IPPROTO_TCP, TCP_REPAIR_QUEUE, (char *)&sockopt, sizeof(sockopt));
	if (result < 0)
	{
		TFE_LOG_ERROR(g_default_logger, "failed at setsockopt(TCP_REPAIR_QUEUE), %d: %s", errno, strerror(errno));
		goto errout;
	}

	sockopt = endpoint->seq;
	result = setsockopt(sockfd, IPPROTO_TCP, TCP_QUEUE_SEQ, (char *)&sockopt, sizeof(sockopt));
	if (result < 0)
	{
		TFE_LOG_ERROR(g_default_logger, "failed at setsockopt(TCP_QUEUE_SEQ), %d: %s", errno, strerror(errno));
		goto errout;
	}

	sockopt = TCP_RECV_QUEUE;
	result = setsockopt(sockfd, IPPROTO_TCP, TCP_REPAIR_QUEUE, (char *)&sockopt, sizeof(sockopt));
	if (result < 0)
	{
		TFE_LOG_ERROR(g_default_logger, "failed at setsockopt(TCP_REPAIR_QUEUE), %d: %s", errno, strerror(errno));
		goto errout;
	}

	sockopt = endpoint->ack;
	result = setsockopt(sockfd, IPPROTO_TCP, TCP_QUEUE_SEQ, (char *)&sockopt, sizeof(sockopt));
	if (result < 0)
	{
		TFE_LOG_ERROR(g_default_logger, "failed at setsockopt(TCP_QUEUE_SEQ), %d: %s", errno, strerror(errno));
		goto errout;
	}

#ifndef TCPOPT_MAXSEG
#define TCPOPT_MAXSEG 2
#endif

#ifndef TCPOPT_WINDOW
#define TCPOPT_WINDOW 3
#endif

#ifndef TCPOPT_SACK_PERMITTED
#define TCPOPT_SACK_PERMITTED 4
#endif

#ifndef TCPOPT_TIMESTAMP
#define TCPOPT_TIMESTAMP 8
#endif

	tcp_repair_opts[nr_tcp_repair_opts].opt_code = TCPOPT_MAXSEG;
	tcp_repair_opts[nr_tcp_repair_opts].opt_val = MIN(endpoint->mss, peer->mss);
	nr_tcp_repair_opts++;

	if (endpoint->sack_perm && peer->sack_perm)
	{
		tcp_repair_opts[nr_tcp_repair_opts].opt_code = TCPOPT_SACK_PERMITTED;
		tcp_repair_opts[nr_tcp_repair_opts].opt_val = 0;
		nr_tcp_repair_opts++;
	}

	if (endpoint->wscale_perm && peer->wscale_perm)
	{
		tcp_repair_opts[nr_tcp_repair_opts].opt_code = TCPOPT_WINDOW;
		tcp_repair_opts[nr_tcp_repair_opts].opt_val = (endpoint->wscale << 16) | peer->wscale;
		nr_tcp_repair_opts++;
	}

	if (endpoint->timestamp_perm && peer->timestamp_perm)
	{
		tcp_repair_opts[nr_tcp_repair_opts].opt_code = TCPOPT_TIMESTAMP;
		tcp_repair_opts[nr_tcp_repair_opts].opt_val = 0;
		nr_tcp_repair_opts++;
	}

	// Bind address and connect to peer endpoint
	result = bind(sockfd, (struct sockaddr *)&endpoint->addr, sizeof(endpoint->addr));
	if (result < 0)
	{
		TFE_LOG_ERROR(g_default_logger, "failed at bind(), %d: %s", errno, strerror(errno));
		goto errout;
	}

	result = connect(sockfd, (struct sockaddr *)&peer->addr, sizeof(peer->addr));
	if (result < 0)
	{
		TFE_LOG_ERROR(g_default_logger, "failed at connect(), %d: %s", errno, strerror(errno));
		goto errout;
	}

	result = setsockopt(sockfd, IPPROTO_TCP, TCP_REPAIR_OPTIONS, (char *)tcp_repair_opts, nr_tcp_repair_opts * sizeof(struct tcp_repair_opt));
	if (result < 0)
	{
		TFE_LOG_ERROR(g_default_logger, "failed at setsockopt(TCP_REPAIR_OPTIONS), %d: %s", errno, strerror(errno));
		goto errout;
	}

	if (endpoint->timestamp_perm && peer->timestamp_perm)
	{
		result = setsockopt(sockfd, IPPROTO_TCP, TCP_TIMESTAMP, &(endpoint->ts_val), sizeof(endpoint->ts_val));
		if (result < 0)
		{
			TFE_LOG_ERROR(g_default_logger, "failed at setsockopt(TCP_TIMESTAMP), %d: %s", errno, strerror(errno));
			goto errout;
		}
	}

	// Perpare Window Setup
	tcp_repair_window.snd_wl1 = peer->seq;
	tcp_repair_window.snd_wnd = peer->window;
	tcp_repair_window.max_window = peer->window;
	tcp_repair_window.rcv_wnd = endpoint->window;
	tcp_repair_window.rcv_wup = endpoint->ack;

	result = setsockopt(sockfd, IPPROTO_TCP, TCP_REPAIR_WINDOW, (char *)&tcp_repair_window, sizeof(tcp_repair_window));
	if (result < 0)
	{
		TFE_LOG_ERROR(g_default_logger, "failed at setsockopt(TCP_REPAIR_WINDOW), %d: %s", errno, strerror(errno));
		goto errout;
	}

	sockopt = 0;
	result = setsockopt(sockfd, IPPROTO_TCP, TCP_REPAIR, (char *)&sockopt, sizeof(sockopt));
	if (result < 0)
	{
		TFE_LOG_ERROR(g_default_logger, "failed at setsockopt(TCP_REPAIR), %d: %s", errno, strerror(errno));
		goto errout;
	}

	return sockfd;

errout:
	if (sockfd > 0)
	{
		close(sockfd);
	}

	return -1;
}