#include <stdio.h>
#include <stdbool.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <errno.h>
#include <assert.h>
#include <libssh/callbacks.h>
#include <sys/select.h>
#include <poll.h>
#include "util.h"


static bool prompt_yn(const char *text) {
	printf("%s", text);
	fflush(stdout);

	bool response;

	char *line = NULL;
	size_t linelen = 0;
	while (true) {
		ssize_t nr = getline(&line, &linelen, stdin);
		if (nr == -1) {
			perror("getline");
			exit(1);
		}

		while (nr > 0 && isspace(line[nr - 1])) nr--;

		for (ssize_t i = 0; i < nr; i++) line[i] = tolower(line[i]);

		if ((nr == 1 && line[0] == 'y') || (nr == 3 && memcmp(line, "yes", 3) == 0)) {
			response = true;
			break;
		}
		if ((nr == 1 && line[0] == 'n') || (nr == 2 && memcmp(line, "no", 2) == 0)) {
			response = false;
			break;
		}

		printf("Please answer with 'y', 'n', 'yes' or 'no'. [y/n]");
		fflush(stdout);
	}

	free(line);
	return response;
}

struct session_data {
	ssh_session session;
	ssh_channel channel;
	bool should_close;
};

void channel_close_cb(ssh_session session, ssh_channel channel, void *sesdata_) {
	(void)session; (void)channel;
	struct session_data *sesdata = (struct session_data*)sesdata_;
	sesdata->should_close = true;
}

void channel_eof_cb(ssh_session session, ssh_channel channel, void *sesdata_) {
	(void)session; (void)channel;
	struct session_data *sesdata = (struct session_data*)sesdata_;
	sesdata->should_close = true;
}

int channel_data_cb(ssh_session session, ssh_channel channel, void *data, uint32_t len, int is_stderr, void *sesdata_) {
	(void)session; (void)channel; (void)is_stderr;
	struct session_data *sesdata = (struct session_data*)sesdata_;

	const char *start = (const char*)data;
	const char *cursor = start;
	const char *end = cursor + len;

	while (cursor < end) {
		ssize_t nw = write(STDOUT_FILENO, cursor, end - cursor);
		if (nw < 0) {
			if (errno == EINTR) continue;
			perror("write(stdout)");
			sesdata->should_close = true;
			return cursor - start;
		}
		assert(nw > 0);
		cursor += nw;
	}

	return len;
}

int channel_write_wontblock_cb(ssh_session session, ssh_channel channel, size_t bytes, void *sesdata_) {
	(void)session; (void)channel; (void)sesdata_;
	fprintf(stderr, "(write won't block for %zu bytes)\n", bytes);
	return 0;
}

int stdin_data_cb(int fd, int revents, void *sesdata_) {
	(void)fd;
	struct session_data *sesdata = (struct session_data*)sesdata_;

	if (revents & POLLIN) {
		char buffer[1024];
		ssize_t nr = read(STDIN_FILENO, buffer, sizeof buffer);
		if (nr < 0) {
			if (errno == EINTR) return 0;
			perror("read(stdin)");
			sesdata->should_close = true;
			return 0;
		}

		if (nr == 0) {  // eof
			sesdata->should_close = true;
			return 0;
		}

		const char *cursor = buffer;
		while (cursor < buffer + nr) {
			int ret = ssh_channel_write(sesdata->channel, cursor, buffer + nr - cursor);
			if (ret == SSH_ERROR) {
				fprintf(stderr, "Error writing to channel: %s\n",
						ssh_get_error(sesdata->channel));
				sesdata->should_close = true;
				return 0;
			}
			assert(ret > 0);
			cursor += ret;
		}
	}
	return 0;
}

