#include "http.h"

#include "log.h"
#include "mem.h"
#include "util.js.h"

#include "picohttpparser.h"
#include "uv.h"

#include <assert.h>
#include <stdlib.h>
#include <string.h>

#if !defined(__APPLE__) && !defined(__OpenBSD__) && !defined(_WIN32)
#include <alloca.h>
#endif

typedef struct _tf_http_connection_t
{
	tf_http_t* http;
	uv_tcp_t tcp;
	uv_shutdown_t shutdown;

	int ref_count;

	const char* method;
	const char* path;
	const char* query;
	int minor_version;

	char headers_buffer[8192];
	size_t headers_buffer_length;
	int parsed_length;

	char buffer[8192];
	size_t buffer_length;

	struct phr_header headers[32];
	int headers_length;
	bool headers_done;

	int flags;
	tf_http_callback_t* callback;
	tf_http_request_t* request;
	void* user_data;

	bool is_websocket;
	void* body;
	size_t body_length;
	size_t content_length;
	bool connection_close;
} tf_http_connection_t;

typedef struct _tf_http_handler_t
{
	const char* pattern;
	int flags;
	tf_http_callback_t* callback;
	void* user_data;
} tf_http_handler_t;

typedef struct _tf_http_t
{
	uv_tcp_t** listeners;
	int listeners_count;

	tf_http_connection_t** connections;
	int connections_count;

	tf_http_handler_t* handlers;
	int handlers_count;

	int pending_closes;
	uv_loop_t* loop;
} tf_http_t;

static const char* _http_connection_get_header(const tf_http_connection_t* connection, const char* name);
static void _http_connection_destroy(tf_http_connection_t* connection, const char* reason);
static const char* _http_status_text(int status);

tf_http_t* tf_http_create(uv_loop_t* loop)
{
	tf_http_t* http = tf_malloc(sizeof(tf_http_t));
	*http = (tf_http_t)
	{
		.loop = loop,
	};
	return http;
}

void _http_allocate_buffer(uv_handle_t* handle, size_t suggestedSize, uv_buf_t* buf)
{
	tf_http_connection_t* connection = handle->data;
	if (!connection->headers_done)
	{
		*buf = uv_buf_init(connection->headers_buffer + connection->headers_buffer_length, sizeof(connection->headers_buffer) - connection->headers_buffer_length);
	}
	else
	{
		*buf = uv_buf_init(connection->buffer + connection->buffer_length, sizeof(connection->buffer) - connection->buffer_length);
	}
}

bool _http_find_handler(tf_http_t* http, const char* path, int flags, tf_http_callback_t** out_callback, void** out_user_data)
{
	for (int i = 0; i < http->handlers_count; i++)
	{
		if (http->handlers[i].flags == flags &&
			(!http->handlers[i].pattern ||
			strcmp(path, http->handlers[i].pattern) == 0 ||
			(strncmp(path, http->handlers[i].pattern, strlen(http->handlers[i].pattern)) == 0 && path[strlen(http->handlers[i].pattern)] == '/')))
		{
			*out_callback = http->handlers[i].callback;
			*out_user_data = http->handlers[i].user_data;
			return true;
		}
	}
	return false;
}

static void _http_on_write(uv_write_t* write, int status)
{
	tf_free(write);
}

static void _http_connection_on_close(uv_handle_t* handle)
{
	tf_http_connection_t* connection = handle->data;
	handle->data = NULL;
	_http_connection_destroy(connection, "handle closed");
}

static void _http_connection_destroy(tf_http_connection_t* connection, const char* reason)
{
	if (connection->tcp.data)
	{
		tf_printf("CLOSE %p: %s\n", connection, reason);
		uv_close((uv_handle_t*)&connection->tcp, _http_connection_on_close);
	}
	else if (connection->ref_count == 0)
	{
		tf_http_t* http = connection->http;
		for (int i = 0; i < http->connections_count; i++)
		{
			if (http->connections[i] == connection)
			{
				http->connections[i] = http->connections[--http->connections_count];
			}
		}
		if (connection->body)
		{
			tf_free(connection->body);
			connection->body = NULL;
		}
		tf_free(connection);
	}
}

