#define _POSIX_C_SOURCE 1

#include <stdio.h>
#include <unistd.h>
#include <time.h>
#include <errno.h>
#include <string.h>
#include <stdlib.h>
#include <limits.h>
#include <fcntl.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <sys/param.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <libgen.h>
#include <getopt.h>
#include <poll.h>
#include <syslog.h>
#include <signal.h>
#include <inttypes.h>

#include "wstationd.h"
#include "imei.h"


// As of the writing of this daemon, SBD maximum message length
// is hardware limited to 1960 bytes, however, the SBD format
// has a maximum size of 64k per burst. We add three bytes
// for the SBD header itself. This value largely
// dictates the memory cost per connection.
#define BUFFERSIZE 65539

// Timeout for each connection, in seconds
#define TIMEOUT 30


static char *exec_name;

// Should this application daemonize?
static int daemonize = 1;
// Hard limit on number of concurrent connections
static int slotlimit = 128;
// Used to handle signaled shutdowns
static int keep_polling = 1;

struct connection 
{
	// Buffer to store the data stream before the dump to file
	unsigned char buff[BUFFERSIZE];
	ssize_t len; // How many bytes have been written to the buffer
	struct timeval tv; // Time connection was made
	struct in_addr addr; // Address connection was made from
};


static int write_connection(struct connection *, int strip);

void shutdown_handler(int sig)
{
	keep_polling = 0;
}

static void print_help()
{
	fprintf(stdout, "usage: %s [-d] [-l limit] [-p port] [-b addr] directory\n", exec_name);
	fprintf(stdout, "  -d, --nodaemon         Do not detach and daemonize\n");
	fprintf(stdout, "  -l, --limit <limit>    Set concurrent connection limit (default: 128)\n");
	fprintf(stdout, "  -p, --port <port>      Port to listen on (default: 10800)\n");
	fprintf(stdout, "  -b, --bind <address>   Address to bind to (default: 0.0.0.0)\n");
	fprintf(stdout, "  -i, --imei <imei>      IMEI to allow - specify multiple times for\n");
	fprintf(stdout, "                         additional IMEIs (default: allow all)\n");
	fprintf(stdout, "  -s, --strip            Strip Iridium SBD header when writing files\n");
	fprintf(stdout, "  -h, --help             Print this message and exit\n");
	fprintf(stdout, "  -v, --version          Print version and exit\n");
}