int main(int argc, char **argv) {
	const char *server_host = NULL;
	int port = 2222;

	if (argc != 2) {
		fprintf(stderr, "Usage: %s <server[:port]>\n", argv[0]);
		fprintf(stderr, "If port is not specified, %d is assumed.\n", port);
		return 1;
	}

	if (!parse_host_port(argv[1], &server_host, &port)) {
		fprintf(stderr, "Cannot parse host:port from argument '%s'\n", argv[1]);
		return 1;
	}

	ssh_session session = ssh_new();
	if (!session) {
		fprintf(stderr, "Could not open SSH session\n");
		goto cleanup_unconnected;
	}

	const char *ciphers_str = "aes256-gcm@openssh.com,aes256-ctr,aes256-cbc";
	bool procconfig = false;
	bool ok = true;
	ok &= ssh_options_set(session, SSH_OPTIONS_PROCESS_CONFIG, &procconfig) == SSH_OK;
	ok &= ssh_options_set(session, SSH_OPTIONS_USER, "tomsg") == SSH_OK;
	ok &= ssh_options_set(session, SSH_OPTIONS_HOST, server_host) == SSH_OK;
	ok &= ssh_options_set(session, SSH_OPTIONS_PORT, &port) == SSH_OK;
	ok &= ssh_options_set(session, SSH_OPTIONS_CIPHERS_C_S, ciphers_str) == SSH_OK;
	ok &= ssh_options_set(session, SSH_OPTIONS_CIPHERS_S_C, ciphers_str) == SSH_OK;
	// int loglevel = SSH_LOG_PROTOCOL;
	// ok &= ssh_options_set(session, SSH_OPTIONS_LOG_VERBOSITY, &loglevel) == SSH_OK;

	if (!ok) {
		fprintf(stderr, "Could not set options on SSH session: %s\n", ssh_get_error(session));
		goto cleanup_unconnected;
	}

	if (ssh_connect(session) != SSH_OK) {
		fprintf(stderr, "Could not connect to %s:%d: %s\n",
				server_host, port, ssh_get_error(session));
		goto cleanup_unconnected;
	}

	ssh_key host_key;
	if (ssh_get_server_publickey(session, &host_key) != SSH_OK) {
		fprintf(stderr, "Could not get host key from session: %s\n", ssh_get_error(session));
		goto cleanup_connected;
	}

	unsigned char *host_key_hash = NULL;
	size_t host_key_hash_length = 0;
	if (ssh_get_publickey_hash(host_key, SSH_PUBLICKEY_HASH_SHA256, &host_key_hash, &host_key_hash_length) != SSH_OK) {
		fprintf(stderr, "Failed to hash host key!\n");
		goto cleanup_connected;
	}

	printf("Server host key hash: ");
	fflush(stdout);
	ssh_print_hash(SSH_PUBLICKEY_HASH_SHA256, host_key_hash, host_key_hash_length);

	bool response = prompt_yn(
			"Does this hash match the one given to you by the server administrator, or by a\n"
			"member that you trust and is already connected to the server? [y/n] ");
	if (!response) {
		printf("Disconnecting.\n");
		goto cleanup_connected;
	}

	printf("Connected.\n");

retry_userauth:
	switch (ssh_userauth_none(session, NULL)) {
		case SSH_AUTH_ERROR:
			fprintf(stderr, "Error authenticating: %s\n", ssh_get_error(session));
			return 1;

		case SSH_AUTH_DENIED:
		case SSH_AUTH_PARTIAL:
			fprintf(stderr, "Server denied authentication.\n");
			return 1;

		case SSH_AUTH_SUCCESS:
			break;

		case SSH_AUTH_AGAIN:
			if (ssh_get_status(session) & (SSH_CLOSED | SSH_CLOSED_ERROR)) {
				fprintf(stderr, "Socket unexpectedly closed!\n");
				return 1;
			}
			goto retry_userauth;
	}

	printf("Authenticated.\n");

	ssh_channel channel = ssh_channel_new(session);
	if (!channel) {
		fprintf(stderr, "Failed to allocate channel: %s\n", ssh_get_error(session));
		goto cleanup_connected;
	}

	printf("Created channel\n");

	if (ssh_channel_open_session(channel) != SSH_OK) {
		fprintf(stderr, "Failed to open channel: %s\n", ssh_get_error(channel));
		goto cleanup_connected;
	}

	printf("Opened channel\n");

	if (ssh_channel_request_subsystem(channel, "tomsg") != SSH_OK) {
		fprintf(stderr, "Server did not allow opening 'tomsg' channel: %s\n", ssh_get_error(channel));
		goto cleanup_connected;
	}

	printf("Obtained tomsg subsystem on channel\n");

	struct session_data *sesdata = malloc(sizeof(struct session_data));
	if (!sesdata) {
		fprintf(stderr, "Out of memory (allocating session_data)!\n");
		return 1;
	}

	sesdata->session = session;
	sesdata->channel = channel;
	sesdata->should_close = false;

	struct ssh_channel_callbacks_struct chan_cb;
	memset(&chan_cb, 0, sizeof chan_cb);
	ssh_callbacks_init(&chan_cb);
	chan_cb.userdata = sesdata;
	chan_cb.channel_close_function = channel_close_cb;
	chan_cb.channel_eof_function = channel_eof_cb;
	chan_cb.channel_data_function = channel_data_cb;
	chan_cb.channel_write_wontblock_function = channel_write_wontblock_cb;

	if (ssh_set_channel_callbacks(channel, &chan_cb) != SSH_OK) {
		fprintf(stderr, "Failed to set channel callbacks\n");
		goto cleanup_connected;
	}

	printf("Set callbacks\n");

	ssh_event event = ssh_event_new();
	if (!event
			|| ssh_event_add_session(event, session) != SSH_OK
			|| ssh_event_add_fd(event, STDIN_FILENO, POLLIN, stdin_data_cb, sesdata) != SSH_OK) {
		fprintf(stderr, "Failed to create ssh event context\n");
		goto cleanup_connected;
	}

	printf("Created event object\n");

	while (!sesdata->should_close) {
		// printf("poll loop...\n");
		ssh_event_dopoll(event, -1);
		int status = ssh_get_status(session);
		if (status & (SSH_CLOSED | SSH_CLOSED_ERROR)) goto cleanup_connected;
		if (status & SSH_READ_PENDING) {
			printf("read pending?\n");
		}
	}

cleanup_connected:
	ssh_disconnect(session);
cleanup_unconnected:
	ssh_free(session);
}