static void _http_builtin_404_handler(tf_http_request_t* request)
{
	const char* k_payload = _http_status_text(404);
	tf_http_respond(request, 404, NULL, 0, k_payload, strlen(k_payload));
}

static void _http_reset_connection(tf_http_connection_t* connection)
{
	connection->headers_done = false;
	connection->headers_buffer_length = 0;
	connection->body_length = 0;
	connection->content_length = 0;
	connection->parsed_length = 0;
}

static void _http_add_body_bytes(tf_http_connection_t* connection, const void* data, size_t size)
{
	if (connection->is_websocket)
	{
		connection->body = tf_realloc(connection->body, connection->body_length + size);
		memcpy((char*)connection->body + connection->body_length, data, size);
		connection->body_length += size;

		uint8_t* p = connection->body;
		while (connection->body_length >= 2)
		{
			uint8_t bits0 = p[0];
			uint8_t bits1 = p[1];
			if ((bits1 & (1 << 7)) == 0)
			{
				/* Unmasked message. */
				_http_connection_destroy(connection, "websocket server received unmasked bytes");
				return;
			}
			uint8_t opcode = bits0 & 0xf;
			bool fin = (bits0 & (1 << 7)) != 0;
			size_t length = bits1 & 0x7f;
			int mask_start = 2;
			if (length == 126)
			{
				length = 0;
				for (int i = 0; i < 2; i++)
				{
					length <<= 8;
					length |= p[2 + i];
				}
				mask_start = 4;
			}
			else if (length == 127)
			{
				length = 0;
				for (int i = 0; i < 8; i++)
				{
					length <<= 8;
					length |= p[2 + i];
				}
				mask_start = 10;
			}
			if (connection->body_length >= mask_start + length + 4)
			{
				uint32_t mask =
					p[mask_start + 0] |
					p[mask_start + 1] << 8 |
					p[mask_start + 2] << 16 |
					p[mask_start + 3] << 24;
				size_t i = 0;
				size_t start = mask_start + 4;
				size_t end = mask_start + 4 + length;
				for (i = start; i < end; i += 4)
				{
					*(uint32_t*)(p + i) ^= mask;
				}
				for (; i < end; i++)
				{
					p[i] ^= ((mask >> (8 * (i % 4))) & 0xff);
				}
				if (fin)
				{
					if (connection->request->on_message)
					{
						connection->request->on_message(connection->request, p + mask_start + 4, length);
					}
					else
					{
						tf_printf("MESSAGE %d [%.*s]\n", opcode, (int)length, p + mask_start + 4);
					}
				}
				size_t total_length = end;
				memmove(connection->body, (char*)connection->body + total_length, connection->body_length - total_length);
				connection->body_length -= total_length;
			}
		}
	}
	else
	{
		size_t fit = tf_min(connection->content_length - connection->body_length, size);
		if (fit > 0)
		{
			memcpy((char*)connection->body + connection->body_length, data, fit);
			connection->body_length += fit;
		}

		if (connection->body_length == connection->content_length)
		{
			if (connection->flags & k_tf_http_handler_flag_websocket)
			{
				connection->is_websocket = true;
			}
			tf_http_request_t* request = tf_malloc(sizeof(tf_http_request_t));
			*request = (tf_http_request_t)
			{
				.connection = connection,
				.phase = k_http_callback_phase_headers_received,
				.method = connection->method,
				.path = connection->path,
				.flags = connection->flags,
				.query = connection->query,
				.body = connection->body,
				.content_length = connection->content_length,
				.headers = connection->headers,
				.headers_count = connection->headers_length,
				.user_data = connection->user_data,
			};
			connection->request = request;

			tf_http_request_ref(request);
			connection->callback(request);
			tf_http_request_release(request);
			if (!connection->is_websocket)
			{
				_http_reset_connection(connection);
			}
		}
	}
}

