#include <assert.h>
#include <errno.h>
#include <fcntl.h>
#include <getopt.h>
#include <limits.h>
#include <signal.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/time.h>
#include <sys/wait.h>
#include <systemd/sd-daemon.h>
#include <unistd.h>
#include <algorithm>
#include <map>
#include <set>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

#include "tlse.h"

#include "acceptor.h"
#include "accesslog.h"
#include "config.h"
#include "input.h"
#include "input_stats.h"
#include "log.h"
#include "sa_compare.h"
#include "serverpool.h"
#include "state.pb.h"
#include "stats.h"
#include "stream.h"
#include "util.h"
#include "version.h"

using namespace std;

AccessLogThread *access_log = nullptr;
ServerPool *servers = nullptr;
volatile bool hupped = false;
volatile bool stopped = false;

typedef pair<string, Input::Encoding> InputKey;

namespace {

struct OrderByConnectionTime {
	bool operator() (const ClientProto &a, const ClientProto &b) const {
		if (a.connect_time_sec() != b.connect_time_sec())
			return a.connect_time_sec() < b.connect_time_sec();
		return a.connect_time_nsec() < b.connect_time_nsec();
	}
};

// An arbitrary ordering.
struct AcceptorConfigCompare {
	bool operator() (const AcceptorConfig &a, const AcceptorConfig &b) const {
		int cmp = a.certificate_chain.compare(b.certificate_chain);
		if (cmp != 0) {
			return cmp < 0;
		}

		cmp = a.private_key.compare(b.private_key);
		if (cmp != 0) {
			return cmp < 0;
		}

		return Sockaddr6Compare()(a.addr, b.addr);
	}
};

}  // namespace

struct InputWithRefcount {
	Input *input;
	int refcount;
};

void hup(int signum)
{
	hupped = true;
	if (signum == SIGINT) {
		stopped = true;
	}
}

void do_nothing(int signum)
{
}

CubemapStateProto collect_state(const timespec &serialize_start,
                                const vector<Acceptor *> acceptors,
                                const multimap<InputKey, InputWithRefcount> inputs,
                                ServerPool *servers)
{
	CubemapStateProto state = servers->serialize();  // Fills streams() and clients().
	state.set_serialize_start_sec(serialize_start.tv_sec);
	state.set_serialize_start_usec(serialize_start.tv_nsec / 1000);

	for (Acceptor *acceptor : acceptors) {	
		state.add_acceptors()->MergeFrom(acceptor->serialize());
	}

	for (const auto &key_and_input_with_refcount : inputs) {
		state.add_inputs()->MergeFrom(key_and_input_with_refcount.second.input->serialize());
	}

	return state;
}

// Find all port statements in the configuration file, and create acceptors for them.
vector<Acceptor *> create_acceptors(
	const Config &config,
	map<AcceptorConfig, Acceptor *, AcceptorConfigCompare> *deserialized_acceptors)
{
	vector<Acceptor *> acceptors;
	for (const AcceptorConfig &acceptor_config : config.acceptors) {
		Acceptor *acceptor = nullptr;
		const auto deserialized_acceptor_it = deserialized_acceptors->find(acceptor_config);
		if (deserialized_acceptor_it != deserialized_acceptors->end()) {
			acceptor = deserialized_acceptor_it->second;
			deserialized_acceptors->erase(deserialized_acceptor_it);
		} else {
			int server_sock = create_server_socket(acceptor_config.addr, TCP_SOCKET);
			acceptor = new Acceptor(server_sock, acceptor_config.addr,
			                        acceptor_config.certificate_chain,
			                        acceptor_config.private_key);
		}
		acceptor->run();
		acceptors.push_back(acceptor);
	}

	// Close all acceptors that are no longer in the configuration file.
	for (auto &config_and_acceptor : *deserialized_acceptors) {
		config_and_acceptor.second->close_socket();
		delete config_and_acceptor.second;
	}

	return acceptors;
}

void create_config_input(const string &src, Input::Encoding encoding, multimap<InputKey, InputWithRefcount> *inputs)
{
	if (src.empty()) {
		return;
	}
	InputKey key(src, encoding);
	if (inputs->count(key) != 0) {
		return;
	}

	InputWithRefcount iwr;
	iwr.input = create_input(src, encoding);
	if (iwr.input == nullptr) {
		log(ERROR, "did not understand URL '%s' or source encoding was invalid, clients will not get any data.",
			src.c_str());
		return;
	}
	iwr.refcount = 0;
	inputs->insert(make_pair(key, iwr));
}

