#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <time.h>

#include <netinet/ip.h>

#include "connection.h"
#include "event.h"
#include "list.h"
#include "logging.h"
#include "sockaddr.h"
#include "tcpsocket.h"

struct flood_ctrl {
	int burst;
	int current;
	int inc_per_sec;
	time_t last_inc;
};

struct connection {
	struct list_head list;

	struct sockaddr_in src_addr;
	struct sockaddr_in dst_addr;

	struct event_fd *src_event;
	struct event_fd *dst_event;
	struct event_timeout *timeout;

	struct flood_ctrl src_flood;
	struct flood_ctrl dst_flood;

	time_t login_time;
};

static LIST_HEAD(connection_list);

static struct connection * create_connection()
{
	struct connection *con = malloc(sizeof(struct connection));
	if (con == NULL) {
		log_print(LOG_WARN, "listen_handler(): out of memory");
		return 0;
	}

	memset(con, 0, sizeof(struct connection));

	time_t now = time(NULL);
	con->login_time = now;

	con->src_flood.burst = 256;
	con->src_flood.inc_per_sec = 1;
	con->src_flood.last_inc = now;

	con->dst_flood.burst = 256;
	con->dst_flood.inc_per_sec = 1;
	con->dst_flood.last_inc = now;

	con->src_flood.current = con->src_flood.burst;
	con->dst_flood.current = con->dst_flood.burst;
	return con;
}

static void destroy_connection(struct connection *con)
{
	if (con->timeout != NULL)
		event_remove_timeout(con->timeout);

	int src_fd = event_get_fd(con->src_event);
	if (src_fd != -1) {
		close(src_fd);
		event_remove_fd(con->src_event);
	}

	int dst_fd = event_get_fd(con->dst_event);
	if (dst_fd != -1) {
		close(dst_fd);
		event_remove_fd(con->dst_event);
	}

	free(con);
}

static int forward_throttle_timeout(void *privdata);

static int forward_handler(int fd, void *privdata)
{
	struct connection *con = (struct connection *)privdata;

	int src_fd = event_get_fd(con->src_event);
	int dst_fd = event_get_fd(con->dst_event);

	struct event_fd *event;
	struct flood_ctrl *flood;
	int outfd;

	/* client -> device */
	if (src_fd == fd) {
		flood = &con->dst_flood;
		outfd = dst_fd;
		event = con->src_event;

	/* device -> client */
	} else if (dst_fd == fd) {
		flood = &con->src_flood;
		outfd = src_fd;
		event = con->dst_event;

	/* unknown fd */
	} else {
		list_del(&con->list);
		destroy_connection(con);
		return -1;
	}

	/* increment quota */
	time_t now = time(NULL);
	flood->current += (flood->inc_per_sec * (now - flood->last_inc));
	flood->last_inc = now;

	/* max burst size */
	if (flood->current > flood->burst)
		flood->current = flood->burst;

	/* read max. buffer size */
	char buf[256];
	int readsize = sizeof(buf);

	/* only read current quota */
	if (readsize > flood->current)
		readsize = flood->current;

	/* no quota left */
	if (readsize == 0) {
		/* disable fd event */
		event_add_readfd(event, 0, NULL, NULL);

		/* setup timer to reenable fd-event */
		struct timeval tv = { .tv_sec = 1, .tv_usec = 0 };
		event_add_timeout(&tv, forward_throttle_timeout, con);

		return 0;
	}

	int len = read(fd, buf, readsize);
	if (len <= 0 && readsize > 0) {
		list_del(&con->list);
		destroy_connection(con);
		return -1;
	}

	/* not forwarding: discard the data */
	if (outfd == -1)
		return 0;

	write(outfd, buf, len);
	flood->current -= len;
	return 0;
}

static int forward_throttle_timeout(void *privdata)
{
	struct connection *con = (struct connection *)privdata;

	/* HACK: enable both directions unconditionally */
	event_add_readfd(con->dst_event, 0, forward_handler, con);
	event_add_readfd(con->src_event, 0, forward_handler, con);

	/* remove timeout again */
	return -1;
}

