/**
 *  @file
 *  @author     Pawel Zubrycki <paw.zubr@gmail.com>
 *  @author     2013 Stefan Radomski (stefan.radomski@cs.tu-darmstadt.de)
 *  @copyright  Simplified BSD
 *
 *  @cond
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the FreeBSD license as published by the FreeBSD
 *  project.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 *
 *  You should have received a copy of the FreeBSD license along with this
 *  program. If not, see <http://www.opensource.org/licenses/bsd-license>.
 *  @endcond
 */

#include "evws.h"
#include <stdlib.h>
#include <string.h>
#include <event2/buffer.h>
#include <event2/bufferevent.h>
#include "uscxml/util/SHA1.h"
#include "uscxml/util/Base64.h"

static int evws_parse_first_line(struct evws_connection *conn, char *line);
static int evws_parse_header_line(char *line, char **skey, char **svalue);

// Callbacks
static void cb_accept(struct evconnlistener *listener, evutil_socket_t fd, struct sockaddr *address, int socklen, void *ctx); 
static void cb_read_handshake(struct bufferevent *bev, void *arg);
static void cb_read_frame(struct bufferevent *bev, void *arg);

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

struct evws *evws_new(struct event_base *base) {
	struct evws *ret_obj = (struct evws*)calloc(1, sizeof(struct evws));
	ret_obj->base = base;
	ret_obj->listener = NULL;

	TAILQ_INIT(&ret_obj->connections);
	TAILQ_INIT(&ret_obj->callbacks);

	return ret_obj;
}

void evws_free(struct evws *ptr) {
	// Tu wiecej czyszczenia
	free(ptr);
}

evutil_socket_t evws_bind_socket(struct evws * ws, unsigned short port) {
	struct sockaddr_in sin;

	// Creating serverside socket
	memset(&sin, 0, sizeof(sin));
	sin.sin_family = AF_INET;
	sin.sin_addr.s_addr = htonl(0);
	sin.sin_port = htons(port);

	if(!(ws->listener = evconnlistener_new_bind(ws->base, cb_accept, ws, LEV_OPT_CLOSE_ON_FREE|LEV_OPT_REUSEABLE|LEV_OPT_THREADSAFE, -1, (struct sockaddr*)&sin, sizeof(sin)))) {
		return 0;
	}
	return evconnlistener_get_fd(ws->listener);
}

int evws_set_cb(struct evws * ws, const char * uri, cb_frame_type message_cb, cb_type connect_cb, void * arg) {
	struct evws_cb *ws_cb;

	TAILQ_FOREACH(ws_cb, &ws->callbacks, next) {
		if (strcmp(ws_cb->uri, uri) == 0)
			return (-1);
	}

	if((ws_cb = (struct evws_cb*)calloc(1, sizeof(struct evws_cb))) == NULL) {
		return (-2);
	}

	ws_cb->uri = (char*)strdup(uri);
	ws_cb->msg_cb = message_cb;
	ws_cb->conn_cb = connect_cb;
	ws_cb->cb_arg = arg;

	TAILQ_INSERT_TAIL(&ws->callbacks, ws_cb, next);

	return (0);
}

cb_frame_type evws_set_gencb(struct evws *ws, cb_frame_type cb, void * arg) {
	cb_frame_type old_cb = ws->gencb;
	ws->gencb = cb;
	ws->gencb_arg = arg;
	return old_cb;
}

// Broadcast data to all buffers associated with pattern
void evws_broadcast(struct evws *ws, const char *uri, enum evws_opcode opcode, const char *data, uint64_t length) {
	struct evws_connection *ws_connection;
	TAILQ_FOREACH(ws_connection, &ws->connections, next) {
		if (strcmp(ws_connection->uri, uri) == 0)
			evws_send_data(ws_connection, opcode, data, length);
	}
}

// Error callback
static void cb_error(struct bufferevent *bev, short what, void *ctx) {
	struct evws_connection *conn = ctx;
	TAILQ_REMOVE(&(conn->ws->connections), conn, next);
	evws_connection_free(conn);
}

//Callback to accept
static void cb_accept(struct evconnlistener *listener, evutil_socket_t fd, struct sockaddr *address, int socklen, void *arg) {
	struct evws *ws = arg;
	struct evws_connection *ws_conn = evws_connection_new(ws, fd);
	
	bufferevent_setcb(ws_conn->bufev, cb_read_handshake, NULL, cb_error, ws_conn);
	bufferevent_enable(ws_conn->bufev, EV_READ);
}