// Find all streams in the configuration file, and create inputs for them.
void create_config_inputs(const Config &config, multimap<InputKey, InputWithRefcount> *inputs)
{
	for (const StreamConfig &stream_config : config.streams) {
		if (stream_config.src == "delete") {
			// Ignored for pre-1.4.0 configuration compatibility.
			continue;
		}
		create_config_input(stream_config.src, Input::Encoding(stream_config.src_encoding), inputs);
	}
	for (const UDPStreamConfig &udpstream_config : config.udpstreams) {
		create_config_input(udpstream_config.src, Input::INPUT_ENCODING_RAW, inputs);
	}
}

void create_streams(const Config &config,
                    const set<string> &deserialized_urls,
                    multimap<InputKey, InputWithRefcount> *inputs)
{
	// HTTP streams.
	set<string> expecting_urls = deserialized_urls;
	for (const StreamConfig &stream_config : config.streams) {
		int stream_index;

		expecting_urls.erase(stream_config.url);

		if (stream_config.src == "delete") {
			// Ignored for pre-1.4.0 configuration compatibility.
			continue;
		}

		if (deserialized_urls.count(stream_config.url) == 0) {
			stream_index = servers->add_stream(stream_config.url,
			                                   stream_config.hls_url,
			                                   stream_config.backlog_size,
			                                   stream_config.prebuffering_bytes,
			                                   Stream::Encoding(stream_config.encoding),
			                                   Stream::Encoding(stream_config.src_encoding),
			                                   stream_config.hls_frag_duration,
			                                   stream_config.hls_backlog_margin,
			                                   stream_config.allow_origin);
		} else {
			stream_index = servers->lookup_stream_by_url(stream_config.url);
			assert(stream_index != -1);
			servers->set_backlog_size(stream_index, stream_config.backlog_size);
			if (!stream_config.hls_url.empty()) {
				servers->register_hls_url(stream_index, stream_config.hls_url);
			}
			servers->set_prebuffering_bytes(stream_index, stream_config.prebuffering_bytes);
			servers->set_encoding(stream_index,
			                      Stream::Encoding(stream_config.encoding));
			servers->set_src_encoding(stream_index,
			                          Stream::Encoding(stream_config.src_encoding));
			servers->set_hls_frag_duration(stream_index, stream_config.hls_frag_duration);
			servers->set_hls_backlog_margin(stream_index, stream_config.hls_backlog_margin);
			servers->set_allow_origin(stream_index, stream_config.allow_origin);
		}

		servers->set_pacing_rate(stream_index, stream_config.pacing_rate);

		string src = stream_config.src;
		Input::Encoding src_encoding = Input::Encoding(stream_config.src_encoding);
		if (!src.empty()) {
			const auto input_it = inputs->find(make_pair(src, src_encoding));
			if (input_it != inputs->end()) {
				input_it->second.input->add_destination(stream_index);
				++input_it->second.refcount;
			}
		}
	}

	// Warn about any streams servers we've lost.
	for (const string &url : expecting_urls) {
		log(WARNING, "stream '%s' disappeared from the configuration file. "
		             "It will not be deleted, but clients will not get any new inputs. "
		             "If you really meant to delete it, set src=delete and reload.",
		             url.c_str());
	}

	// UDP streams.
	for (const UDPStreamConfig &udpstream_config : config.udpstreams) {
		int stream_index = servers->add_udpstream(
			udpstream_config.dst,
			udpstream_config.pacing_rate,
			udpstream_config.ttl,
			udpstream_config.multicast_iface_index);

		string src = udpstream_config.src;
		if (!src.empty()) {
			auto input_it = inputs->find(make_pair(src, Input::INPUT_ENCODING_RAW));
			assert(input_it != inputs->end());
			input_it->second.input->add_destination(stream_index);
			++input_it->second.refcount;
		}
	}

	// HTTP gen204 endpoints.
	for (const Gen204Config &ping_config : config.pings) {
		servers->add_gen204(ping_config.url, ping_config.allow_origin);
	}
}
	
void open_logs(const vector<LogConfig> &log_destinations)
{
	for (const LogConfig &log_destination : log_destinations) {
		if (log_destination.type == LogConfig::LOG_TYPE_FILE) {
			add_log_destination_file(log_destination.filename);
		} else if (log_destination.type == LogConfig::LOG_TYPE_CONSOLE) {
			add_log_destination_console();
		} else if (log_destination.type == LogConfig::LOG_TYPE_SYSLOG) {
			add_log_destination_syslog();
		} else {
			assert(false);
		}
	}
	start_logging();
}
	
