#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <threads.h>
#include <assert.h>
#include <libssh/callbacks.h>
#include "sshnc.h"


// - We never use ssh_disconnect, just ssh_free. I believe the only added value
//   of ssh_disconnect is that it sends a polite quit message to the server,
//   but it has the disadvantage of then closing the socket immediately --
//   without checking on the channels, which ssh_free closes neatly.


static thread_local char libssh_additional_error_description[1024];

static size_t min_size_t(size_t a, size_t b) { return a < b ? a : b; }

static void clear_additional_error(void) {
	libssh_additional_error_description[0] = '\0';
}

static void store_additional_error(const ssh_session session) {
	const char *const ptr = ssh_get_error(session);
	const size_t available = sizeof libssh_additional_error_description;
	const size_t length = min_size_t(strlen(ptr), available - 1);
	memcpy(libssh_additional_error_description, ptr, length);
	libssh_additional_error_description[length] = '\0';
}

struct session_data {
	ssh_session session;
	ssh_channel channel;

	bool should_close;
	enum sshnc_retval close_reason;

	// Data is appended whenever some is received
	unsigned char recvdata[4096];
	size_t recvlen;
};

struct sshnc_client {
	// If is_closed, all fields in this structure have been freed. This is to
	// facilitate closing the connection on error, but later still accomodate
	// sshnc_close().
	bool is_closed;
	ssh_event event;
	struct session_data *sesdata;

	const char *libssh_error_descr;

	struct ssh_channel_callbacks_struct chan_cb;
};

const char* sshnc_print_hash(const unsigned char *hash, size_t length) {
	static char buffer[64];

    char *const fingerprint = ssh_get_fingerprint_hash(
			SSH_PUBLICKEY_HASH_SHA256, (unsigned char*)hash, length);
    if (fingerprint == NULL) return NULL;

	const size_t fingerprint_len = strlen(fingerprint);
	assert(fingerprint_len + 1 <= sizeof buffer);
	memcpy(buffer, fingerprint, fingerprint_len + 1);

	ssh_string_free_char(fingerprint);

	return buffer;
}

const char* sshnc_strerror(enum sshnc_retval code) {
	const char *description = NULL;

	switch (code) {
#define CASE(code, descr) case code: description = #code ": " descr; break;
		CASE(SSHNC_OK, "Success")
		CASE(SSHNC_EOF, "EOF")
		CASE(SSHNC_AGAIN, "Non-blocking read with no data, try again")
		CASE(SSHNC_ERR_CONNECT, "Could not connect to host")
		CASE(SSHNC_ERR_UNTRUSTED, "Hostkey checker rejected key")
		CASE(SSHNC_ERR_AUTH, "Error authenticating to server")
		CASE(SSHNC_ERR_DENIED, "Server did not accept 'none' authentication")
		CASE(SSHNC_ERR_CLOSED, "Server unexpectedly closed connection")
		CASE(SSHNC_ERR_SUBSYSTEM, "Server did not accept the subsystem channel")
		CASE(SSHNC_ERR_SESSION, "Could not open libssh session")
		CASE(SSHNC_ERR_CHANNEL, "Could not open libssh channel")
		CASE(SSHNC_ERR_OPTIONS, "Could not set libssh options")
		CASE(SSHNC_ERR_GETKEY, "Could not get key from libssh")
		CASE(SSHNC_ERR_CALLBACKS, "Sshnc would not accept our callbacks structure")
		CASE(SSHNC_ERR_EVENT, "Could not create libssh event poller")
		CASE(SSHNC_ERR_WRITE, "Could not write to ssh channel")
		CASE(SSHNC_ERR_POLL, "Could not poll the socket for activity")
#undef CASE
	}

	static char buffer[2048];
	if (description == NULL) {
		snprintf(buffer, sizeof buffer, "sshnc_strerror: unknown code %d", code);
	} else if (libssh_additional_error_description[0] != '\0') {
		snprintf(buffer, sizeof buffer, "%s (%s)",
				description, libssh_additional_error_description);
	} else {
		strcpy(buffer, description);
	}
	return buffer;
}

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

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

static 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 *const sesdata = (struct session_data*)sesdata_;

	const size_t remaining_space = sizeof sesdata->recvdata - sesdata->recvlen;
	const size_t consumed = min_size_t(len, remaining_space);

	memcpy(sesdata->recvdata + sesdata->recvlen, data, consumed);
	sesdata->recvlen += consumed;
	return consumed;
}

enum sshnc_retval sshnc_connect(
	const char *hostname,
	int port,
	const char *username,
	const char *subsystem,
	sshnc_hostkey_checker_t checker,
	void *userdata,
	struct sshnc_client **clientp  // output
) {
	clear_additional_error();

	const ssh_session session = ssh_new();
	if (!session) {
		return SSHNC_ERR_SESSION;
	}

#define RETURN(errorcode) \
	do { ssh_free(session); *clientp = NULL; return (errorcode); } while(0)

	const char *ciphers_str = "aes256-gcm@openssh.com,aes256-ctr,aes256-cbc";
	const 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, username) == SSH_OK;
	ok &= ssh_options_set(session, SSH_OPTIONS_HOST, hostname) == 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) {
		store_additional_error(session);
		RETURN(SSHNC_ERR_OPTIONS);
	}

	if (ssh_connect(session) != SSH_OK) {
		store_additional_error(session);
		RETURN(SSHNC_ERR_CONNECT);
	}

	ssh_key host_key;
	if (ssh_get_server_publickey(session, &host_key) != SSH_OK) {
		store_additional_error(session);
		RETURN(SSHNC_ERR_GETKEY);
	}

	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) {
		store_additional_error(session);
		RETURN(SSHNC_ERR_GETKEY);
	}

	if (!checker(host_key_hash, host_key_hash_length, userdata)) {
		RETURN(SSHNC_ERR_UNTRUSTED);
	}

	// Now we're connected, let's do authentication.