int evws_parse_first_line(struct evws_connection *conn, char *line) {
	char *method;
	char *uri;
	char *version;

	/* Parse the request line */
	method = strsep(&line, " ");
	if (line == NULL)
		return (-1);
	uri = strsep(&line, " ");
	if (line == NULL)
		return (-1);
	version = strsep(&line, " ");
	if (line != NULL)
		return (-1);

	if ((conn->uri = strdup(uri)) == NULL) {
		return (-1);
	}

	return (0);
}

// Callback to read handshake
void cb_read_handshake(struct bufferevent *bev, void *arg) {
	struct evws_connection *ws_conn = arg;
	char *line, *skey, *svalue;
	struct evbuffer *buffer = bufferevent_get_input(bev);
	size_t line_length;
	char *key = NULL;
	char *host = NULL;
	char *origin = NULL;

	switch(ws_conn->state) {
	case 0:
		line = evbuffer_readln(buffer, &line_length, EVBUFFER_EOL_CRLF);
		evws_parse_first_line(ws_conn, line);
		ws_conn->state = EVWS_FIRSTLINE_READ;
		free(line);
	case EVWS_FIRSTLINE_READ:
		while ((line = evbuffer_readln(buffer, &line_length, EVBUFFER_EOL_CRLF)) != NULL) {
			if (*line == '\0') { /* Last header - Done */
				free(line);
				ws_conn->state = EVWS_HEADER_READ;
				break;
			}
			evws_parse_header_line(line, &skey, &svalue);
			if(strcmp(skey, "Sec-WebSocket-Key") == 0) {
				key = strdup(svalue);
			} else if(strcmp(skey, "Host") == 0) {
				host = strdup(svalue);
			} else if(strcmp(skey, "Origin") == 0) {
				origin = strdup(svalue);
			}
			struct evws_header *header = evws_header_new(skey, svalue);
			TAILQ_INSERT_TAIL(&ws_conn->headers, header, next);
			free(line);
		}
	default:
		break;
	};

	// -- SHA1
	
	SHA1Context sha1;
	SHA1Reset(&sha1);
	SHA1Input(&sha1, (const unsigned char*)key, 24);
	SHA1Input(&sha1, (const unsigned char*)WS_GUID, 36);
	SHA1Result(&sha1);
	
	char chksumSha1[21];
  int i;
	for (i = 0; i < 5; i++) {
		chksumSha1[i * 4 + 0] = (sha1.Message_Digest[i] >> 24) & 0xff;
		chksumSha1[i * 4 + 1] = (sha1.Message_Digest[i] >> 16) & 0xff;
		chksumSha1[i * 4 + 2] = (sha1.Message_Digest[i] >> 8) & 0xff;
		chksumSha1[i * 4 + 3] = (sha1.Message_Digest[i] >> 0) & 0xff;
	}
//	printf("%s\n", chksumSha1);

	// -- BASE64
	
	int md5End = 0;
	char chksumBase64[200];
	base64_encodestate* base64Ctx = malloc(sizeof(base64_encodestate));
	base64_init_encodestate(base64Ctx);
	md5End += base64_encode_block(chksumSha1, 20, chksumBase64, base64Ctx);
	md5End += base64_encode_blockend(&chksumBase64[md5End], base64Ctx);
	// blockend writes unneccessary \n
	chksumBase64[md5End - 1] = 0;

	free(base64Ctx);
	free(key);
	free(host);
	free(origin);
	evbuffer_add_printf(bufferevent_get_output(ws_conn->bufev),
											"HTTP/1.1 101 Web Socket Protocol Handshake\r\n"
											"Upgrade: WebSocket\r\n"
											"Connection: Upgrade\r\n"
											"Sec-WebSocket-Accept: %s\r\n"
											"Server: uSCXML\r\n"
											"Access-Control-Allow-Origin: *\r\n"
											"Access-Control-Allow-Credentials: true\r\n"
											"Access-Control-Allow-Headers: content-type\r\n"
											"Access-Control-Allow-Headers: authorization\r\n"
											"Access-Control-Allow-Headers: x-websocket-extensions\r\n"
											"Access-Control-Allow-Headers: x-websocket-version\r\n"
											"Access-Control-Allow-Headers: x-websocket-protocol\r\n"
											"\r\n",
											chksumBase64
	);
	bufferevent_setcb(ws_conn->bufev, cb_read_frame, NULL, cb_error, ws_conn);

	TAILQ_INSERT_TAIL(&(ws_conn->ws->connections), ws_conn, next);
	{
		struct evws_cb *ws_cb;
		TAILQ_FOREACH(ws_cb, &ws_conn->ws->callbacks, next) {
			if (strcmp(ws_cb->uri, ws_conn->uri) == 0) {
				if(ws_cb->conn_cb != NULL)
					ws_cb->conn_cb(ws_conn, NULL, 0, ws_cb->cb_arg);
				return;
			}
		}
	}
}