static void _http_on_read(uv_stream_t* stream, ssize_t read_size, const uv_buf_t* buffer)
{
	tf_http_connection_t* connection = stream->data;
	if (read_size > 0)
	{
		if (!connection->headers_done)
		{
			connection->headers_buffer_length += read_size;

			const char* method = NULL;
			size_t method_length = 0;
			const char* path = NULL;
			size_t path_length = 0;
			size_t header_count = sizeof(connection->headers) / sizeof(*connection->headers);

			int parse_result = phr_parse_request(connection->headers_buffer, connection->headers_buffer_length, &method, &method_length, &path, &path_length, &connection->minor_version, connection->headers, &header_count, connection->parsed_length);
			connection->parsed_length = connection->headers_buffer_length;
			if (parse_result > 0)
			{
				connection->headers_done = true;
				connection->headers_length = header_count;
				connection->method = method;
				((char*)connection->method)[method_length] = '\0';
				connection->path = path;
				((char*)connection->path)[path_length] = '\0';
				char* q = strchr(connection->path, '?');
				if (q)
				{
					*q = '\0';
					connection->query = q + 1;
				}

				connection->connection_close = connection->minor_version == 0;

				for (int i = 0; i < (int)header_count; i++)
				{
					for (size_t j = 0; j < connection->headers[i].name_len; j++)
					{
						if (connection->headers[i].name[j] >= 'A' &&
							connection->headers[i].name[j] <= 'Z')
						{
							((char*)connection->headers[i].name)[j] += 'a' - 'A';
						}
					}
					((char*)connection->headers[i].name)[connection->headers[i].name_len] = '\0';
					((char*)connection->headers[i].value)[connection->headers[i].value_len] = '\0';
					if (strcasecmp(connection->headers[i].name, "content-length") == 0)
					{
						connection->content_length = strtoull(connection->headers[i].value, NULL, 10);
					}
					else if (strcasecmp(connection->headers[i].name, "connection") == 0)
					{
						if (strcasecmp(connection->headers[i].value, "close") == 0)
						{
							connection->connection_close = true;
						}
					}
				}

				if (connection->content_length)
				{
					connection->body = tf_malloc(connection->content_length);
				}

				int flags = _http_connection_get_header(connection, "upgrade") ? k_tf_http_handler_flag_websocket : 0;
				connection->flags = flags;
				if (!_http_find_handler(connection->http, connection->path, flags, &connection->callback, &connection->user_data) || !connection->callback)
				{
					connection->callback = _http_builtin_404_handler;
				}

				_http_add_body_bytes(connection, connection->headers_buffer + parse_result, connection->headers_buffer_length - parse_result);
			}
			else if (parse_result == -2)
			{
				/* Incomplete.  Will try again next time. */
			}
			else
			{
				tf_printf("phr_parse_request: %d\n", parse_result);
				_http_connection_destroy(connection, "failed to parse request headers");
			}
		}
		else
		{
			connection->buffer_length += read_size;
			_http_add_body_bytes(connection, connection->buffer, connection->buffer_length);
			connection->buffer_length = 0;
		}
	}
	else
	{
		_http_connection_destroy(connection, uv_strerror(read_size));
	}
}

static void _http_on_connection(uv_stream_t* stream, int status)
{
	tf_http_t* http = stream->data;
	tf_http_connection_t* connection = tf_malloc(sizeof(tf_http_connection_t));
	*connection = (tf_http_connection_t) { .http = http, .tcp = { .data = connection } };
	int r = uv_tcp_init(connection->http->loop, &connection->tcp);
	if (r)
	{
		tf_printf("uv_tcp_init: %s\n", uv_strerror(r));
		tf_free(connection);
		return;
	}

	r = uv_accept(stream, (uv_stream_t*)&connection->tcp);
	if (r)
	{
		tf_printf("uv_accept: %s\n", uv_strerror(r));
		uv_close((uv_handle_t*)&connection->tcp, NULL);
		tf_free(connection);
		return;
	}

	r = uv_read_start((uv_stream_t*)&connection->tcp, _http_allocate_buffer, _http_on_read);
	if (r)
	{
		tf_printf("uv_read-start: %s\n", uv_strerror(r));
		uv_close((uv_handle_t*)&connection->tcp, NULL);
		tf_free(connection);
		return;
	}

	http->connections = tf_realloc(http->connections, sizeof(tf_http_connection_t*) * (http->connections_count + 1));
	http->connections[http->connections_count++] = connection;
}

