#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netdb.h>
#include <unistd.h>
#include <pthread.h>

struct workerinfo {
	pthread_t thr;
	const struct addrinfo *ai;
};

static struct addrinfo *head;
static unsigned nr_workers;
static struct workerinfo *workers;
static int keep_going;
static pthread_cond_t work_cond;
static pthread_mutex_t work_mutex;

static void usage(void)
{
	fprintf(stderr, "usage: connect-monster <host/port>...\n");
	exit(EXIT_FAILURE);
}

static int try_connect(const struct addrinfo *ai)
{
	int fd;

	fd = socket(ai->ai_family, ai->ai_socktype, ai->ai_protocol);
	if (fd < 0) {
		perror("socket()");
		exit(EXIT_FAILURE); // TODO: not nice!
	}

	if (connect(fd, ai->ai_addr, ai->ai_addrlen)) {
		perror("connect()");
		exit(EXIT_FAILURE); // TODO: not nice!
	}

	close(fd);
	return 1;
}

static void add(const struct addrinfo *ai)
{
	struct addrinfo *new = malloc(sizeof(*new));
	char host[NI_MAXHOST];
	char serv[NI_MAXSERV];
	int err;

	*new = *ai;
	new->ai_next = head;
	if (new->ai_canonname)
		new->ai_canonname = strdup(new->ai_canonname);
	head = new;

	err = getnameinfo(new->ai_addr, new->ai_addrlen, host, sizeof(host), serv, sizeof(serv), NI_NUMERICHOST | NI_NUMERICSERV);
	if (err) {
		fprintf(stderr, "ERROR:getnameinfo():%s\n", gai_strerror(err));
		exit(EXIT_FAILURE);
	}
	fprintf(stderr,"Adding %s/%s (family=%d socktype=%d proto=%d)\n",
		host, serv,
		ai->ai_family, ai->ai_socktype, ai->ai_protocol);
}

static void find(const char *arg)
{
	char *port = strrchr(arg, '/');
	char *node;
	struct addrinfo *res;
	struct addrinfo *curr;
	struct addrinfo hints;
	int err;

	if (!port)
		port = strrchr(arg, ':');
	if (!port)
		usage();
	node = strndup(arg, port - arg);
	++port;

	/* setup some hints */
	memset(&hints, 0, sizeof(hints));
	hints.ai_next = NULL;
	hints.ai_family = AF_INET; /* or AF_UNSPEC or AF_INET6 */
	hints.ai_socktype = SOCK_STREAM;
	hints.ai_protocol = IPPROTO_TCP;
	hints.ai_flags = AI_CANONNAME;

	err = getaddrinfo(node, port, &hints, &res);
	free(node);
	if (err)
	{
		fprintf(stderr, "ERROR:getaddrinfo():%s:%s\n", arg, gai_strerror(err));
		exit(EXIT_FAILURE);
	}
	for (curr = res; curr != NULL; curr = curr->ai_next)
		add(curr);
	freeaddrinfo(res);
}

static unsigned count_ai(const struct addrinfo *ai)
{
	const struct addrinfo *curr;
	unsigned total = 0;

	for (curr = ai; curr != NULL; curr = curr->ai_next)
		total++;
	return total;
}

static void *worker(void *p)
{
	struct workerinfo *info = p;
	intptr_t count = 0;

	/* hold until start_workers() is called */
	pthread_mutex_lock(&work_mutex);
	pthread_cond_wait(&work_cond, &work_mutex);
	pthread_mutex_unlock(&work_mutex);

	while (keep_going) {
		try_connect(info->ai);
		count++;
	}

	return (void*)count;
}

static void create_workers(void)
{
	unsigned i;
	struct addrinfo *curr;

	for (curr = head, i = 0; i < nr_workers; i++, curr = curr->ai_next) {
		workers[i].ai = curr;
		pthread_create(&workers[i].thr, NULL, worker, &workers[i]);
	}
}

static void start_workers(void)
{
	keep_going = 1;

	pthread_cond_broadcast(&work_cond);
}

static void stop_workers(void)
{
	unsigned i;
	unsigned count = 0;

	keep_going = 0;

	for (i = 0; i < nr_workers; i++) {
		void *res;

		pthread_join(workers[i].thr, &res);
		workers[i].ai = NULL;
		count += (intptr_t)res;
	}
	fprintf(stderr, "connected %d times\n", count);
}

int main(int argc, char **argv)
{
	int i;

	for (i = 1; i < argc; i++)
		find(argv[i]);

	pthread_mutex_init(&work_mutex, NULL);
	pthread_cond_init(&work_cond, NULL);

	nr_workers = count_ai(head);
	if (!nr_workers)
		usage();
	workers = calloc(nr_workers, sizeof(*workers));
	create_workers();
	start_workers();
	sleep(10);
	stop_workers();

	return 0;
}