#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <stdatomic.h>
#include <string.h>
#include <ctype.h>
#include <errno.h>
#include <assert.h>
#include <pthread.h>
#include <sys/select.h>
#include <libssh/server.h>
#include <libssh/callbacks.h>
#include "../global.h"


#define CHECK(obj_, expr_) do { \
		if (!(expr_)) { \
			fprintf(stderr, "libssh error! expression: " #expr_ "\nError description: %s\n", \
					ssh_get_error((obj_))); \
			exit(1); \
		} \
	} while (0)

void xxd(FILE *stream, const void *buf_, size_t length) {
	unsigned char *buf = (unsigned char*)buf_;

	for (size_t cursor = 0; cursor < length;) {
		fprintf(stream, "%08zx:", cursor);

		for (int i = 0; i < 16; i++) {
			if (i % 2 == 0) fprintf(stream, " ");
			if (i % 8 == 0) fprintf(stream, " ");
			if (cursor + i < length) fprintf(stream, "%02x", (unsigned)buf[cursor + i]);
			else fprintf(stream, "  ");
		}

		fprintf(stream, "  |");

		for (int i = 0; i < 16 && cursor + i < length; i++) {
			if (isprint(buf[cursor + i])) fprintf(stream, "%c", buf[cursor + i]);
			else fprintf(stream, ".");
		}

		fprintf(stream, "|\n");

		cursor += 16;
	}
}

// struct sessions {
//     // Always NULL-terminated
//     ssh_session *list;
//     ssh_session *outlist;  // same length as 'list', contents not managed
//     size_t cap, len;
// };

// static struct sessions sessions_make() {
//     size_t cap = 2;
//     return (struct sessions) {
//         .list = malloc(cap, ssh_session),
//         .outlist = malloc(cap, ssh_session),
//         .cap = cap,
//         .len = 0,
//     };
// }

// static void sessions_add(struct sessions *ss, ssh_session ses) {
//     if (ss->len + 1 >= ss->cap) {
//         ss->cap *= 2;
//         ss->list = realloc(ss->list, ss->cap, ssh_session);
//         ss->outlist = realloc(ss->outlist, ss->cap, ssh_session);
//     }
//     ss->list[ss->len++] = ses;
//     ss->list[ss->len] = NULL;
// }

// static void sessions_remove(struct sessions *ss, size_t index) {
//     assert(0 <= index && index < ss->len);
//     if (index < ss->len - 1) ss->list[index] = ss->list[ss->len - 1];
//     ss->len--;
//     ss->list[ss->len] = NULL;
// }

static atomic_int g_thread_count;

struct thread_data {
	int thread_id;
	ssh_session session;
	ssh_channel channel;  // NULL before channel has been opened
	bool should_close;

	struct ssh_server_callbacks_struct server_cb;
	struct ssh_channel_callbacks_struct chan_cb;
};

///////// CHANNEL CALLBACKS //////////

int channel_subsystem_request_cb(ssh_session session, ssh_channel channel, const char *subsystem, void *tdata_) {
	(void)session;
	(void)channel;
	struct thread_data *tdata = (struct thread_data*)tdata_;
	if (strcmp(subsystem, "tomsg") == 0) {
		printf("[%d] subsystem request: <%s>, allowing\n", tdata->thread_id, subsystem);
		return 0;
	} else {
		printf("[%d] subsystem request: <%s>, denying!\n", tdata->thread_id, subsystem);
		return 1;
	}
}

void channel_close_cb(ssh_session session, ssh_channel channel, void *tdata_) {
	(void)session; (void)channel;
	struct thread_data *tdata = (struct thread_data*)tdata_;
	printf("[%d] channel close!\n", tdata->thread_id);
}

int channel_shell_request_cb(ssh_session session, ssh_channel channel, void *tdata_) {
	(void)session; (void)channel;
	struct thread_data *tdata = (struct thread_data*)tdata_;
	printf("[%d] shell request, denying\n", tdata->thread_id);
	return 1;
}

