#include <assert.h>
#include <errno.h>
#include <stddef.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <unistd.h>

#include <aml.h>
#include <json_object.h>

#include "conn.h"
#include "vali.h"
#include "util.h"

#if HAVE_AML_V1
#define AML_HANDLER_ARG struct aml_handler *
#else
#define AML_HANDLER_ARG void *
#endif

struct vali_service {
	struct aml *event_loop;
	struct vali_service_call_handler call_handler;
	struct array listen_handlers; // struct aml_handler *
	void *user_data;
};

struct vali_service_conn {
	struct vali_conn base;
	struct aml_handler *handler;
	struct vali_service *service;
	struct vali_service_call *pending_call;
};

struct vali_service_call {
	struct vali_service_conn *conn; // NULL if disconnected
	char *method;
	struct json_object *params;
	bool oneway, more;
	void *user_data;
};

struct vali_service *vali_service_create(void) {
	struct vali_service *service = calloc(1, sizeof(*service));
	if (service == NULL) {
		return NULL;
	}

	service->event_loop = aml_new();
	if (service->event_loop == NULL) {
		free(service);
		return NULL;
	}

	return service;
}

void vali_service_destroy(struct vali_service *service) {
	// TODO: close client connections
	struct aml_handler **listen_handlers = service->listen_handlers.data;
	size_t listen_handlers_len = service->listen_handlers.size / sizeof(listen_handlers[0]);
	for (size_t i = 0; i < listen_handlers_len; i++) {
		struct aml_handler *handler = listen_handlers[i];
		aml_stop(service->event_loop, handler);
		close(aml_get_fd(handler));
		aml_unref(handler);
	}
	array_finish(&service->listen_handlers);
	aml_unref(service->event_loop);
	free(service);
}

void *vali_service_get_user_data(struct vali_service *service) {
	return service->user_data;
}

void vali_service_set_user_data(struct vali_service *service, void *user_data) {
	service->user_data = user_data;
}

int vali_service_get_fd(struct vali_service *service) {
	return aml_get_fd(service->event_loop);
}

void vali_service_set_call_handler(struct vali_service *service, struct vali_service_call_handler handler) {
	service->call_handler = handler;
}

static bool conn_dequeue_request(struct vali_service_conn *conn, struct vali_service_call **out) {
	*out = NULL;

	struct json_object *req;
	if (!conn_dequeue(&conn->base, &req)) {
		return false;
	} else if (req == NULL) {
		return true;
	}

	const char *method = json_object_get_string(json_object_object_get(req, "method"));
	struct json_object *params = json_object_object_get(req, "parameters");
	bool oneway = json_object_get_boolean(json_object_object_get(req, "oneway"));
	bool more = json_object_get_boolean(json_object_object_get(req, "more"));
	if (method == NULL) {
		goto err_req;
	}

	struct vali_service_call *call = calloc(1, sizeof(*call));
	if (call == NULL) {
		goto err_req;
	}

	*call = (struct vali_service_call){
		.conn = conn,
		.method = strdup(method),
		.params = json_object_get(params),
		.oneway = oneway,
		.more = more,
	};
	if (call->method == NULL) {
		free(call);
		goto err_req;
	}
	json_object_put(req);
	*out = call;
	return true;

err_req:
	json_object_put(req);
	return false;
}

static void conn_destroy(struct vali_service_conn *conn) {
	if (conn->pending_call != NULL) {
		conn->pending_call->conn = NULL;
	}
	aml_stop(conn->service->event_loop, conn->handler);
	aml_unref(conn->handler);
	conn_finish(&conn->base);
	free(conn);
}

static void conn_update_event_mask(struct vali_service_conn *conn) {
	uint32_t mask = 0;
	if (conn->pending_call == NULL) {
		mask |= AML_EVENT_READ;
	}
	if (conn->base.out.size > 0) {
		mask |= AML_EVENT_WRITE;
	}
	aml_set_event_mask(conn->handler, mask);
}