int evws_parse_header_line(char *line, char **skey, char **svalue)
{
	*svalue = line;
	*skey = strsep(svalue, ":");
	if (*svalue == NULL)
		return -1;

	*svalue += strspn(*svalue, " ");

	return (0);
}

// Callback to read sent data
void cb_read_frame(struct bufferevent *bev, void *arg)
{
	struct evws_connection *conn = arg;
	struct evws *ws = conn->ws;
	char readbuf[1024];
	int size = 0;
	struct evbuffer *buffer = bufferevent_get_input(bev);
	while ((size = evbuffer_remove(buffer, readbuf, sizeof(readbuf))) > 0) {
		char* dataPtr = readbuf;

NEXT_FRAME:
		if (conn->frame == NULL) {
			// data starts with the first byte of a new frame
			conn->frame = evws_frame_new();
			conn->frame->fin = (dataPtr[0] >> 7) & 1;
			conn->frame->rsv = (dataPtr[0] >> 4) & 7;
			conn->frame->opcode = dataPtr[0] & 0xfu;
			conn->frame->state = EVWS_OPCODE_READ;
			dataPtr++;
		}

		if (size - (dataPtr - readbuf) == 0)
			continue;
		
		if (conn->frame->state == EVWS_OPCODE_READ) { // we already read the first byte
			conn->frame->has_mask = (dataPtr[0] >> 7) & 1;
			if (conn->frame->has_mask)
				conn->frame->mask_bytes = 4;
			conn->frame->size = dataPtr[0] & 0x7fu;
			dataPtr++;
			
			if(conn->frame->size == 126) {
				conn->frame->payload_bytes = 2;
				conn->frame->size = 0;
			} else if(conn->frame->size == 127) {
				conn->frame->payload_bytes = 8;
				conn->frame->size = 0;
			} else {
				conn->frame->payload_bytes = 0;
			}
			conn->frame->state = EVWS_PAYLOAD_LENGTH_LENGTH_READ;
		}
		
		if (size - (dataPtr - readbuf) == 0)
			continue;
		
		if (conn->frame->state == EVWS_PAYLOAD_LENGTH_LENGTH_READ) { // we already read the first byte
			
			while(conn->frame->payload_bytes > 0 && // need to add more bytes to the payload length
						(size - (dataPtr - readbuf) > 0)) { // and bytes still available
				// length is in MSB network order - shift to left and add
				conn->frame->size = (conn->frame->size << 8) + dataPtr[0];
				conn->frame->payload_bytes--;
				dataPtr++;
			}
			if (conn->frame->payload_bytes == 0) {
				conn->frame->state = EVWS_PAYLOAD_LENGTH_READ;
				if (!conn->frame->has_mask)
					conn->frame->state = EVWS_MASKING_KEY_READ;
			}
			conn->frame->data = (char*)malloc(conn->frame->size);
		}
		
		if (size - (dataPtr - readbuf) == 0)
			continue;
		
		if (conn->frame->state == EVWS_PAYLOAD_LENGTH_READ && // we already read the complete payload
				conn->frame->has_mask) {
			while (conn->frame->mask_bytes > 0 && // bytes for the frame mask
						 (size - (dataPtr - readbuf) > 0)) { // still some available
				conn->frame->mask[4 - conn->frame->mask_bytes] = dataPtr[0];
				conn->frame->mask_bytes--;
				dataPtr++;
				if (conn->frame->mask_bytes == 0)
					conn->frame->state = EVWS_MASKING_KEY_READ;
			}
		}
		
		if (size - (dataPtr - readbuf) == 0)
			continue;
		
		if (conn->frame->state == EVWS_MASKING_KEY_READ) { // we read all of the header
			size_t remaining = size - (dataPtr - readbuf);
			remaining = (remaining > conn->frame->size - conn->frame->payload_read ? conn->frame->size - conn->frame->payload_read : remaining);
			
			if (conn->frame->has_mask) {
        int i;
				for (i = 0; i < remaining; i++) {
					dataPtr[i] = dataPtr[i] ^ conn->frame->mask[conn->frame->mask_index];
					conn->frame->mask_index++;
					conn->frame->mask_index = conn->frame->mask_index % 4;
				}
			}
			memcpy(conn->frame->data + conn->frame->payload_read, dataPtr, remaining);
			conn->frame->payload_read += remaining;
			dataPtr += remaining;
		}

		if (conn->frame->payload_read == conn->frame->size) {
			// done reading this frame - invoke callbacks
			struct evws_cb *ws_cb;
			TAILQ_FOREACH(ws_cb, &ws->callbacks, next) {
				if (strcmp(ws_cb->uri, conn->uri) == 0) {
					if(ws_cb->msg_cb != NULL)
						ws_cb->msg_cb(conn, conn->frame, ws_cb->cb_arg);
					continue;
				}
			}
			ws->gencb(conn, conn->frame, ws->gencb_arg);
			
			evws_frame_free(conn->frame);
			conn->frame = NULL;
			if (size - (dataPtr - readbuf) != 0) // there is more data in this packet
				goto NEXT_FRAME;
		}
	}
}