int main(int argc, char* argv[])
{
	const struct option longopts[] = {
		{ "imei", required_argument, NULL, 'i' },
		{ "limit", required_argument, NULL, 'l' },
		{ "nodaemon", required_argument, NULL, 'd' },
		{ "port", required_argument, NULL, 'p' },
		{ "bind", required_argument, NULL, 'b' },
		{ "strip", no_argument, NULL, 's' },
		{ "version", no_argument, NULL, 'v' },
		{ "help", no_argument, NULL, 'h' },
		{ NULL, 0, NULL, 0 }
	};

	int flags = 0;
	int sockfd = 0;
	int port = 10800; // Default port
	uint32_t address = INADDR_ANY; // Default listen address

	struct connection** conns = NULL;
	struct pollfd* fds = NULL;

	imei_set* imeis = NULL;
	uint64_t imei = 0;

	int strip = 0;

	if((exec_name = basename(argv[0])) == NULL){
		fprintf(stderr, "%s: cannot get basename - %s\n",
			exec_name, strerror(errno)
		);
		goto shutdown_error;
	}

	for(int ch; (ch = getopt_long(argc, argv, "i:l:b:p:dsvh0", longopts, NULL)) != -1;){
		switch(ch){
			case 'd':
				daemonize = 0;
			break;

			case 'l':
				slotlimit = (int)strtol(optarg, NULL, 10);
				if(slotlimit < 1 || slotlimit > 1024){
					fprintf(stderr, "%s: invalid limit (must be between 1 and 1024)\n", exec_name);
					print_help();
					goto shutdown_error;
				}
			break;

			case 'i':
				imei = imei_uint64(
					(unsigned char *)argv[optind-1],
					strlen(argv[optind-1])
				);
				if(imei == UINT64_MAX){
					fprintf(stderr, "%s: invalid IMEI \"%s\"\n", exec_name, argv[optind-1]);
					goto shutdown_error;
				}

				if(imeis == NULL){
					imeis = imei_set_new(255);
					if(imeis == NULL){
						fprintf(stderr, "%s: cannot create IMEI set (out of memory?)\n", exec_name);
						goto shutdown_error;
					}
				}

				// TODO: Probably should be some error handling here
				imei_set_add(&imeis, imei);
			break;

			case 's':
				strip = 1;
			break;
			
			case 'p':
				port = (int)strtol(optarg, NULL, 10);
				if(port < 1 || port > 65535){
					fprintf(stderr, "%s: invalid port number\n", exec_name);
					print_help();
					goto shutdown_error;
				}
			break;

			case 'b':
				if(inet_pton(AF_INET, optarg, &address) == -1){
					fprintf(stderr, "%s: invalid bind address\n", exec_name);
					print_help();
					goto shutdown_error;
				}
			break;

			case 'v':
				fprintf(stdout, "wstationd v%s (%s %s)\n", VERSION,__DATE__, __TIME__);
			goto shutdown_clean;

			case '?':
			case 'h':
				print_help();
			goto shutdown_clean;
		}
	}

	// Accept final argument, the directory
	if(argc == (optind+1)){
		// Fail if we can't chdir to that directory
		if(chdir(argv[argc-1]) == -1){
			fprintf(stderr, "%s: cannot change directory - %s\n",
				exec_name, strerror(errno)
			);
			goto shutdown_error;
		}
	} else {
		print_help();
		goto shutdown_error;
	}

	// Handle signals
	signal(SIGINT, shutdown_handler);

	// Daemonize if requested
	if(daemonize){
		pid_t pid = fork();

		// Fork failure
		if(pid == -1){
			fprintf(stderr, "%s: cannot fork - %s\n",
				exec_name, strerror(errno)
			);
			goto shutdown_error;
		}

		// Fork success, exit parent process
		if(pid > 0){ exit(EXIT_SUCCESS); }

		// create new session for daemon
		if(setsid() == -1){
			fprintf(stderr, "%s: cannot setsid - %s\n",
				exec_name, strerror(errno)
			);
			goto shutdown_error;
		}
		
		// Ensure that standard descriptors are
		// unavailable after fork() and setsid()
		int fd = open("/dev/null", O_RDWR, 0);
		if(fd != -1){
			dup2(fd, fileno(stdin));
			dup2(fd, fileno(stdout));
			dup2(fd, fileno(stderr));
		}

		// Open syslog
		openlog(exec_name, LOG_PID, LOG_DAEMON);
	}

	// Change default mask
	umask(0);

	// Establish the socket
	sockfd = socket(AF_INET, SOCK_STREAM, 0);
	if(sockfd == -1){
		if(daemonize){
			syslog(LOG_ERR, "cannot open socket - %s",
				strerror(errno)
			);
		} else {
			fprintf(stderr, "%s: cannot open socket - %s\n",
				exec_name, strerror(errno)
			);
		}
		goto shutdown_error;
	}

	// This sets REUSE on the socket so it's easily reallocated if
	// this program dies
	flags = 1;
	setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, &flags, sizeof(flags));

	struct sockaddr_in sock;
	memset(&sock, 0, sizeof(sock));
	sock.sin_family = AF_INET;
	sock.sin_addr.s_addr = address;
	sock.sin_port = htons(port);

	// Bind our socket to the specified address and port
	if(bind(sockfd, (struct sockaddr*)&sock, sizeof(sock)) < 0){
		if(daemonize){
			syslog(LOG_ERR, "cannot bind - %s",
				strerror(errno)
			);
		} else {
			fprintf(stderr, "%s: cannot bind - %s\n",
				exec_name, strerror(errno)
			);
		}
		goto shutdown_error;
	}

	// Listen on port
	if(listen(sockfd, 10) != 0){
		if(daemonize){
			syslog(LOG_ERR, "cannot listen - %s",
				strerror(errno)
			);
		} else {
			fprintf(stderr, "%s: cannot listen - %s\n",
				exec_name, strerror(errno)
			);
		}
		goto shutdown_error;
	}

	fds = malloc(sizeof(struct pollfd) * (slotlimit + 1));
	conns = malloc(sizeof(struct connection*) * slotlimit);
	if(conns == NULL || fds == NULL){
		if(daemonize){ syslog(LOG_ERR, "out of memory"); }
		else { fprintf(stderr, "%s: out of memory\n", exec_name); }
	}

	memset(conns, 0, sizeof(struct connection*) * slotlimit);
	memset(fds, 0, sizeof(struct pollfd) * (slotlimit + 1));

	// Add sockfd
	fds[0].fd = sockfd;
	fds[0].events = POLLIN;
	fds[0].revents = 0;

	while(keep_polling){
		if(poll(fds, (slotlimit + 1), 1000) == -1){
			// Ignore interrupt generated errors
			if(errno != EINTR){
				if(daemonize){
					syslog(LOG_ERR, "socket poll error - %s",
						strerror(errno)
					);
				} else {
					fprintf(stderr, "%s: socket poll error - %s\n",
						exec_name, strerror(errno)
					);
				}
				goto shutdown_error;
			}
		}

		// Handle existing connections
		int slotfirst = 0, slotfree = 0;
		for(int i = 1; i <= slotlimit; i++){
			// Handle event
			if(fds[i].revents & POLLIN){
				ssize_t sz = recv(
					fds[i].fd, conns[i-1]->buff,
					(BUFFERSIZE - conns[i-1]->len), 0
				);

				if(sz > 0) conns[i-1]->len += sz;
				else {
					if(conns[i-1]->len > 25){
						uint16_t sbd_sz = *(uint16_t*)(&conns[i-1]->buff[1]);
						// Convert and add three bytes for the SBD header
						ssize_t expected = (ssize_t)htons(sbd_sz) + 3;

						int ok = 1;
						// Is the reported size the same as the actual size
						// written in bytes?
						if(expected != conns[i-1]->len) ok = 0;

						// Is the protocol byte in place?
						if(ok && conns[i-1]->buff[0] != 1) ok = 0;

						// Parse connection's IMEI
						uint64_t imei = imei_uint64((conns[i-1]->buff + 10), 15);
						// Check for invalid IMEIs
						if(ok && imei == UINT64_MAX) ok = 0;

						// If everything is ok so far, and 
						// we have a list of approved IMEIs, 
						// check it for this IMEI
						if(ok && imeis != NULL && !imei_set_search(imeis, imei)){
							ok = 0;
							if(daemonize){
								syslog(LOG_NOTICE, "IMEI rejected: %.*s",
									15, (conns[i-1]->buff + 10)
								);
							} else {
								fprintf(stderr, "%s: IMEI rejected: %.*s\n",
									exec_name, 15, (conns[i-1]->buff + 10)
								);
							}
						}

						if(ok) write_connection(conns[i - 1], strip);
					}

					free(conns[i - 1]);
					conns[i - 1] = NULL;

					close(fds[i].fd);
					fds[i].fd = 0;
					fds[i].events = 0;
				}

			// Sweep up dead connections
			} else if(fds[i].fd){
				struct timeval tv;
				gettimeofday(&tv, NULL);

				if((tv.tv_sec - conns[i - 1]->tv.tv_sec) >= TIMEOUT){
					free(conns[i - 1]);
					conns[i - 1] = NULL;

					close(fds[i].fd);
					fds[i].fd = fds[i].events = 0;
				}
			}

			// Count total free slots
			if(!fds[i].fd){
				++slotfree;
				// Note the first available slot
				if(!slotfirst) slotfirst = i;
			}

			// Clear events
			fds[i].revents = 0;
		}

		// Handle new connections
		if((fds[0].revents & POLLIN) && slotfree){
			struct sockaddr_in client;
			socklen_t client_sz = sizeof(client);

			int fd = accept(sockfd, (struct sockaddr*)&client, &client_sz);
			if(fd == -1){
				if(daemonize){
					syslog(LOG_ERR, "accept failed - %s",
						strerror(errno)
					);
				} else {
					fprintf(stderr, "%s: accept failed - %s\n",
						exec_name, strerror(errno)
					);
				}
				continue;
			}
			
			// Get current flags before setting nonblock
			if((flags = fcntl(fd, F_GETFL)) == -1){
				if(daemonize){
					syslog(LOG_ERR, "cannot get connection flags - %s",
						strerror(errno)
					);
				} else {
					fprintf(stderr, "%s: cannot get connection flags - %s\n",
						exec_name, strerror(errno)
					);
				}
				close(fd);
				continue;
			}

			// Set nonblocking on socket
			if(fcntl(fd, F_SETFL, flags | O_NONBLOCK) == -1){
				if(daemonize){
					syslog(LOG_ERR, "cannot set connection flags - %s",
						strerror(errno)
					);
				} else {
					fprintf(stderr, "%s: cannot set connection flags - %s\n",
						exec_name, strerror(errno)
					);
				}
				close(fd);
				continue;
			}

			// Allocate memory for the connection's state. 
			conns[slotfirst - 1] = malloc(sizeof(struct connection));
			if(conns[slotfirst - 1] == NULL){
				if(daemonize){
					syslog(LOG_ERR, "out of memory");
				} else {
					fprintf(stderr, "%s: out of memory\n", exec_name);
				}
				close(fd);
				continue;
			}

			conns[slotfirst - 1]->len = 0; // reset buffer
			conns[slotfirst - 1]->addr = client.sin_addr; // save source ip
			gettimeofday(&conns[slotfirst - 1]->tv, NULL); // mark connection time

			// Add to poll list
			fds[slotfirst].fd = fd;
			fds[slotfirst].events = POLLIN;
			fds[slotfirst].revents = 0;
			slotfree--;

			fds[0].revents = 0;
		}

		// Reset polling event on listener socket - 
		// if there's no free slots, stop polling the listener
		if(slotfree) fds[0].events = POLLIN;
		else fds[0].events = 0;
	}

	shutdown_clean:
	if(sockfd > 0) close(sockfd);
	if(fds != NULL){ free(fds); fds = NULL; }
	if(conns != NULL){
		for(int i = 0; i < slotlimit; i++){
			if(conns[i] != NULL) free(conns[i]);
		}
		free(conns); conns = NULL;
	}
	return 0;

	shutdown_error:
	if(sockfd > 0) close(sockfd);
	if(fds != NULL){ free(fds); fds = NULL; }
	if(conns != NULL){
		for(int i = 0; i < slotlimit; i++){
			if(conns[i] != NULL) free(conns[i]);
		}
		free(conns); conns = NULL;
	}
	return 1;
}