static void conn_handle_event(AML_HANDLER_ARG data) {
	struct vali_service_conn *conn = aml_get_userdata(data);
	uint32_t revents = aml_get_revents(conn->handler);

	bool eof = false;
	if (revents & AML_EVENT_READ) {
		if (!conn_receive(&conn->base, &eof)) {
			goto error;
		}

		while (conn->pending_call == NULL) {
			struct vali_service_call *call = NULL;
			if (!conn_dequeue_request(conn, &call)) {
				goto error;
			} else if (call == NULL) {
				break;
			}

			conn->pending_call = call;
			conn->service->call_handler.func(call, conn->service->call_handler.user_data);
		}
	}

	if (revents & AML_EVENT_WRITE) {
		if (!conn_flush(&conn->base)) {
			goto error;
		}
	}

	conn_update_event_mask(conn);

	if (eof && conn->base.in.size == 0 && conn->base.out.size == 0 && conn->pending_call == NULL) {
		conn_destroy(conn);
	}

	return;

error:
	conn_destroy(conn);
}

struct vali_service_conn *vali_service_create_conn(struct vali_service *service, int fd) {
	if (!set_cloexec(fd) || !set_nonblock(fd)) {
		return NULL;
	}

	struct vali_service_conn *conn = calloc(1, sizeof(*conn));
	if (conn == NULL) {
		return NULL;
	}

	// TODO: limit conn buffer size
	conn_init(&conn->base, fd);
	conn->service = service;

	conn->handler = aml_handler_new(fd, conn_handle_event, conn, NULL);
	if (conn->handler == NULL) {
		goto err_conn;
	}

	if (aml_start(service->event_loop, conn->handler) < 0) {
		goto err_handler;
	}

	aml_set_event_mask(conn->handler, AML_EVENT_READ);

	return conn;

err_handler:
	aml_unref(conn->handler);
err_conn:
	free(conn);
	return NULL;
}

void vali_service_conn_disconnect(struct vali_service_conn *conn) {
	// TODO: handle errors somehow?
	shutdown(conn->base.fd, SHUT_RDWR);
}

int vali_service_conn_get_fd(struct vali_service_conn *conn) {
	return conn->base.fd;
}

static void listener_handle_readable(AML_HANDLER_ARG data) {
	struct aml_handler *listen_handler = data;
	struct vali_service *service = aml_get_userdata(listen_handler);

	int client_fd = accept(aml_get_fd(listen_handler), NULL, NULL);
	if (client_fd < 0) {
		return;
	}

	struct vali_service_conn *conn = vali_service_create_conn(service, client_fd);
	if (conn == NULL) {
		close(client_fd);
		return;
	}
}

