#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <stdatomic.h>
#include <string.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)

// 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

	struct ssh_server_callbacks_struct server_cb;
	struct ssh_channel_callbacks_struct chan_cb;
};

int 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_;
	printf("[%d] subsystem request: <%s>\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\n", tdata->thread_id);
}

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

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)) {
				tdata->channel = chan;
				return chan;
			}
			ssh_channel_close(chan);
			ssh_channel_free(chan);
		}
	}
	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;

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

	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 (true) {
		// 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_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_HOSTKEY, "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;

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