#include <stdio.h>
#include <stdbool.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <poll.h>
#include <unistd.h>
#include "util.h"
#include "sshnc.h"


static bool prompt_yn(const char *text) {
	printf("%s [y/n] ", 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;
}

static bool hostkey_checker(const unsigned char *hash, size_t length, void *userdata) {
	(void)userdata;
	printf("Server host key hash: %s\n", sshnc_print_hash(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?");
	if (!response) {
		printf("Disconnecting.\n");
	}

	return response;
}

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;
	}

	struct sshnc_client *client;
	enum sshnc_retval ret = sshnc_connect(
			server_host, port, "tomsg", "tomsg", hostkey_checker, NULL, &client);

	if (ret != SSHNC_OK) {
		fprintf(stderr, "Could not connect: %s\n", sshnc_strerror(ret));
		return 1;
	}

	struct pollfd polls[2];
	polls[0] = (struct pollfd){
		.fd = sshnc_poll_fd(client),
		.events = POLLIN,
	};
	polls[1] = (struct pollfd){
		.fd = STDIN_FILENO,
		.events = POLLIN,
	};

	while (true) {
		int pollret = poll(polls, sizeof polls / sizeof polls[0], -1);
		if (pollret < 0) {
			perror("poll");
			goto cleanup;
		}

		if (polls[0].revents & (POLLERR | POLLNVAL)) {
			fprintf(stderr, "Error reading from SSH socket\n");
			goto cleanup;
		}
		if (polls[1].revents & (POLLERR | POLLNVAL)) {
			fprintf(stderr, "Error reading from stdin\n");
			goto cleanup;
		}

		if (polls[0].revents & (POLLIN | POLLHUP)) {
			// Get all data currently available
			bool should_exit = false;
			while (true) {
				char buffer[4096];
				size_t length = 0;
				ret = sshnc_maybe_recv(client, sizeof buffer, buffer, &length);
				if (ret == SSHNC_AGAIN) break;
				if (ret == SSHNC_EOF) {
					should_exit = true;
					break;
				} else if (ret == SSHNC_OK) {
					fwrite(buffer, 1, length, stdout);
				} else {
					fprintf(stderr, "Error on SSH recv: %s\n", sshnc_strerror(ret));
					goto cleanup;
				}
			}

			if (should_exit) break;
		}

		if (polls[1].revents & (POLLIN | POLLHUP)) {
			char buffer[4096];
			ssize_t nr = read(STDIN_FILENO, buffer, sizeof buffer);
			if (nr < 0) {
				perror("Error reading from stdin");
				goto cleanup;
			}
			if (nr == 0) {
				break;
			}

			ret = sshnc_send(client, buffer, nr);
			if (ret == SSHNC_EOF) {
				break;
			} else if (ret != SSHNC_OK) {
				fprintf(stderr, "Error on SSH send: %s\n", sshnc_strerror(ret));
				goto cleanup;
			}
		}
	}

	sshnc_close(client);
	return 0;

cleanup:
	sshnc_close(client);
	return 1;
}