void channel_eof_cb(ssh_session session, ssh_channel channel, void *tdata_) {
	(void)session; (void)channel;
	struct thread_data *tdata = (struct thread_data*)tdata_;
	printf("[%d] eof on channel, setting close flag\n", tdata->thread_id);
	tdata->should_close = true;
}

int channel_data_cb(ssh_session session, ssh_channel channel, void *data, uint32_t len, int is_stderr, void *tdata_) {
	(void)is_stderr; (void)data; (void)channel; (void)session;
	struct thread_data *tdata = (struct thread_data*)tdata_;
	printf("[%d] data on channel (length %u):\n", tdata->thread_id, len);
	xxd(stdout, data, len);
	printf("[%d] echoing back!\n", tdata->thread_id);
	if (ssh_channel_write(channel, data, len) == SSH_ERROR) {
		printf("[%d] write to channel failed! Setting close flag\n", tdata->thread_id);
		tdata->should_close = true;
	}
	return len;
}

void channel_signal_cb(ssh_session session, ssh_channel channel, const char *signal, void *tdata_) {
	(void)channel; (void)session;
	printf("[%d] signal SIG%s\n", ((struct thread_data*)tdata_)->thread_id, signal);
}

void channel_exit_status_cb(ssh_session session, ssh_channel channel, int exit_status, void *tdata_) {
	(void)channel; (void)session;
	printf("[%d] exit status %d\n", ((struct thread_data*)tdata_)->thread_id, exit_status);
}

void channel_exit_signal_cb(ssh_session session, ssh_channel channel, const char *signal, int core, const char *errmsg, const char *lang, void *tdata_) {
	(void)lang; (void)errmsg; (void)core; (void)channel; (void)session;
	printf("[%d] exit signal %s\n", ((struct thread_data*)tdata_)->thread_id, signal);
}

int channel_pty_request_cb(ssh_session session, ssh_channel channel, const char *term, int width, int height, int pxwidth, int pwheight, void *tdata_) {
	(void)pwheight; (void)pxwidth; (void)channel; (void)session;
	printf("[%d] pty request (term %s, %dx%d), denying\n", ((struct thread_data*)tdata_)->thread_id, term, width, height);
	return -1;
}

void channel_auth_agent_req_cb(ssh_session session, ssh_channel channel, void *tdata_) {
	(void)channel; (void)session;
	printf("[%d] auth agent request\n", ((struct thread_data*)tdata_)->thread_id);
}

void channel_x11_req_cb(ssh_session session, ssh_channel channel, int single_connection, const char *auth_protocol, const char *auth_cookie, uint32_t screen_number, void *tdata_) {
	(void)screen_number; (void)auth_cookie; (void)auth_protocol; (void)single_connection; (void)channel; (void)session;
	printf("[%d] X11 REQUEST WTF\n", ((struct thread_data*)tdata_)->thread_id);
}

int channel_pty_window_change_cb(ssh_session session, ssh_channel channel, int width, int height, int pxwidth, int pwheight, void *tdata_) {
	(void)pwheight; (void)pxwidth; (void)channel; (void)session;
	printf("[%d] pty window change (%dx%d), denying\n", ((struct thread_data*)tdata_)->thread_id, width, height);
	return -1;
}

int channel_exec_request_cb(ssh_session session, ssh_channel channel, const char *command, void *tdata_) {
	(void)channel; (void)session;
	printf("[%d] exec request (<%s>), denying\n", ((struct thread_data*)tdata_)->thread_id, command);
	return 1;
}

int channel_env_request_cb(ssh_session session, ssh_channel channel, const char *env_name, const char *env_value, void *tdata_) {
	(void)channel; (void)session;
	printf("[%d] environment request (<%s> = <%s>), denying\n", ((struct thread_data*)tdata_)->thread_id, env_name, env_value);
	return 1;
}

int channel_write_wontblock_cb(ssh_session session, ssh_channel channel, size_t bytes, void *tdata_) {
	(void)channel; (void)session;
	printf("[%d] write won't block for %zu bytes notification\n", ((struct thread_data*)tdata_)->thread_id, bytes);
	return 0;
}

