#include <ctype.h>

#define RFC6455_MAGIC_KEY "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

static size_t socks = 0;
static websocket* sock = NULL;

int ws_close(websocket* ws){
	size_t p;

	if(ws->fd >= 0){
		close(ws->fd);
		ws->fd = -1;
	}

	if(ws->peer >= 0){
		close(ws->peer);
		ws->peer = -1;
	}

	for(p = 0; p < ws->protocols; p++){
		free(ws->protocol[p]);
	}
	ws->protocols = 0;
	ws->protocol = NULL;

	ws->read_buffer_offset = 0;

	free(ws->request_path);
	ws->request_path = NULL;

	free(ws->socket_key);
	ws->socket_key = NULL;

	ws->websocket_version = 0;
	ws->want_upgrade = 0;

	return 0;
}

int ws_accept(int listen_fd){
	size_t n = 0;

	websocket ws = {
		.fd = accept(listen_fd, NULL, NULL),
		.peer = -1
	};

	//try to find a slot to occupy
	for(n = 0; n < socks; n++){
		if(sock[n].fd == -1){
			break;
		}
	}

	//none found, need to extend
	if(n == socks){
		sock = realloc(sock, (socks + 1) * sizeof(websocket));
		if(!sock){
			close(ws.fd);
			fprintf(stderr, "Failed to allocate memory\n");
			return 1;
		}
		socks++;
	}

	sock[n] = ws;

	return 0;
}

int ws_handle_new(websocket* ws){
	size_t u;
	char* path, *proto;

	if(!strncmp((char*) ws->read_buffer, "GET ", 4)){
		path = (char*) ws->read_buffer + 4;
		for(u = 0; path[u] && !isspace(path[u]); u++){
		}
		path[u] = 0;
		proto = path + u + 1;
	}
	//TODO handle other methods
	else{
		fprintf(stderr, "Unknown HTTP method in request\n");
		return 1;
	}

	if(strncmp(proto, "HTTP/", 5)){
		fprintf(stderr, "Malformed HTTP initiation\n");
		return 1;
	}

	ws->state = ws_http;
	ws->request_path = strdup(path);
	return 0;
}

int ws_upgrade_http(websocket* ws){
	if(ws->websocket_version == 13
			&& ws->socket_key
			&& ws->want_upgrade){
		if(network_send_str(ws->fd, "HTTP/1.1 101 Upgrading\r\n")
				|| network_send_str(ws->fd, "Upgrade: websocket\r\n")
				|| network_send_str(ws->fd, "Connection: Upgrade\r\n")){
			//TODO might want to disconnect here
			return 1;
		}

		//calculate the websocket key which for some reason is defined as
		//base64(sha1(concat(trim(client-key), "258EAFA5-E914-47DA-95CA-C5AB0DC85B11")))
		//which requires not one but 2 unnecessarily complex operations
		size_t encode_offset = 0;
		struct sha1_ctx ws_accept_ctx;
		struct base64_encode_ctx ws_accept_encode;
		uint8_t ws_accept_digest[SHA1_DIGEST_SIZE];
		char ws_accept_key[BASE64_ENCODE_LENGTH(SHA1_DIGEST_SIZE) + 3] = "";
		sha1_init(&ws_accept_ctx);
		sha1_update(&ws_accept_ctx, strlen(ws->socket_key), (uint8_t*) ws->socket_key);
		sha1_update(&ws_accept_ctx, strlen(RFC6455_MAGIC_KEY), (uint8_t*) RFC6455_MAGIC_KEY);
		sha1_digest(&ws_accept_ctx, sizeof(ws_accept_digest), (uint8_t*) &ws_accept_digest);
		base64_encode_init(&ws_accept_encode);
		encode_offset = base64_encode_update(&ws_accept_encode, (uint8_t*) ws_accept_key, SHA1_DIGEST_SIZE, ws_accept_digest);
		encode_offset += base64_encode_final(&ws_accept_encode, (uint8_t*) ws_accept_key + encode_offset);
		memcpy(ws_accept_key + encode_offset, "\r\n\0", 3);

		//send websocket accept key
		if(network_send_str(ws->fd, "Sec-WebSocket-Accept: ")
					|| network_send_str(ws->fd, ws_accept_key)){
			return 1;
		}

		//TODO Sec-Websocket-Protocol
		//TODO find/connect peer

		ws->state = ws_open;
		return network_send_str(ws->fd, "\r\n") ? 1 : 0;
	}
	//TODO RFC 4.2.2.4: An unsupported version must be answered with HTTP 426
	return 1;
}