retry_userauth:
	switch (ssh_userauth_none(session, NULL)) {
		case SSH_AUTH_ERROR:
			store_additional_error(session);
			RETURN(SSHNC_ERR_AUTH);

		case SSH_AUTH_DENIED:
		case SSH_AUTH_PARTIAL:
			RETURN(SSHNC_ERR_DENIED);

		case SSH_AUTH_SUCCESS:
			break;

		case SSH_AUTH_AGAIN:
			if (ssh_get_status(session) & (SSH_CLOSED | SSH_CLOSED_ERROR)) {
				RETURN(SSHNC_ERR_CLOSED);
			}
			goto retry_userauth;
	}

	// We're authenticated; open channel and set it up.

	const ssh_channel channel = ssh_channel_new(session);
	if (!channel) {
		store_additional_error(session);
		RETURN(SSHNC_ERR_CHANNEL);
	}

	if (ssh_channel_open_session(channel) != SSH_OK ||
			ssh_channel_request_subsystem(channel, subsystem) != SSH_OK) {
		store_additional_error(session);
		RETURN(SSHNC_ERR_SUBSYSTEM);
	}

	// Fully connected, now set up libssh to our wishes.

	struct sshnc_client *const client = malloc(sizeof(struct sshnc_client));
	client->is_closed = false;
	client->libssh_error_descr = NULL;

	client->sesdata = malloc(sizeof(struct session_data));
	*client->sesdata = (struct session_data){
		.session = session,
		.channel = channel,
		.should_close = false,
		.close_reason = SSHNC_OK,
		.recvlen = 0,
	};

	memset(&client->chan_cb, 0, sizeof client->chan_cb);
	ssh_callbacks_init(&client->chan_cb);
	client->chan_cb.userdata = client->sesdata;
	client->chan_cb.channel_close_function = channel_close_cb;
	client->chan_cb.channel_eof_function = channel_eof_cb;
	client->chan_cb.channel_data_function = channel_data_cb;

	if (ssh_set_channel_callbacks(channel, &client->chan_cb) != SSH_OK) {
		store_additional_error(session);
		RETURN(SSHNC_ERR_CALLBACKS);
	}

	client->event = ssh_event_new();
	if (!client->event || ssh_event_add_session(client->event, session) != SSH_OK) {
		store_additional_error(session);
		RETURN(SSHNC_ERR_EVENT);
	}

	*clientp = client;
	return SSHNC_OK;

#undef RETURN
}

static void sshnc_close_nofree(struct sshnc_client *client) {
	ssh_event_free(client->event);
	ssh_free(client->sesdata->session);
	free(client->sesdata);
	client->is_closed = true;
}

void sshnc_close(struct sshnc_client *client) {
	if (!client->is_closed) {
		sshnc_close_nofree(client);
	}
	free(client);
}

int sshnc_poll_fd(struct sshnc_client *client) {
	if (client->is_closed) return -1;
	return ssh_get_fd(client->sesdata->session);
}

enum sshnc_retval sshnc_send(
	struct sshnc_client *client,
	const char *data,
	size_t length
) {
	clear_additional_error();

	if (client->is_closed) return SSHNC_EOF;

	size_t cursor = 0;
	while (cursor < length) {
		const int ret = ssh_channel_write(client->sesdata->channel, data + cursor, length - cursor);

		if (client->sesdata->should_close) {
			return client->sesdata->close_reason;
		}

		if (ret == SSH_ERROR) {
			if (ssh_channel_is_closed(client->sesdata->channel)
					|| !ssh_is_connected(client->sesdata->session)) {
				return SSHNC_EOF;
			} else {
				store_additional_error(client->sesdata->session);
				return SSHNC_ERR_WRITE;
			}
		}

		cursor += ret;
	}

	return SSHNC_OK;
}

enum sshnc_retval sshnc_maybe_recv(
	struct sshnc_client *client,
	size_t capacity,
	char *data,  // output
	size_t *length  // output
) {
	clear_additional_error();

	*length = 0;  // in case we error along the way

	if (client->is_closed) return SSHNC_EOF;

	const int ret = ssh_event_dopoll(client->event, -1);
	if (ret == SSH_ERROR) {
		return SSHNC_ERR_POLL;
	}

	const int status = ssh_get_status(client->sesdata->session);
	if (client->sesdata->should_close
			|| (status & (SSH_CLOSED | SSH_CLOSED_ERROR))) {
		sshnc_close_nofree(client);
		return SSHNC_EOF;
	}

	if (client->sesdata->recvlen == 0) {
		return SSHNC_AGAIN;
	}

	struct session_data *const ses = client->sesdata;
	const size_t consumed = min_size_t(ses->recvlen, capacity);

	memcpy(data, ses->recvdata, consumed);
	*length = consumed;

	memmove(ses->recvdata, ses->recvdata + consumed, ses->recvlen - consumed);
	ses->recvlen -= consumed;

	return SSHNC_OK;
}