////////// SERVER CALLBACKS //////////

int auth_none_cb(ssh_session session, const char *user, void *tdata_) {
	(void)session;
	struct thread_data *tdata = (struct thread_data*)tdata_;
	printf("[%d] auth none (user <%s>), accepting\n", tdata->thread_id, user);
	return SSH_AUTH_SUCCESS;
}

int auth_password_cb(ssh_session session, const char *user, const char *password, void *tdata_) {
	(void)session;
	struct thread_data *tdata = (struct thread_data*)tdata_;
	printf("[%d] auth password (user <%s> password <%s>), denying\n", tdata->thread_id, user, password);
	return SSH_AUTH_DENIED;
}

int auth_gssapi_mic_cb(ssh_session session, const char *user, const char *principal, void *tdata_) {
	(void)session;
	struct thread_data *tdata = (struct thread_data*)tdata_;
	printf("[%d] auth gssapi (user <%s> principal <%s>), denying\n", tdata->thread_id, user, principal);
	return SSH_AUTH_DENIED;
}

int auth_pubkey_cb(ssh_session session, const char *user, struct ssh_key_struct *pubkey, char signature_state, void *tdata_) {
	(void)session; (void)pubkey; (void)signature_state;
	struct thread_data *tdata = (struct thread_data*)tdata_;
	printf("[%d] auth pubkey (user <%s>), denying\n", tdata->thread_id, user);
	return SSH_AUTH_DENIED;
}

int service_request_cb(ssh_session session, const char *service, void *tdata_) {
	(void)session;
	struct thread_data *tdata = (struct thread_data*)tdata_;
	if (strcmp(service, "ssh-userauth") == 0) {
		printf("[%d] ssh-userauth service request, allowing through\n", tdata->thread_id);
		return 0;
	} else {
		printf("[%d] service request <%s>, not allowing\n", tdata->thread_id, service);
		return -1;
	}
}

ssh_channel chan_open_request_cb(ssh_session session, void *tdata_) {
	struct thread_data *tdata = (struct thread_data*)tdata_;
	if (tdata->channel == NULL) {
		ssh_channel chan = ssh_channel_new(session);
		if (chan != NULL) {
			if (ssh_set_channel_callbacks(chan, &tdata->chan_cb) == SSH_OK) {
				printf("[%d] channel open request, allowing\n", tdata->thread_id);
				tdata->channel = chan;
				return chan;
			}
			ssh_channel_close(chan);
			ssh_channel_free(chan);
		}
	}
	printf("[%d] channel open request, denying!\n", tdata->thread_id);
	return NULL;
}

// int msg_callback(ssh_session session, ssh_message msg, void *tdata_) {
//     (void)session;
//     struct thread_data *tdata = (struct thread_data*)tdata_;
//     const int tid = tdata->thread_id;

//     const int subtype = ssh_message_subtype(msg);

//     switch (ssh_message_type(msg)) {
//         case SSH_REQUEST_AUTH:
//             printf("[%d] message callback: type auth (subtype %d)\n", tid, subtype);
//             if (subtype == SSH_AUTH_METHOD_NONE) {
//                 if (ssh_message_auth_reply_success(msg, false) == SSH_OK) {
//                     return 0;  // handled
//                 } else {
//                     printf("[%d]   failed to reply success for auth method none\n", tid);
//                 }
//             }
//             break;

//         case SSH_REQUEST_CHANNEL_OPEN:
//             printf("[%d] message callback: type channel open (subtype %d)\n", tid, subtype);
//             if (tdata->channel == NULL && subtype == SSH_CHANNEL_SESSION) {
//                 ssh_channel chan = ssh_message_channel_request_open_reply_accept(msg);
//                 chan_cb.userdata = tdata;