bool dry_run_config(const string &argv0, const string &config_filename)
{
	char *argv0_copy = strdup(argv0.c_str());
	char *config_filename_copy = strdup(config_filename.c_str());

	pid_t pid = fork();
	switch (pid) {
	case -1:
		log_perror("fork()");
		free(argv0_copy);
		free(config_filename_copy);
		return false;
	case 0:
		// Child.
		execlp(argv0_copy, argv0_copy, "--test-config", config_filename_copy, nullptr);
		log_perror(argv0_copy);
		_exit(1);
	default:
		// Parent.
		break;
	}
		
	free(argv0_copy);
	free(config_filename_copy);

	int status;
	pid_t err;
	do {
		err = waitpid(pid, &status, 0);
	} while (err == -1 && errno == EINTR);

	if (err == -1) {
		log_perror("waitpid()");
		return false;
	}	

	return (WIFEXITED(status) && WEXITSTATUS(status) == 0);
}

void find_all_streams(const Config &config, set<string> *all_urls)
{
	for (const StreamConfig &stream_config : config.streams) {
		if (stream_config.src == "delete") {
			log(WARNING, "stream '%s' has src=delete; ignoring it. Since Cubemap 1.4.0, you do not "
			             "need to set src=delete to delete streams anymore; just delete them from "
			             "the configuration file.",
				     stream_config.url.c_str());
		} else {
			all_urls->insert(stream_config.url);
		}
	}
}