void evws_send_data(struct evws_connection *conn, enum evws_opcode opcode, const char *data, uint64_t length) {
	char *sendbuf = malloc(length + 10); // payload + header without masking key
	if(sendbuf == NULL)
		return;
	
	char* writePtr = sendbuf;
	
	writePtr[0] = (1 << 7); // set fin header and zero out RSV
	writePtr[0] += opcode; // set opcode
	writePtr++;
	writePtr[0] = (0 << 7); // we don't mask replies
	
	if (length < 126) {
		writePtr[0] += (uint8_t)length;
		writePtr++;
	} else if(length < (1 << 16)) {
		writePtr[0] = 126;
		writePtr[1] = (((uint16_t)length) >> 8) & 0xff;
		writePtr[2] = ((uint16_t)length) & 0xff;
		writePtr += 3;
	} else if(length < (1ull << 63)) {
		writePtr[0] = 126;
		// integer division and bitmask ought to be endian agnostic
		writePtr[1] = (length / 0x0100000000000000) & 0xff;
		writePtr[2] = (length / 0x0001000000000000) & 0xff;
		writePtr[3] = (length / 0x0000010000000000) & 0xff;
		writePtr[4] = (length / 0x0000000100000000) & 0xff;
		writePtr[5] = (length / 0x0000000001000000) & 0xff;
		writePtr[6] = (length / 0x0000000000010000) & 0xff;
		writePtr[7] = (length / 0x0000000000000100) & 0xff;
		writePtr[8] = (length / 0x0000000000000001) & 0xff;
		writePtr += 9;
	}
	
	memcpy(writePtr, data, length);
	writePtr += length;
	
	struct evbuffer *buffer =  bufferevent_get_output(conn->bufev);
	evbuffer_add(buffer, sendbuf, writePtr - sendbuf);
	free(sendbuf);
}


// --- new and free pairs ---

struct evws_connection* evws_connection_new(struct evws *ws, evutil_socket_t fd) {
	struct evws_connection* conn = calloc(1, sizeof(struct evws_connection));
	conn->ws = ws;
	conn->fd = fd;
	conn->uri = NULL;

	conn->bufev = bufferevent_socket_new(ws->base, fd, BEV_OPT_CLOSE_ON_FREE);
	conn->state = 0;
	conn->frame = NULL;
	TAILQ_INIT(&conn->headers);
	return conn;
}

void evws_connection_free(struct evws_connection *conn) {
	struct evws_header *header;
	bufferevent_free(conn->bufev);
	if(conn->uri != NULL)
		free(conn->uri);

	TAILQ_FOREACH(header, &conn->headers, next) {
		evws_header_free(header);
	}
	free(conn);
}

struct evws_frame *evws_frame_new() {
	struct evws_frame *frame = calloc(1, sizeof(struct evws_frame));
	return frame;
}

void evws_frame_free(struct evws_frame *frame) {
	if (frame->data)
		free(frame->data);
	free(frame);
}

struct evws_header *evws_header_new(char *key, char *value)
{
	struct evws_header *head = calloc(1, sizeof(struct evws_header));
	head->key = strdup(key);
	head->value = strdup(value);
	return head;
}

void evws_header_free(struct evws_header *header) {
	if(header->key != NULL)
		free(header->key);
	// @Note: segfault when freeing value, some strange value
	if(header->value != NULL)
		free(header->value);
	free(header);
}

char *evws_find_header(const struct wsheadersq *q, const char *key) {
	struct evws_header *hdr;
	char * ret = NULL;
	TAILQ_FOREACH(hdr, q, next) {
		if(strcmp(hdr->key, key) == 0) {
			ret = hdr->value;
			break;
		}
	}
	return ret;
}