static int connect_handler(int fd, void *privdata)
{
	struct connection *con = (struct connection *)privdata;

	if (tcp_connect_error(fd)) {
		list_del(&con->list);
		destroy_connection(con);
		return -1;
	}

	log_print(LOG_INFO, "forwarding to %s", get_sockaddr_buf(&con->dst_addr));

	/* everything from dst will be forwarded */
	event_add_readfd(con->dst_event, 0, forward_handler, con);

	/* remove connect_handler */
	event_add_writefd(con->dst_event, 0, NULL, NULL);

	/* remove timeout */
	event_remove_timeout(con->timeout);
	con->timeout = NULL;
	return 0;
}

static int connect_timeout(void *privdata)
{
	struct connection *con = (struct connection *)privdata;

	log_print(LOG_INFO, "connect to %s timed out", get_sockaddr_buf(&con->dst_addr));

	list_del(&con->list);
	destroy_connection(con);

	/* singleshot timer */
	return -1;
}

int client_handler(int fd, void *privdata)
{
	struct connection *con = (struct connection *)privdata;
	char buf[256];

	int len = read(fd, buf, sizeof(buf));
	if (len <= 0) {
		list_del(&con->list);
		destroy_connection(con);
		return -1;
	}

	if (parse_sockaddr(buf, &con->dst_addr) != 0) {
		list_del(&con->list);
		destroy_connection(con);
		return -1;
	}

	// TODO: check destination

	int dst_fd = tcp_connect_nonblock(&con->dst_addr);
	if (dst_fd < 0) {
		list_del(&con->list);
		destroy_connection(con);
		return -1;
	}

	log_print(LOG_INFO, "connecting to %s", get_sockaddr_buf(&con->dst_addr));

	struct timeval tv;
	tv.tv_sec = 2;
	tv.tv_usec = 0;
	con->timeout = event_add_timeout(&tv, connect_timeout, con);

	/* wait for connect() (non-blocking) */
	con->dst_event = event_add_writefd(NULL, dst_fd, connect_handler, con);

	/* everything from src will be forwarded */
	event_add_readfd(con->src_event, 0, forward_handler, con);

	return 0;
}

int admin_handler(int fd, void *privdata)
{
	struct connection *con = (struct connection *)privdata;
	char buf[256];

	int len = read(fd, buf, sizeof(buf));
	if (len <= 0) {
		list_del(&con->list);
		destroy_connection(con);
		return -1;
	}

	if (strncmp(buf, "list", 4) == 0) {
		struct connection *entry;
		list_for_each_entry(entry, &connection_list, list) {
			len = snprintf(buf, sizeof(buf), "%d,", event_get_fd(entry->src_event));
			len += get_sockaddr(buf + len, sizeof(buf) - len, &entry->src_addr);
			len += snprintf(buf + len, sizeof(buf) - len, ",%d,", event_get_fd(entry->dst_event));
			len += get_sockaddr(buf + len, sizeof(buf) - len, &entry->dst_addr);
			len += snprintf(buf + len, sizeof(buf) - len, ",%ld\n", (time(NULL) - entry->login_time));
			write(fd, buf, len);
		}

	} else if (strncmp(buf, "kill", 4) == 0) {


	} else if (strncmp(buf, "quit", 4) == 0) {
		list_del(&con->list);
		destroy_connection(con);
		return -1;
	}

	return 0;
}

int listen_handler(int fd, void *privdata)
{
	struct connection *con = create_connection();
	if (con == NULL) {
		log_print(LOG_WARN, "listen_handler(): out of memory");
		return 0;
	}

	unsigned int i = sizeof(con->src_addr);
	int src_fd = accept(fd, (struct sockaddr *)&con->src_addr, &i);
	if (src_fd < 0) {
		free(con);
		return 0;
	}

	log_print(LOG_INFO, "accepted connection from %s", get_sockaddr_buf(&con->src_addr));

	list_add_tail(&con->list, &connection_list);

	con->src_event = event_add_readfd(NULL, src_fd, privdata, con);
	return 0;
}