#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netdb.h>
#include <string.h>
#include <errno.h>
#include "net.h"
#include "fd.h"
#include "util.h"

int getport(const char *s, unsigned short *svalp)
{
	struct servent *serv;
	if (!getushort(s, svalp)) return 0;
	if (errno!=EINVAL) return -1;
	serv=getservbyname(s, "tcp");
	if (!serv)
	{
		errno=EINVAL;
		return -1;
	}
	*svalp=ntohs(serv->s_port);
	return 0;
}

void lookup(struct sockaddr_in *s, const char *hostname, unsigned short port,
	char **addrs, int *addrlenp, int *addrnump)
{
	struct hostent *host;
	int i, max, len;
	char *data;

	s->sin_family=AF_INET;
	s->sin_port=htons(port);
	if (!hostname)
	{
		*addrs=NULL;
		*addrlenp=0;
		*addrnump=0;
		return;
	}

	host=gethostbyname(hostname);
	if (!host) die("lookup %s: %s\n", hostname, hstrerror(h_errno));
	for (i=0; host->h_addr_list[i]; i++);
	max=i;
	len=host->h_length;
	data=malloc(len*max);
	for (i=0; i<max; i++)
		memcpy(&data[i*len], host->h_addr_list[i], len);
	*addrs=data;
	*addrlenp=len;
	*addrnump=max;
}

fd_t *connection(const char *shost, unsigned short sport,
	const char *dhost, unsigned short dport)
{
	struct sockaddr_in src, dst;
	char *srcaddrs, *dstaddrs;
	int srclen, dstlen;
	int srcnum, dstnum;
	int i;

	memset(&src, 0, sizeof(struct sockaddr_in));
	memset(&dst, 0, sizeof(struct sockaddr_in));

	lookup(&src, shost, sport, &srcaddrs, &srclen, &srcnum);
	lookup(&dst, dhost, dport, &dstaddrs, &dstlen, &dstnum);

	src.sin_family=AF_INET;
	if (shost)
	{
		if (!srcnum) die("No address found for source hostname\n");
		if (srclen!=sizeof(in_addr_t)) 
			die("BUG! Length mismatch %i vs %li\n",
				srclen, (long)sizeof(in_addr_t));
		memcpy(&src.sin_addr, srcaddrs, sizeof(in_addr_t));
	}
	else
	{
		src.sin_addr.s_addr=INADDR_ANY;
	}

	if (!dhost) 
		die("Desination hostname must be specified when "
			"initiating a connection\n");
	if (!dstnum) die("No address found for destination hostname\n");
	if (!dport) 
		die("Desination port must be specified when "
			"initiating a connection\n");
	if (dstlen!=sizeof(in_addr_t)) 
		die("BUG! Length mismatch %i vs %li\n", 
			dstlen, (long)sizeof(in_addr_t));
	for (i=0; i<dstnum; i++)
	{
		int sock;
		memcpy(&dst.sin_addr, &dstaddrs[i*sizeof(in_addr_t)], sizeof(in_addr_t));

		if ((sock=socket(PF_INET, SOCK_STREAM, IPPROTO_TCP))==-1)
			die("socket: %s\n", strerror(errno));
		if (shost||sport) 
		{
			if (bind(sock, (struct sockaddr *)&src, sizeof(src)))
				die("bind: %s\n", strerror(errno));
		}
		if (!connect(sock, (struct sockaddr *)&dst, sizeof(dst)))
		{
			fd_t *fd;
			fd=fd_alloc();
			fd_init(fd, sock, 1, 1, "network connection");
			return fd;
		}
		warn("connect: %s\n", strerror(errno));
		close(sock);
	}
	if (sport)
		die("Couldn't establish connection from %s:%i to %s:%i\n", 
			shost? shost:"*", (int)sport, dhost, (int)dport);
	die("Couldn't establish connection from %s:* to %s:%i\n", 
		shost? shost:"*", dhost, (int)dport);
	return NULL; /* NOTREACHED */
}

connection_t *connection_alloc(void)
{
	return xmalloc(sizeof(connection_t));
}

connection_t *connection_init(connection_t *conn, fd_t *fd, 
	struct sockaddr *src, struct sockaddr *dst)
{
	conn->srcaddr=*src;
	conn->dstaddr=*dst;
	conn->fd=fd;
	conn->tonet=conn->fromnet=NULL;
	return conn;
}

void connection_try(protoconnection_t *protoconn);

