/*
 * Check that IORING_OP_CONNECT works, with and without other side
 * being open.
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <errno.h>
#include <fcntl.h>
#include <unistd.h>
#include <poll.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/tcp.h>

#include "liburing.h"

#include <pthread.h>

static int create_socket(void)
{
	int fd;

	fd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	if (fd == -1) {
		perror("socket()");
		return -1;
	}

	return fd;
}

static int submit_and_wait(struct io_uring *ring, int fd, int *res, int sleep_for)
{
	struct io_uring_cqe *cqe;
	int ret;

	ret = io_uring_submit(ring);
	if (ret != 1) {
		fprintf(stderr, "io_using_submit: got %d\n", ret);
		return 1;
	}
	if (sleep_for) {
	    sleep(sleep_for);
	}
	shutdown(fd, SHUT_RDWR);

	while (!io_uring_cq_ready(ring)) {}
	ret = io_uring_peek_cqe(ring, &cqe);
	if (ret < 0) {
	    printf("peek cqe?\n");
	    return -1;
	}

	*res = cqe->res;
	io_uring_cqe_seen(ring, cqe);
	return 0;
}

static void *
thread_start(void *arg)
{
	int fd = (int)(uint64_t)arg;
	sleep(1);
	shutdown(fd, SHUT_RDWR);
	return NULL;
}

static int pthread_connect_socket(int fd, int *code)
{
	struct sockaddr_in addr;
	int ret;

	memset(&addr, 0, sizeof(addr));
	addr.sin_family = AF_INET;
	addr.sin_port = 0;
	addr.sin_addr.s_addr = 0x10010ac;

	pthread_attr_t attr;
	int s = pthread_attr_init(&attr);
	if (s != 0) {
	    perror("pthread_attr_init");
	    return -1;
	}

	pthread_t thr;
	s = pthread_create(&thr, &attr, &thread_start, (void*)(uint64_t)fd);
	if (s != 0) {
	    perror("pthread_create");
	    return -1;
	}

	ret = connect(fd, (struct sockaddr*)&addr, sizeof(addr));

	return (ret < 0);
}

static int connect_socket(struct io_uring *ring, int fd, int *code, int sleep_for)
{
	struct io_uring_sqe *sqe;
	struct sockaddr_in addr;
	int res;

	memset(&addr, 0, sizeof(addr));
	addr.sin_family = AF_INET;
	addr.sin_port = 0;
	addr.sin_addr.s_addr = 0x10010ac;

	sqe = io_uring_get_sqe(ring);
	if (!sqe) {
		fprintf(stderr, "unable to get sqe\n");
		return -1;
	}

	io_uring_prep_connect(sqe, fd, (struct sockaddr*)&addr, sizeof(addr));
	sqe->user_data = 1;
	io_uring_sqe_set_flags(sqe, IOSQE_ASYNC);

	return submit_and_wait(ring, fd, &res, sleep_for);
}

static int test_connect(struct io_uring *ring, int use_uring, int sleep_for)
{
	int connect_fd;
	int ret, code = 0;

	connect_fd = create_socket();
	if (connect_fd == -1)
		return -1;

	if (use_uring) {
		ret = connect_socket(ring, connect_fd, &code, sleep_for);
	} else {
		ret = pthread_connect_socket(connect_fd, &code);
	}

	if (ret == -1)
		goto err;

	if (code != 0) {
		fprintf(stderr, "connect failed with %d\n", code);
		goto err;
	}

	close(connect_fd);

	return 0;

err:
	close(connect_fd);
	return -1;
}

int main(int argc, char *argv[])
{
	struct io_uring ring;
	int ret;
	int sleep_for;
   
	if (argc == 1) {
	    sleep_for = 0;
	} else {
	    sleep_for = atoi(argv[1]);
	}

	ret = io_uring_queue_init(8, &ring, 0);
	if (ret == -1) {
		perror("io_uring_queue_setup()");
		return 1;
	}

	ret = test_connect(&ring, 0, sleep_for);
	if (ret == 0) {
	    printf("shutting down a socket trying to connect works with pthread\n");
	} else {
	    return -1;
	}
	
	ret = test_connect(&ring, 1, sleep_for);
	if (ret == 0) {
	    printf("shutting down a socket trying to connect with uring works too, waited %d s before shutdown\n", sleep_for);
	} else {
	    return -1;
	}

	io_uring_queue_exit(&ring);
	return 0;
}