//                 if (chan && ssh_set_channel_callbacks(chan, &chan_cb) == SSH_OK) {
//                     tdata->channel = chan;
//                     return 0;  // handled
//                 } else {
//                     if (chan) {
//                         ssh_channel_close(chan);
//                         ssh_channel_free(chan);
//                     }
//                     printf("[%d]   failed to accept channel open request for session\n", tid);
//                 }
//             }
//             break;

//         case SSH_REQUEST_CHANNEL:
//             printf("[%d] message callback: type channel (subtype %d)\n", tid, subtype);
//             break;

//         case SSH_REQUEST_SERVICE:
//             printf("[%d] message callback: type service\n", tid);
//             break;

//         case SSH_REQUEST_GLOBAL:
//             printf("[%d] message callback: type global (subtype %d)\n", tid, subtype);
//             break;
//     }

//     return 1;  // not handled
// }

void* thread_entry(void *tdata_) {
	struct thread_data *tdata = (struct thread_data*)tdata_;
	const int tid = tdata->thread_id;
	const ssh_session session = tdata->session;
	ssh_event event = NULL;

	printf("[%d] Thread started\n", tid);

	// ssh_set_message_callback(session, msg_callback, tdata);

	memset(&tdata->server_cb, 0, sizeof tdata->server_cb);
	ssh_callbacks_init(&tdata->server_cb);
	tdata->server_cb.userdata = tdata;
	tdata->server_cb.auth_none_function = auth_none_cb;
	tdata->server_cb.channel_open_request_session_function = chan_open_request_cb;
	tdata->server_cb.auth_password_function = auth_password_cb;
	tdata->server_cb.auth_gssapi_mic_function = auth_gssapi_mic_cb;
	tdata->server_cb.auth_pubkey_function = auth_pubkey_cb;
	tdata->server_cb.gssapi_select_oid_function = (void*)0x424242;  // just crash if it attemps to invoke these
	tdata->server_cb.gssapi_accept_sec_ctx_function = (void*)0x424242;
	tdata->server_cb.gssapi_verify_mic_function = (void*)0x424242;
	tdata->server_cb.service_request_function = service_request_cb;

	memset(&tdata->chan_cb, 0, sizeof tdata->chan_cb);
	ssh_callbacks_init(&tdata->chan_cb);
	tdata->chan_cb.userdata = tdata;
	tdata->chan_cb.channel_subsystem_request_function = channel_subsystem_request_cb;
	tdata->chan_cb.channel_close_function = channel_close_cb;
	tdata->chan_cb.channel_shell_request_function = channel_shell_request_cb;
	tdata->chan_cb.channel_eof_function = channel_eof_cb;
	tdata->chan_cb.channel_data_function = channel_data_cb;
	tdata->chan_cb.channel_signal_function = channel_signal_cb;
	tdata->chan_cb.channel_exit_status_function = channel_exit_status_cb;
	tdata->chan_cb.channel_exit_signal_function = channel_exit_signal_cb;
	tdata->chan_cb.channel_pty_request_function = channel_pty_request_cb;
	tdata->chan_cb.channel_auth_agent_req_function = channel_auth_agent_req_cb;
	tdata->chan_cb.channel_x11_req_function = channel_x11_req_cb;
	tdata->chan_cb.channel_pty_window_change_function = channel_pty_window_change_cb;
	tdata->chan_cb.channel_exec_request_function = channel_exec_request_cb;
	tdata->chan_cb.channel_env_request_function = channel_env_request_cb;
	tdata->chan_cb.channel_write_wontblock_function = channel_write_wontblock_cb;

	if (ssh_set_server_callbacks(session, &tdata->server_cb) != SSH_OK) {
		printf("[%d] Failed setting server callbacks: %s\n", tid, ssh_get_error(session));
		goto cleanup;
	}

	ssh_set_auth_methods(session, SSH_AUTH_METHOD_NONE);

	if (ssh_handle_key_exchange(session) != SSH_OK) {
		printf("[%d] Key exchange failed: %s\n", tid, ssh_get_error(session));
		goto cleanup;
	}
	printf("[%d] Handled key exchange\n", tid);

	event = ssh_event_new();
	if (!event || ssh_event_add_session(event, session) != SSH_OK) {
		printf("[%d] Failed to create ssh event context\n", tid);
		goto cleanup;
	}

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

cleanup:
	if (event) ssh_event_free(event);
	if (tdata->channel) {
		ssh_channel_close(tdata->channel);
		ssh_channel_free(tdata->channel);
	}
	if (session) {
		ssh_disconnect(session);
		ssh_free(session);
		printf("[%d] Disconnected\n", tid);
	}
	free(tdata);
	int num_threads = atomic_fetch_sub(&g_thread_count, 1);
	printf("[%d] Exiting! (%d threads remaining)\n", tid, num_threads - 1);
	pthread_exit(NULL);
}