void tf_http_listen(tf_http_t* http, int port)
{
	uv_tcp_t* tcp = tf_malloc(sizeof(uv_tcp_t));
	*tcp = (uv_tcp_t) { .data = http };
	int r = uv_tcp_init(http->loop, tcp);
	if (r)
	{
		tf_printf("uv_tcp_init: %s\n", uv_strerror(r));
	}

	if (r == 0)
	{
		struct sockaddr_in addr =
		{
			.sin_family = AF_INET,
			.sin_addr = { .s_addr = INADDR_ANY },
			.sin_port = ntohs(port),
		};
		r = uv_tcp_bind(tcp, (struct sockaddr*)&addr, 0);
		if (r)
		{
			tf_printf("uv_tcp_bind: %s\n", uv_strerror(r));
		}
	}

	if (r == 0)
	{
		r = uv_listen((uv_stream_t*)tcp, 16, _http_on_connection);
		if (r)
		{
			tf_printf("uv_listen: %s\n", uv_strerror(r));
		}
	}

	if (r == 0)
	{
		http->listeners = tf_realloc(http->listeners, sizeof(uv_tcp_t*) * (http->listeners_count + 1));
		http->listeners[http->listeners_count++] = tcp;
	}
}

void tf_http_add_handler(tf_http_t* http, const char* pattern, int flags, tf_http_callback_t* callback, void* user_data)
{
	http->handlers = tf_realloc(http->handlers, sizeof(tf_http_handler_t) * (http->handlers_count + 1));
	http->handlers[http->handlers_count++] = (tf_http_handler_t)
	{
		.pattern = tf_strdup(pattern),
		.flags = flags,
		.callback = callback,
		.user_data = user_data,
	};
}

static void _http_free_on_close(uv_handle_t* handle)
{
	handle->data = NULL;
	tf_free(handle);
}

void tf_http_destroy(tf_http_t* http)
{
	for (int i = 0; i < http->listeners_count; i++)
	{
		uv_close((uv_handle_t*)http->listeners[i], _http_free_on_close);
	}
	tf_free(http->listeners);
	http->listeners = NULL;
	http->listeners_count = 0;

	for (int i = 0; i < http->handlers_count; i++)
	{
		if (http->handlers[i].pattern)
		{
			tf_free((void*)http->handlers[i].pattern);
			http->handlers[i].pattern = NULL;
		}
	}
	tf_free(http->handlers);
	http->handlers_count = 0;

	tf_free(http->connections);
	http->connections_count = 0;

	tf_free(http);
}

static const char* _http_status_text(int status)
{
	switch (status)
	{
	case 101: return "Switching Protocols";
	case 200: return "OK";
	case 303: return "See other";
	case 304: return "Not Modified";
	case 400: return "Bad Request";
	case 401: return "Unauthorized";
	case 403: return "Forbidden";
	case 404: return "File not found";
	case 500: return "Internal server error";
	default: return "Unknown";
	}
}

static void _http_on_shutdown(uv_shutdown_t* request, int status)
{
	request->data = NULL;
}

static void _http_write(tf_http_connection_t* connection, const void* data, size_t size)
{
	uv_write_t* write = tf_malloc(sizeof(uv_write_t) + size);
	*write = (uv_write_t) { .data = connection };
	memcpy(write + 1, data, size);
	int r = uv_write(write, (uv_stream_t*)&connection->tcp, &(uv_buf_t) { .base = (void*)(write + 1), .len = size }, 1, _http_on_write);
	if (r)
	{
		tf_printf("uv_write: %s\n", uv_strerror(r));
	}
}