bool vali_service_listen_unix(struct vali_service *service, const char *path) {
	if (service->call_handler.func == NULL) {
		abort();
	}

	struct sockaddr_un addr;
	if (!set_sockaddr_un(&addr, path)) {
		errno = -EINVAL;
		return false;
	}

	int fd = socket(AF_UNIX, SOCK_STREAM, 0);
	if (fd < 0) {
		return false;
	}

	if (!set_cloexec(fd) || !set_nonblock(fd)) {
		goto err_fd;
	}

	if (bind(fd, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
		goto err_fd;
	}

	if (listen(fd, 128) < 0) {
		goto err_fd;
	}

	if (!vali_service_listen_fd(service, fd)) {
		goto err_fd;
	}

	return true;

err_fd:
	close(fd);
	return false;
}

bool vali_service_listen_fd(struct vali_service *service, int fd) {
	struct aml_handler *handler = aml_handler_new(fd, listener_handle_readable, service, NULL);
	if (handler == NULL) {
		return false;
	}

	if (aml_start(service->event_loop, handler) < 0) {
		goto err_handler;
	}

	struct aml_handler **handler_ptr = array_add(&service->listen_handlers, sizeof(*handler_ptr));
	if (handler_ptr == NULL) {
		goto err_stop;
	}
	*handler_ptr = handler;

	aml_set_event_mask(handler, AML_EVENT_READ);

	return true;

err_stop:
	aml_stop(service->event_loop, handler);
err_handler:
	aml_unref(handler);
	return false;
}

bool vali_service_dispatch(struct vali_service *service) {
	if (aml_poll(service->event_loop, -1) < 0) {
		return false;
	}
	aml_dispatch(service->event_loop);
	return true;
}

void *vali_service_call_get_user_data(const struct vali_service_call *call) {
	return call->user_data;
}

void vali_service_call_set_user_data(struct vali_service_call *call, void *user_data) {
	call->user_data = user_data;
}

struct vali_service *vali_service_call_get_service(struct vali_service_call *call) {
	return call->conn->service;
}

struct vali_service_conn *vali_service_call_get_conn(struct vali_service_call *call) {
	return call->conn;
}

const char *vali_service_call_get_method(const struct vali_service_call *call) {
	return call->method;
}

struct json_object *vali_service_call_get_parameters(const struct vali_service_call *call) {
	return call->params;
}

bool vali_service_call_is_oneway(const struct vali_service_call *call) {
	return call->oneway;
}

bool vali_service_call_is_more(const struct vali_service_call *call) {
	return call->more;
}

static void service_call_destroy(struct vali_service_call *call) {
	if (call->conn != NULL) {
		assert(call->conn->pending_call == call);
		call->conn->pending_call = NULL;
	}
	json_object_put(call->params);
	free(call->method);
	free(call);
}

static void service_call_disconnect(struct vali_service_call *call) {
	struct vali_service_conn *conn = call->conn;
	service_call_destroy(call);
	if (conn != NULL) {
		vali_service_conn_disconnect(conn);
	}
}

static void service_call_close_with_response(struct vali_service_call *call, struct json_object *resp) {
	struct vali_service_conn *conn = call->conn;
	if (conn == NULL) {
		json_object_put(resp);
		return;
	}

	if (call->oneway) {
		json_object_put(resp);
	} else {
		if (!conn_enqueue(&conn->base, resp)) {
			vali_service_conn_disconnect(conn);
		}
	}

	service_call_destroy(call);
	conn_update_event_mask(conn);
}

void vali_service_call_close_with_reply(struct vali_service_call *call, struct json_object *params) {
	assert(params == NULL || json_object_get_type(params) == json_type_object);

	struct json_object *resp = NULL;
	if (call->oneway) {
		json_object_put(params);
	} else {
		resp = json_object_new_object();
		if (resp == NULL) {
			service_call_disconnect(call);
			return;
		}
		json_object_object_add(resp, "parameters", params);
	}

	service_call_close_with_response(call, resp);
}

void vali_service_call_close_with_error(struct vali_service_call *call, const char *error, struct json_object *params) {
	struct json_object *resp = NULL;
	if (call->oneway) {
		json_object_put(params);
	} else {
		resp = json_object_new_object();
		struct json_object *error_obj = json_object_new_string(error);
		if (resp == NULL || error_obj == NULL) {
			service_call_disconnect(call);
			return;
		}
		json_object_object_add(resp, "error", error_obj);
		json_object_object_add(resp, "parameters", params);
	}

	service_call_close_with_response(call, resp);
}

void vali_service_call_reply(struct vali_service_call *call, struct json_object *params) {
	assert(call->more);
	assert(params == NULL || json_object_get_type(params) == json_type_object);

	if (call->conn == NULL) {
		return;
	}

	struct json_object *resp = NULL;
	if (call->oneway) {
		json_object_put(params);
	} else {
		resp = json_object_new_object();
		struct json_object *continues_obj = json_object_new_boolean(true);
		if (resp == NULL || continues_obj == NULL) {
			vali_service_conn_disconnect(call->conn);
			return;
		}
		json_object_object_add(resp, "parameters", params);
		json_object_object_add(resp, "continues", continues_obj);
	}

	if (!conn_enqueue(&call->conn->base, resp)) {
		vali_service_conn_disconnect(call->conn);
	}
	conn_update_event_mask(call->conn);
}