int ws_handle_http(websocket* ws){
	char* header, *value;
	ssize_t p;

	if(!ws->read_buffer[0]){
		return ws_upgrade_http(ws);
	}
	else if(isspace(ws->read_buffer[0])){
		//i hate header folding
		//TODO disconnect client
		return 1;
	}
	else{
		header = (char*) ws->read_buffer;
		value = strchr(header, ':');
		if(!value){
			//TODO disconnect
			return 1;
		}
		*value = 0;
		value++;
		for(; isspace(*value); value++){
		}
	}

	if(!strcmp(header, "Sec-WebSocket-Version")){
		ws->websocket_version = strtoul(value, NULL, 10);
	}
	else if(!strcmp(header, "Sec-WebSocket-Key")){
		//right-trim the key
		for(p = strlen(value) - 1; p >= 0 && isspace(value[p]); p--){
			value[p] = 0;
		}
		ws->socket_key = strdup(value);
	}
	else if(!strcmp(header, "Upgrade") && !strcmp(value, "websocket")){
		ws->want_upgrade = 1;
	}
	//TODO parse websocket protocol offers
	return 0;
}

int ws_frame(websocket* ws){
	//TODO check for complete WS frame
	//read/handle/forward data
	fprintf(stderr, "Incoming websocket data\n");
	return 0;
}

int ws_data(websocket* ws){
	ssize_t bytes_read, n, bytes_left = sizeof(ws->read_buffer) - ws->read_buffer_offset;
	int rv = 0;

	bytes_read = recv(ws->fd, ws->read_buffer + ws->read_buffer_offset, bytes_left - 1, 0);
	if(bytes_read < 0){
		fprintf(stderr, "Failed to receive from websocket: %s\n", strerror(errno));
		ws_close(ws);
		return 0;
	}
	else if(bytes_read == 0){
		//client closed connection
		ws_close(ws);
		return 0;
	}

	//terminate new data
	ws->read_buffer[ws->read_buffer_offset + bytes_read] = 0;
	
	switch(ws->state){
		case ws_new:
		case ws_http:
			//scan for newline, handle line
			for(n = 0; n < bytes_read - 1; n++){
				if(!strncmp((char*) ws->read_buffer + ws->read_buffer_offset + n, "\r\n", 2)){
					//terminate line
					ws->read_buffer[ws->read_buffer_offset + n] = 0;

					if(ws->state == ws_new){
						rv |= ws_handle_new(ws);
					}
					else{
						rv |= ws_handle_http(ws);
					}
					//TODO handle rv

					//remove from buffer
					bytes_read -= (n + 2);
					memmove(ws->read_buffer, ws->read_buffer + ws->read_buffer_offset + n + 2, bytes_read);
					ws->read_buffer_offset = 0;

					//restart from the beginning
					n = -1;
				}
			}
			break;
		case ws_open:
			rv |= ws_frame(ws);
	}

	//update read buffer offset
	ws->read_buffer_offset = bytes_read;

	//disconnect spammy clients
	if(sizeof(ws->read_buffer) - ws->read_buffer_offset < 2){
		fprintf(stderr, "Disconnecting misbehaving client\n");
		ws_close(ws);
		return 0;
	}
	return rv;
}

void ws_cleanup(){
	size_t n;
	for(n = 0; n < socks; n++){
		ws_close(sock + n);
	}

	free(sock);
	sock = NULL;
	socks = 0;
}