void tf_http_request_send(tf_http_request_t* request, const void* data, size_t size)
{
	_http_write(request->connection, data, size);
}

void tf_http_respond(tf_http_request_t* request, int status, const char** headers, int headers_count, const void* body, size_t content_length)
{
	const char* status_text = _http_status_text(status);
	/* HTTP/1.x 200 OK\r\n */
	bool sent_content_length = false;
	int headers_length = 8 + 1 + 3 + 1 + strlen(status_text) + 2;
	if (headers)
	{
		for (int i = 0; i < headers_count * 2; i += 2)
		{
			/* Key: Value\r\n */
			headers_length += strlen(headers[i]) + 2 + strlen(headers[i + 1]) + 2;
			if (strcasecmp(headers[i], "content-length") == 0)
			{
				sent_content_length = true;
			}
		}
	}
	/* \r\n */
	headers_length += 2;

	char content_length_buffer[32] = { 0 };
	int content_length_buffer_length = 0;
	if (!sent_content_length && status != 101)
	{
		content_length_buffer_length = snprintf(content_length_buffer, sizeof(content_length_buffer), "Content-Length: %zd\r\n", content_length);
		headers_length += content_length_buffer_length;
	}

	uv_write_t* write = tf_malloc(sizeof(uv_write_t) + headers_length + content_length + 1);
	*write = (uv_write_t) { .data = request->connection };
	char* buffer = (char*)(write + 1);
	int offset = snprintf(buffer, headers_length + 1, "HTTP/1.%d %03d %s\r\n", request->connection->minor_version, status, status_text);
	if (headers)
	{
		for (int i = 0; i < headers_count * 2; i += 2)
		{
			offset += snprintf(buffer + offset, headers_length + 1 - offset, "%s: %s\r\n", headers[i], headers[i + 1]);
		}
	}
	if (!sent_content_length)
	{
		memcpy(buffer + offset, content_length_buffer, content_length_buffer_length);
		offset += content_length_buffer_length;
	}
	offset += snprintf(buffer + offset, headers_length + 1 - offset, "\r\n");
	assert(offset == headers_length);
	if (content_length)
	{
		memcpy(buffer + offset, body, content_length);
	}
	int r = uv_write(write, (uv_stream_t*)&request->connection->tcp, &(uv_buf_t) { .base = buffer, .len = headers_length + content_length }, 1, _http_on_write);
	if (r)
	{
		tf_printf("uv_write: %s\n", uv_strerror(r));
	}

	if (request->connection->connection_close &&
		!request->connection->shutdown.data)
	{
		request->connection->shutdown.data = request->connection;
		uv_shutdown(&request->connection->shutdown, (uv_stream_t*)&request->connection->tcp, _http_on_shutdown);
	}
}

size_t tf_http_get_body(const tf_http_request_t* request, const void** out_data)
{
	*out_data = request->connection->body;
	return request->connection->content_length;
}

void tf_http_request_ref(tf_http_request_t* request)
{
	request->ref_count++;
	request->connection->ref_count++;
}

void tf_http_request_release(tf_http_request_t* request)
{
	if (--request->connection->ref_count == 0)
	{
		/* Reset the connection? */
	}
	if (--request->ref_count == 0)
	{
		tf_free(request);
	}
}

static const char* _http_connection_get_header(const tf_http_connection_t* connection, const char* name)
{
	for (int i = 0; i < connection->headers_length; i++)
	{
		if (strcasecmp(connection->headers[i].name, name) == 0)
		{
			return connection->headers[i].value;
		}
	}
	return NULL;
}

const char* tf_http_request_get_header(tf_http_request_t* request, const char* name)
{
	return _http_connection_get_header(request->connection, name);
}