static int write_connection(struct connection *conn, int strip)
{
	// Build a filename for the file we're writing based
	// on the time and the source IP address
	char timebuff[20];
	struct tm timeptr;
	localtime_r(&conn->tv.tv_sec, &timeptr);
	strftime(&timebuff[0], 20, "%Y%m%d%H%M%S", &timeptr);

	char fnbuff[50];
	snprintf(&fnbuff[0], 50, "%s.%ld-%s.dat",
		timebuff, (long)conn->tv.tv_usec,
		inet_ntoa(conn->addr)
	);

	int ofd;
	if((ofd = open(fnbuff, O_WRONLY | O_CREAT, 0644)) == -1){
		if(daemonize){
			syslog(LOG_ERR, "cannot open %s - %s",
				fnbuff, strerror(errno)
			);
		} else {
			fprintf(stderr, "%s: cannot open %s - %s\n",
				exec_name, fnbuff, strerror(errno)
			);
		}
		return 1;
	}


	unsigned char *buff = (strip ? (conn->buff + 3) : conn->buff);
	ssize_t len = (strip ? (conn->len - 3) : conn->len);


	// Write out the buffer
	if(write(ofd, buff, len) == -1){
		if(daemonize){
			syslog(LOG_ERR, "write error on %s - %s",
				fnbuff, strerror(errno)
			);
		} else {
			fprintf(stderr, "%s: write error on %s - %s\n",
				exec_name, fnbuff, strerror(errno)
			);
		}
		close(ofd);
		return 1;
	}

	close(ofd);
	return 0;
}