void connection_check(void *object, void *arg1 UNUSED, int arg2 UNUSED)
{
	protoconnection_t *protoconn=(protoconnection_t *)object;
	int sock=protoconn->fd->fd_no;
	int err;
	int errlen=sizeof(err);

	if (getsockopt(sock, SOL_SOCKET, SO_ERROR, &err, &errlen))
		die("getsockopt: %s\n", strerror(errno));
	if (err==0)
	{
		connection_t *conn=connection_alloc();
		connection_init(conn, protoconn->fd, 
			&protoconn->srcaddr, 
			&protoconn->dstaddr[protoconn->thisdst]);
		(*protoconn->notify_connected)(conn);
		fd_unrequest_write(&protoconn->connection_ready);
		fd_del_writer(protoconn->fd, &protoconn->fd_write);
		free(protoconn);
		return;
	}
	if (err==EALREADY) return;
	warn("getsockopt (error while connecting): %s\n", strerror(err));
	if (++protoconn->thisdst==protoconn->dstaddrs)
	{
		die("Couldn't establish connection\n");
		/* XXX */
	}
	fd_unrequest_write(&protoconn->connection_ready);
	fd_del_writer(protoconn->fd, &protoconn->fd_write);
	fd_close(protoconn->fd);
	free(protoconn->fd);
	connection_try(protoconn);
}

void connection_try(protoconnection_t *protoconn)
{
	int sock;
	fd_t *fd;
	struct sockaddr *dstaddr=&protoconn->dstaddr[protoconn->thisdst];

	if ((sock=socket(PF_INET, SOCK_STREAM, IPPROTO_TCP))==-1)
		die("socket: %s\n", strerror(errno));
	fd=fd_alloc();
	fd_init(fd, sock, 1, 1, "network connection");
	fd_set_nonblock(fd);
	if (protoconn->dobind) 
	{
		if (bind(sock, (struct sockaddr *)&protoconn->srcaddr, 
			sizeof(struct sockaddr_in)))
			die("bind: %s\n", strerror(errno));
	}
	if (!connect(sock, (struct sockaddr *)dstaddr,
		sizeof(struct sockaddr_in)))
	{
		connection_t *conn=connection_alloc();
		connection_init(conn, protoconn->fd, 
			&protoconn->srcaddr, 
			&protoconn->dstaddr[protoconn->thisdst]);
		(*protoconn->notify_connected)(conn);
		free(protoconn);
		return;
	}
	if (errno!=EINPROGRESS)
	{
		warn("connect: %s\n", strerror(errno));
		if (++protoconn->thisdst==protoconn->dstaddrs)
		{
			die("Couldn't establish connection\n");
			/* XXX */
		}
	}
	notifier_init(&protoconn->connection_ready, protoconn, connection_check);
	fd_add_writer(fd, &protoconn->fd_write);
	fd_request_write(fd, &protoconn->connection_ready);
	protoconn->fd=fd;
}

void connection_create(const char *shost, unsigned short sport, 
	const char *dhost, unsigned short dport,
	void (*notify_connected)(connection_t *connection))
{
	struct sockaddr_in src, dst; /* XXX Remove sock from lookup */
	char *srcaddrs, *dstaddrs;
	int srclen, dstlen;
	int srcnum, dstnum;
	int i;
	struct sockaddr_in *sockdst;
	protoconnection_t *protoconn;

	memset(&src, 0, sizeof(struct sockaddr_in));
	memset(&dst, 0, sizeof(struct sockaddr_in));

	lookup(&src, shost, sport, &srcaddrs, &srclen, &srcnum);
	lookup(&dst, dhost, dport, &dstaddrs, &dstlen, &dstnum);

	src.sin_family=AF_INET;
	if (shost)
	{
		if (!srcnum) die("No address found for source hostname\n");
		if (srclen!=sizeof(in_addr_t)) 
			die("BUG! Length mismatch %i vs %li\n",
				srclen, (long)sizeof(in_addr_t));
		memcpy(&src.sin_addr, srcaddrs, sizeof(in_addr_t));
	}
	else
		src.sin_addr.s_addr=INADDR_ANY;

	if (!dhost) 
		die("Desination hostname must be specified when "
			"initiating a connection\n");
	if (!dstnum) die("No address found for destination hostname\n");
	if (!dport) 
		die("Desination port must be specified when "
			"initiating a connection\n");
	if (dstlen!=sizeof(in_addr_t)) 
		die("BUG! Length mismatch %i vs %li\n", 
			dstlen, (long)sizeof(in_addr_t));

	sockdst=xmalloc(sizeof(struct sockaddr_in)*dstnum);
	for (i=0; i<dstnum; i++)
	{
		
		sockdst[i].sin_family=AF_INET;
		memcpy(&sockdst[i].sin_addr, &dstaddrs[i*dstlen], 
			sizeof(in_addr_t));
		sockdst[i].sin_port=htons(dport);
	}

	protoconn=xmalloc(sizeof(protoconnection_t));
	memcpy(&protoconn->srcaddr, &src, sizeof(struct sockaddr_in));
	protoconn->dstaddr=(struct sockaddr *)sockdst;
	protoconn->dstaddrs=dstnum;
	protoconn->thisdst=0;
	protoconn->dobind=shost||sport;
	protoconn->notify_connected=notify_connected;
	free(srcaddrs);
	free(dstaddrs);
	connection_try(protoconn);
}