int main(int argc, char **argv)
{
	signal(SIGHUP, hup);
	signal(SIGINT, hup);
	signal(SIGUSR1, do_nothing);  // Used in internal signalling.
	signal(SIGPIPE, SIG_IGN);

	tls_init();
	
	// Parse options.
	int state_fd = -1;
	bool test_config = false;
	for ( ;; ) {
		static const option long_options[] = {
			{ "state", required_argument, 0, 's' },
			{ "test-config", no_argument, 0, 't' },
			{ 0, 0, 0, 0 }
		};
		int option_index = 0;
		int c = getopt_long(argc, argv, "s:t", long_options, &option_index);
     
		if (c == -1) {
			break;
		}
		switch (c) {
		case 's':
			state_fd = atoi(optarg);
			break;
		case 't':
			test_config = true;
			break;
		default:
			fprintf(stderr, "Unknown option '%s'\n", argv[option_index]);
			exit(1);
		}
	}

	string config_filename = "cubemap.config";
	if (optind < argc) {
		config_filename = argv[optind++];
	}

	// Canonicalize argv[0] and config_filename.
	char argv0_canon[PATH_MAX];
	char config_filename_canon[PATH_MAX];

	if (realpath("/proc/self/exe", argv0_canon) == nullptr) {
		log_perror(argv[0]);
		exit(1);
	}
	if (realpath(config_filename.c_str(), config_filename_canon) == nullptr) {
		log_perror(config_filename.c_str());
		exit(1);
	}

	// Now parse the configuration file.
	Config config;
	if (!parse_config(config_filename_canon, &config)) {
		exit(1);
	}
	if (test_config) {
		exit(0);
	}
	
	// Ideally we'd like to daemonize only when we've started up all threads etc.,
	// but daemon() forks, which is not good in multithreaded software, so we'll
	// have to do it here.
	if (config.daemonize) {
		if (daemon(0, 0) == -1) {
			log_perror("daemon");
			exit(1);
		}
	}

start:
	// Open logs as soon as possible.
	open_logs(config.log_destinations);

	log(INFO, "Cubemap " SERVER_VERSION " starting.");
	if (config.access_log_file.empty()) {
		// Create a dummy logger.
		access_log = new AccessLogThread();
	} else {
		access_log = new AccessLogThread(config.access_log_file);
	}
	access_log->run();

	servers = new ServerPool(config.num_servers);

	// Find all the streams that are to be kept.
	set<string> all_urls;
	find_all_streams(config, &all_urls);

	CubemapStateProto loaded_state;
	timespec serialize_start;
	set<string> deserialized_urls;
	map<AcceptorConfig, Acceptor *, AcceptorConfigCompare> deserialized_acceptors;
	multimap<InputKey, InputWithRefcount> inputs;  // multimap due to older versions without deduplication.
	if (state_fd != -1) {
		log(INFO, "Deserializing state from previous process...");
		string serialized;
		if (!read_tempfile_and_close(state_fd, &serialized)) {
			exit(1);
		}
		if (!loaded_state.ParseFromString(serialized)) {
			log(ERROR, "Failed deserialization of state.");
			exit(1);
		}

		serialize_start.tv_sec = loaded_state.serialize_start_sec();
		serialize_start.tv_nsec = loaded_state.serialize_start_usec() * 1000ull;

		// Deserialize the streams.
		unordered_map<string, string> stream_headers_for_url;  // See below.
		for (const StreamProto &stream : loaded_state.streams()) {
			if (all_urls.count(stream.url()) == 0) {
				// Delete the stream backlogs.
				log(INFO, "Deleting stream '%s'.", stream.url().c_str());
				for (const int fd : stream.data_fds()) {
					safe_close(fd);
				}
			} else {
				vector<int> data_fds;
				for (const int fd : stream.data_fds()) {
					data_fds.push_back(fd);
				}

				servers->add_stream_from_serialized(stream, data_fds);
				deserialized_urls.insert(stream.url());

				stream_headers_for_url.insert(make_pair(stream.url(), stream.stream_header()));
			}
		}

		// Deserialize the inputs. Note that we don't actually add them to any stream yet.
		for (const InputProto &serialized_input : loaded_state.inputs()) {
			InputWithRefcount iwr;
			iwr.input = create_input(serialized_input);
			iwr.refcount = 0;

			Input::Encoding src_encoding = serialized_input.is_metacube_encoded() ?
				Input::INPUT_ENCODING_METACUBE :
				Input::INPUT_ENCODING_RAW;
			InputKey key(serialized_input.url(), src_encoding);
			inputs.insert(make_pair(key, iwr));
		} 

		// Deserialize the acceptors.
		for (const AcceptorProto &serialized_acceptor : loaded_state.acceptors()) {
			AcceptorConfig config;
			config.addr = extract_address_from_acceptor_proto(serialized_acceptor);
			config.certificate_chain = serialized_acceptor.certificate_chain();
			config.private_key = serialized_acceptor.private_key();
			deserialized_acceptors.insert(make_pair(
				config,
				new Acceptor(serialized_acceptor)));
		}

		log(INFO, "Deserialization done.");
	}

	// Add any new inputs coming from the config.
	create_config_inputs(config, &inputs);
	
	// Find all streams in the configuration file, create them, and connect to the inputs.
	create_streams(config, deserialized_urls, &inputs);
	vector<Acceptor *> acceptors = create_acceptors(config, &deserialized_acceptors);

	// Make all the servers create TLS contexts for every TLS keypair we have.
	for (Acceptor *acceptor : acceptors) {
		if (acceptor->is_tls()) {
			servers->create_tls_context_for_acceptor(acceptor);
		}
	}

	// Allocate strings for the short responses.
	vector<shared_ptr<const string>> short_response_pool;
	for (const ShortResponsePool &str : loaded_state.short_response_pool()) {
		short_response_pool.emplace_back(new string(str.header_or_short_response()));
	}
	
	// Put back the existing clients. It doesn't matter which server we
	// allocate them to, so just do round-robin. However, we need to sort them
	// by connection time first, since add_client_serialized() expects that.
	sort(loaded_state.mutable_clients()->begin(),
	     loaded_state.mutable_clients()->end(),
	     OrderByConnectionTime());
	for (int i = 0; i < loaded_state.clients_size(); ++i) {
		if (!loaded_state.clients(i).url().empty() &&
		    all_urls.count(loaded_state.clients(i).url()) == 0) {
			// Belongs to a dead stream (not keepalive), so we just have to close.
			safe_close(loaded_state.clients(i).sock());
		} else {
			servers->add_client_from_serialized(loaded_state.clients(i), short_response_pool);
		}
	}
	
	short_response_pool.clear();  // No longer needed; the clients have their own refcounts now.

	// Put back the HLS zombies. There's no really good allocation here
	// except round-robin; it would be marginally more efficient to match it
	// to the client (since that would have them deleted immediately when
	// the client requests the next fragment, instead of being later weeded
	// out during statistics collection), but it's not a big deal.
	for (const HLSZombieProto &zombie_proto : loaded_state.hls_zombies()) {
		servers->add_hls_zombie_from_serialized(zombie_proto);
	}

	servers->run();

	// Now delete all inputs that are longer in use, and start the others.
	for (auto input_it = inputs.begin(); input_it != inputs.end(); ) {
		if (input_it->second.refcount == 0) {
			if (input_it->first.second == Input::INPUT_ENCODING_RAW) {
				log(WARNING, "Raw input '%s' no longer in use, closing.",
				    input_it->first.first.c_str());
			} else {
				assert(input_it->first.second == Input::INPUT_ENCODING_METACUBE);
				log(WARNING, "Metacube input '%s' no longer in use, closing.",
				    input_it->first.first.c_str());
			}
			input_it->second.input->close_socket();
			delete input_it->second.input;
			inputs.erase(input_it++);
		} else {
			input_it->second.input->run();
			++input_it;
		}
	}

	// Start writing statistics.
	unique_ptr<StatsThread> stats_thread;
	if (!config.stats_file.empty()) {
		stats_thread.reset(new StatsThread(config.stats_file, config.stats_interval));
		stats_thread->run();
	}

	unique_ptr<InputStatsThread> input_stats_thread;
	if (!config.input_stats_file.empty()) {
		vector<Input*> inputs_no_refcount;
		for (const auto &key_and_input_with_refcount : inputs) {
			inputs_no_refcount.push_back(key_and_input_with_refcount.second.input);
		}

		input_stats_thread.reset(new InputStatsThread(config.input_stats_file, config.input_stats_interval, inputs_no_refcount));
		input_stats_thread->run();
	}

	timespec server_start;
	int err = clock_gettime(CLOCK_MONOTONIC, &server_start);
	assert(err != -1);
	if (state_fd != -1) {
		// Measure time from we started deserializing (below) to now, when basically everything
		// is up and running. This is, in other words, a conservative estimate of how long our
		// “glitch” period was, not counting of course reconnects if the configuration changed.
		double glitch_time = server_start.tv_sec - serialize_start.tv_sec +
			1e-9 * (server_start.tv_nsec - serialize_start.tv_nsec);
		log(INFO, "Re-exec happened in approx. %.0f ms.", glitch_time * 1000.0);
	}

	sd_notify(0, "READY=1");

	while (!hupped) {
		usleep(100000);
	}

	if (stopped) {
		sd_notify(0, "STOPPING=1");
	} else {
		sd_notify(0, "RELOADING=1");
	}

	// OK, we've been HUPed. Time to shut down everything, serialize, and re-exec.
	err = clock_gettime(CLOCK_MONOTONIC, &serialize_start);
	assert(err != -1);

	if (input_stats_thread != nullptr) {
		input_stats_thread->stop();
		input_stats_thread.reset();
	}
	if (stats_thread != nullptr) {
		stats_thread->stop();
		stats_thread.reset();
	}
	for (Acceptor *acceptor : acceptors) {
		acceptor->stop();
	}
	for (const auto &key_and_input_with_refcount : inputs) {
		key_and_input_with_refcount.second.input->stop();
		if (stopped) {
			key_and_input_with_refcount.second.input->close_socket();
		}
	}
	servers->stop();

	CubemapStateProto state;
	if (stopped) {
		log(INFO, "Shutting down.");
	} else {
		log(INFO, "Serializing state and re-execing...");
		state = collect_state(
			serialize_start, acceptors, inputs, servers);
		string serialized;
		state.SerializeToString(&serialized);
		state_fd = make_tempfile(serialized);
		if (state_fd == -1) {
			exit(1);
		}
	}
	delete servers;

	access_log->stop();
	delete access_log;
	shut_down_logging();

	if (stopped) {
		exit(0);
	}

	// OK, so the signal was SIGHUP. Check that the new config is okay, then exec the new binary.
	if (!dry_run_config(argv0_canon, config_filename_canon)) {
		open_logs(config.log_destinations);
		log(ERROR, "%s --test-config failed. Restarting old version instead of new.", argv[0]);
		hupped = false;
		shut_down_logging();
		goto start;
	}
	 
	char buf[16];
	sprintf(buf, "%d", state_fd);

	// Unset the close-on-exec flag for the state fd.
	// (This can't leak into a child, since there's only one thread left.)
	fcntl(state_fd, F_SETFD, 0);

	for ( ;; ) {
		execlp(argv0_canon, argv0_canon, config_filename_canon, "--state", buf, nullptr);
		open_logs(config.log_destinations);
		log_perror("execlp");
		log(ERROR, "re-exec of %s failed. Waiting 0.2 seconds and trying again...", argv0_canon);
		shut_down_logging();
		usleep(200000);
	}
}