int main(void) {
	if (ssh_init() != SSH_OK) {
		fprintf(stderr, "Could not initialise libssh\n");
		return 1;
	}

	ssh_key host_key;
	if (ssh_pki_import_privkey_file("host_key", NULL, NULL, NULL, &host_key) != SSH_OK) {
		fprintf(stderr, "Failed to read host private key file 'host_key'\n");
		return 1;
	}

	size_t host_key_hash_length = 0;
	unsigned char *host_key_hash = NULL;
	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");
		return 1;
	}

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

	ssh_bind srvbind = ssh_bind_new();
	CHECK(srvbind, srvbind);

	bool procconfig = false;
	CHECK(srvbind, ssh_bind_options_set(srvbind, SSH_BIND_OPTIONS_PROCESS_CONFIG, &procconfig) == SSH_OK);
	int port = 2222;
	CHECK(srvbind, ssh_bind_options_set(srvbind, SSH_BIND_OPTIONS_BINDPORT, &port) == SSH_OK);
	CHECK(srvbind, ssh_bind_options_set(srvbind, SSH_BIND_OPTIONS_IMPORT_KEY, host_key) == SSH_OK);
	const char *ciphers_str = "aes256-gcm@openssh.com,aes256-ctr,aes256-cbc";
	CHECK(srvbind, ssh_bind_options_set(srvbind, SSH_BIND_OPTIONS_CIPHERS_C_S, ciphers_str) == SSH_OK);
	CHECK(srvbind, ssh_bind_options_set(srvbind, SSH_BIND_OPTIONS_CIPHERS_S_C, ciphers_str) == SSH_OK);

	CHECK(srvbind, ssh_bind_listen(srvbind) == SSH_OK);
	printf("Listening for SSH connections on port %d\n", port);

	// int srvbind_fd = ssh_bind_get_fd(srvbind);

	// struct sessions sessions = sessions_make();

	pthread_attr_t thread_attrs;
	assert(pthread_attr_init(&thread_attrs) == 0);
	assert(pthread_attr_setdetachstate(&thread_attrs, PTHREAD_CREATE_DETACHED) == 0);

	int next_thread_id = 0;
	atomic_store(&g_thread_count, 0);

	while (true) {
		// fd_set inset;
		// FD_ZERO(&inset);
		// FD_SET(srvbind_fd, &inset);

		// int ret = ssh_select(sessions.list, sessions.outlist, srvbind_fd + 1, &inset, NULL);
		// if (ret == SSH_EINTR) continue;
		// if (ret == SSH_ERROR) {
		//     fprintf(stderr, "ssh_select reported error!\n");
		//     return 1;
		// }

		ssh_session session = ssh_new();
		CHECK(session, session);

		CHECK(srvbind, ssh_bind_accept(srvbind, session) == SSH_OK);
		int num_existing_threads = atomic_fetch_add(&g_thread_count, 1);
		printf("Accepted connection, spinning up thread (currently %d threads)\n",
				num_existing_threads + 1);

		struct thread_data *tdata = calloc(1, struct thread_data);
		tdata->thread_id = next_thread_id++;
		tdata->session = session;
		tdata->channel = NULL;
		tdata->should_close = false;

		pthread_t thread;
		assert(pthread_create(&thread, &thread_attrs, thread_entry, tdata) == 0);
	}
}