listener_t *listener_alloc(void)
{
	return xmalloc(sizeof(listener_t));
}

void listener_accept(void *object, void *arg1 UNUSED, int arg2 UNUSED)
{
	listener_t *listener=(listener_t *)object;
	struct sockaddr addr;
	socklen_t addrlen=sizeof(addr);
	fd_t *fd;
	connection_t *conn;

	if (fd_accept(&fd, listener->fd, &addr, &addrlen, 
		"network connection"))
		return;
	conn=connection_alloc();
	connection_init(conn, fd, &addr, &listener->boundaddr);
	if (listener->notify_accepted)
		(*listener->notify_accepted)(listener, conn);
}

void listener_misc(void *object, void *arg1 UNUSED, int arg2)
{
	listener_t *listener=(listener_t *)object;

	if (arg2==FD_CLOSE)
	{
		if (listener->readok.list.next==&listener->readok.list)
			fd_unrequest_read(&listener->readok);
	}
}

listener_t *listener_init(listener_t *listener, fd_t *fd, 
	struct sockaddr *addr, 
	void (*notify_accepted)(listener_t *listener, 
		connection_t *connection))
{
	listener->fd=fd;
	listener->boundaddr=*addr;
       	notifier_init(&listener->readok, listener, listener_accept);
	notifier_init(&listener->miscevent, listener, listener_misc);
	listener->notify_accepted=notify_accepted;
	return listener;
}

void listener_create(const char *shost, unsigned short sport,
	void (*notify_accepted)(struct listener_s *listener, 
		connection_t *connection),
	listener_t ***listenersp, int *listenernump)
{
	struct sockaddr_in src;
	char *srcaddrs;
	int srclen;
	int srcnum;
	in_addr_t any=INADDR_ANY;
	listener_t **listeners;
	int num=0, i;

	memset(&src, 0, sizeof(struct sockaddr_in));
	if (!sport) 
		die("Source port must be specified when "
			"listening for a connection\n");

	lookup(&src, shost, sport, &srcaddrs, &srclen, &srcnum);

	src.sin_family=AF_INET;
	if (shost)
	{
		if (!srcnum) die("No address found for source hostname\n");
		if (srclen!=sizeof(in_addr_t)) 
			die_hard("BUG! Length mismatch %i vs %li\n",
				srclen, (long)sizeof(in_addr_t));
		memcpy(&src.sin_addr, srcaddrs, sizeof(in_addr_t));
	}
	else
	{
		srcaddrs=(char *)&any;
		srclen=sizeof(any);
		srcnum=1;
	}

	listeners=xmalloc(sizeof(listener_t *)*srcnum);
	for (i=0; i<srcnum; i++)
	{
		fd_t *fd;
		int sock;
		if ((sock=socket(PF_INET, SOCK_STREAM, IPPROTO_TCP))==-1)
			die("socket: %s\n", strerror(errno));
		if (bind(sock, (struct sockaddr *)&src, sizeof(src)))
		{
			warn("bind: %s\n", strerror(errno));
			close(sock);
			continue;
		}
		if (listen(sock, 5))
		{
			warn("listen: %s\n", strerror(errno));
			close(sock);
			continue;
		}
		fd=fd_alloc();
		fd_init(fd, sock, 1, 0, "network listener");
		fd_set_nonblock(fd);
		listeners[num]=listener_alloc();
		listener_init(listeners[num], fd, 
			(struct sockaddr *)&src, notify_accepted);
		fd_add_reader(fd, &listeners[num]->fd_list);
		fd_request_read(fd, &listeners[num]->readok);
		fd_request_misc(fd, &listeners[num]->miscevent);
		num++;
	}
	if (num==0)
	{
		warn("Couldn't listen on %s:%i\n", 
			shost? shost:"*", (int)sport);
		free(listeners);
		listeners=NULL;
	}
	*listenersp=listeners;
	*listenernump=num;
}