#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <string.h>
#include <unistd.h>
#include <netdb.h>
#include <errno.h>
#include <sys/socket.h>
#include <sys/select.h>
#include <netinet/in.h>
#include "termio/termio.h"

#include <vector>
#include <string>

using namespace std;

#define PORT 1234

#define STR_(x) #x
#define STR(x) STR_(x)

__attribute__((noreturn))
static void die_perror(const char *msg) {
	perror(msg);
	exit(1);
}

static int create_server_socket(void) {
	int sock = socket(AF_INET, SOCK_STREAM, 0);
	if (sock < 0) die_perror("socket");

	int one = 1;
	setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &one, sizeof one);

	struct sockaddr_in name;
	memset(&name, 0, sizeof name);
	name.sin_family = AF_INET;
	name.sin_addr.s_addr = htonl(INADDR_ANY);
	name.sin_port = htons(PORT);
	if (::bind(sock, (struct sockaddr*)&name, sizeof name) < 0) die_perror("bind");

	if (listen(sock, 16) < 0) die_perror("listen");

	return sock;
}

static void writeall(int sock, const char *data, size_t len) {
	size_t cursor = 0;
	while (cursor < len) {
		ssize_t nw = write(sock, data + cursor, len - cursor);
		if (nw <= 0) die_perror("write");
		cursor += nw;
	}
}

static void print_buffer(const char *buffer, size_t len) {
	Size tsize = querytermsize();
	if (len + 1 >= tsize.w) {
		printf("\x1B[2K\x1B[G%s", buffer + len - tsize.w + 1);
	} else {
		printf("\x1B[2K\x1B[G%s", buffer);
	}
	fflush(stdout);
}

static void communicate(int sock) {
	printf("Connected.\n");

	string buffer;
	string inbuffer;

	initkeyboard(false);
	atexit(endkeyboard);

	while (true) {
		Size tsize = querytermsize();

		fd_set inset;
		FD_ZERO(&inset);
		FD_SET(0, &inset);
		FD_SET(sock, &inset);
		int ret = select(sock + 1, &inset, NULL, NULL, NULL);
		if (ret < 0) {
			if (errno == EINTR) continue;
			die_perror("select");
		}

		if (FD_ISSET(0, &inset)) {
			int key = tgetkey();
			if (key == -1) {  // EOF
				return;
			} else if (key == KEY_LF || key == KEY_CR) {
				buffer.push_back('\n');
				writeall(sock, buffer.data(), buffer.size());
				buffer.pop_back();
				printf("\x1B[2K\x1B[G%s\n", buffer.data());
				fflush(stdout);
				buffer.clear();
			} else if (key == KEY_BACKSPACE) {
				if (buffer.empty()) {
					bel();
				} else {
					buffer.pop_back();
					print_buffer(buffer.data(), buffer.size());
				}
			} else if (32 <= key) {
				buffer.push_back(key);
				print_buffer(buffer.data(), buffer.size());
			} else {
				bel();
			}
		}

		if (FD_ISSET(sock, &inset)) {
			size_t prevsize = inbuffer.size();
			inbuffer.resize(prevsize + 1024);

			ssize_t nr = read(sock, (void*)(inbuffer.data() + prevsize), 1024);
			if (nr < 0) die_perror("read");
			if (nr == 0) return;  // eof

			inbuffer.resize(prevsize + nr);
			for (size_t idx; (idx = inbuffer.find('\n')) != string::npos; ) {
				inbuffer[idx] = '\0';
				printf("\x1B[2K\x1B[G> \x1B[36m%s\x1B[0m\n", inbuffer.data());
				inbuffer.erase(inbuffer.begin(), inbuffer.begin() + idx + 1);
			}

			print_buffer(buffer.data(), buffer.size());
		}
	}
}

int main(int argc, char **argv) {
	if (argc == 2 && strcmp(argv[1], "-s") == 0) {
		int lsock = create_server_socket();

		printf("Listening on port %d\n", PORT);
		int sock = accept(lsock, NULL, NULL);
		if (sock == -1) die_perror("accept");

		communicate(sock);
		close(sock);
	} else if (argc == 2) {
		struct addrinfo *addrs;
		struct addrinfo hints;
		memset(&hints, 0, sizeof hints);
		hints.ai_family = AF_UNSPEC;
		hints.ai_socktype = SOCK_STREAM;
		hints.ai_protocol = IPPROTO_TCP;
		if (getaddrinfo(argv[1], STR(PORT), &hints, &addrs) < 0) die_perror("getaddrinfo");

		int sock = socket(AF_INET, SOCK_STREAM, 0);
		if (sock < 0) die_perror("socket");

		bool success = false;
		for (struct addrinfo *addr = addrs; addr; addr = addr->ai_next) {
			if (connect(sock, addr->ai_addr, addr->ai_addrlen) < 0) continue;
			success = true;
			break;
		}

		if (!success) {
			perror("connect");
			freeaddrinfo(addrs);
			exit(1);
		}

		communicate(sock);
		close(sock);
	}
}