#include "taskstub.js.h"

#include "log.h"
#include "mem.h"
#include "packetstream.h"
#include "serialize.h"
#include "task.h"
#include "util.js.h"

#include <string.h>
#include <stdio.h>

#ifdef _WIN32
#include <io.h>
#include <windows.h>
#include <ws2tcpip.h>
#else
#include <unistd.h>
#endif

static JSClassID _classId;
static char _executable[1024];

typedef struct _tf_taskstub_t
{
	taskid_t _id;
	JSValue _object;

	JSValue _on_exit;
	JSValue _on_error;
	JSValue _on_print;

	tf_task_t* _owner;
	tf_packetstream_t* _stream;
	uv_process_t _process;
	bool _finalized;
} tf_taskstub_t;

void tf_taskstub_startup()
{
	static bool initialized;
	if (!initialized)
	{
		JS_NewClassID(&_classId);
		size_t size = sizeof(_executable);
		uv_exepath(_executable, &size);
		tf_printf("exepath is %s\n", _executable);
		initialized = true;
	}
}

static JSValue _taskstub_activate(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv);
static JSValue _taskstub_execute(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv);
static JSValue _taskstub_setImports(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv);
static JSValue _taskstub_getExports(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv);
static JSValue _taskstub_kill(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv);
static JSValue _taskstub_get_on_exit(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv);
static JSValue _taskstub_set_on_exit(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv);
static JSValue _taskstub_get_on_error(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv);
static JSValue _taskstub_set_on_error(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv);
static JSValue _taskstub_get_on_print(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv);
static JSValue _taskstub_set_on_print(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv);
static JSValue _taskstub_loadFile(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv);
static void _taskstub_on_process_exit(uv_process_t* process, int64_t status, int terminationSignal);
static void _taskstub_finalizer(JSRuntime* runtime, JSValue value);

static void _tf_taskstub_run_sandbox_thread(void* data)
{
	uv_file fd = (uv_file)(intptr_t)data;
	tf_task_t* task = tf_task_create();
	tf_task_set_one_proc(task, true);
	tf_task_configure_from_fd(task, fd);
	/* The caller will trigger tf_task_activate with a message. */
	tf_task_run(task);
	tf_task_destroy(task);
}

static JSValue _taskstub_create(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv)
{
	tf_task_t* parent = tf_task_get(context);
	tf_taskstub_t* stub = tf_malloc(sizeof(tf_taskstub_t));
	memset(stub, 0, sizeof(*stub));
	stub->_stream = tf_packetstream_create();

	JSValue taskObject = JS_NewObjectClass(context, _classId);
	JS_SetOpaque(taskObject, stub);
	stub->_owner = parent;
	stub->_on_exit = JS_UNDEFINED;
	stub->_on_error = JS_UNDEFINED;
	stub->_on_print = JS_UNDEFINED;
	stub->_object = taskObject;

	JSAtom atom = JS_NewAtom(context, "onExit");
	JS_DefinePropertyGetSet(
		context, taskObject, atom, JS_NewCFunction(context, _taskstub_get_on_exit, "getOnExit", 0), JS_NewCFunction(context, _taskstub_set_on_exit, "setOnExit", 0), 0);
	JS_FreeAtom(context, atom);

	atom = JS_NewAtom(context, "onError");
	JS_DefinePropertyGetSet(
		context, taskObject, atom, JS_NewCFunction(context, _taskstub_get_on_error, "getOnError", 0), JS_NewCFunction(context, _taskstub_set_on_error, "setOnError", 0), 0);
	JS_FreeAtom(context, atom);

	atom = JS_NewAtom(context, "onPrint");
	JS_DefinePropertyGetSet(
		context, taskObject, atom, JS_NewCFunction(context, _taskstub_get_on_print, "getOnPrint", 0), JS_NewCFunction(context, _taskstub_set_on_print, "setOnPrint", 0), 0);
	JS_FreeAtom(context, atom);

	JS_SetPropertyStr(context, taskObject, "activate", JS_NewCFunction(context, _taskstub_activate, "activate", 0));
	JS_SetPropertyStr(context, taskObject, "execute", JS_NewCFunction(context, _taskstub_execute, "execute", 1));
	JS_SetPropertyStr(context, taskObject, "setImports", JS_NewCFunction(context, _taskstub_setImports, "setImports", 1));
	JS_SetPropertyStr(context, taskObject, "getExports", JS_NewCFunction(context, _taskstub_getExports, "getExports", 0));
	JS_SetPropertyStr(context, taskObject, "kill", JS_NewCFunction(context, _taskstub_kill, "kill", 0));
	JS_SetPropertyStr(context, taskObject, "loadFile", JS_NewCFunction(context, _taskstub_loadFile, "loadFile", 1));

	taskid_t id = k_task_parent_id;
	if (parent)
	{
		id = tf_task_allocate_task_id(parent, (tf_taskstub_t*)stub);
	}
	stub->_id = id;

	char arg1[] = "sandbox";
	char* command_argv[] = { _executable, arg1, 0 };

	JSValue result = JS_NULL;
	if (tf_task_get_one_proc(parent))
	{
		uv_os_sock_t fds[2];
		int pipe_result = uv_socketpair(SOCK_STREAM, 0, fds, 0, 0);
		if (pipe_result)
		{
			tf_printf("uv_socketpair failed: %s\n", uv_strerror(pipe_result));
		}

		uv_pipe_t* pipe = tf_packetstream_get_pipe(stub->_stream);
		memset(pipe, 0, sizeof(*pipe));
		pipe_result = uv_pipe_init(tf_task_get_loop(parent), pipe, 1);
		if (pipe_result != 0)
		{
			tf_printf("uv_pipe_init failed: %s\n", uv_strerror(pipe_result));
		}
		pipe_result = uv_pipe_open(pipe, fds[0]);
		if (pipe_result != 0)
		{
			tf_printf("uv_pipe_open failed: %s\n", uv_strerror(pipe_result));
		}

		uv_thread_t* thread = tf_malloc(sizeof(uv_thread_t));
		uv_thread_create(thread, _tf_taskstub_run_sandbox_thread, (void*)(intptr_t)fds[1]);

		tf_packetstream_set_on_receive(stub->_stream, tf_task_on_receive_packet, stub);
		tf_packetstream_start(stub->_stream);
		result = taskObject;
	}
	else
	{
		uv_pipe_t* pipe = tf_packetstream_get_pipe(stub->_stream);
		memset(pipe, 0, sizeof(*pipe));
		if (uv_pipe_init(tf_task_get_loop(parent), pipe, 1) != 0)
		{
			tf_printf("uv_pipe_init failed\n");
		}

		uv_stdio_container_t io[3];
		io[0].flags = UV_CREATE_PIPE | UV_READABLE_PIPE | UV_WRITABLE_PIPE;
		io[0].data.stream = (uv_stream_t*)pipe;
		io[1].flags = UV_INHERIT_FD;
		io[1].data.fd = STDOUT_FILENO;
		io[2].flags = UV_INHERIT_FD;
		io[2].data.fd = STDERR_FILENO;

		uv_process_options_t options = { 0 };
		options.args = command_argv;
		options.exit_cb = _taskstub_on_process_exit;
		options.stdio = io;
		options.stdio_count = sizeof(io) / sizeof(*io);
		options.file = command_argv[0];
		options.env = (char*[]) { "ASAN_OPTIONS=detect_leaks=0", NULL };

		stub->_process.data = stub;
		int spawn_result = uv_spawn(tf_task_get_loop(parent), &stub->_process, &options);
		if (spawn_result == 0)
		{
			tf_packetstream_set_on_receive(stub->_stream, tf_task_on_receive_packet, stub);
			tf_packetstream_start(stub->_stream);
			result = taskObject;
		}
		else
		{
			tf_printf("uv_spawn failed: %s\n", uv_strerror(spawn_result));
			JS_FreeValue(context, taskObject);
		}
	}
	return JS_DupValue(context, result);
}

static void _taskstub_gc_mark(JSRuntime* rt, JSValueConst value, JS_MarkFunc mark_func)
{
	tf_taskstub_t* stub = JS_GetOpaque(value, _classId);
	if (stub)
	{
		JS_MarkValue(rt, stub->_on_exit, mark_func);
		JS_MarkValue(rt, stub->_on_error, mark_func);
		JS_MarkValue(rt, stub->_on_print, mark_func);
	}
}

JSValue tf_taskstub_register(JSContext* context)
{
	JSClassDef def = {
		.class_name = "TaskStub",
		.finalizer = &_taskstub_finalizer,
		.gc_mark = _taskstub_gc_mark,
	};
	if (JS_NewClass(JS_GetRuntime(context), _classId, &def) != 0)
	{
		tf_printf("Failed to register TaskStub class.\n");
	}
	return JS_NewCFunction2(context, _taskstub_create, "TaskStub", 0, JS_CFUNC_constructor, 0);
}

taskid_t tf_taskstub_get_id(const tf_taskstub_t* stub)
{
	return stub->_id;
}

JSValue tf_taskstub_get_task_object(const tf_taskstub_t* stub)
{
	return stub->_object;
}

tf_packetstream_t* tf_taskstub_get_stream(const tf_taskstub_t* stub)
{
	return stub->_stream;
}

tf_task_t* tf_taskstub_get_owner(const tf_taskstub_t* stub)
{
	return stub->_owner;
}

tf_taskstub_t* tf_taskstub_create_parent(tf_task_t* task, uv_file file)
{
	JSValue parentObject = JS_NewObjectClass(tf_task_get_context(task), _classId);
	tf_taskstub_t* parentStub = tf_malloc(sizeof(tf_taskstub_t));
	memset(parentStub, 0, sizeof(tf_taskstub_t));
	parentStub->_stream = tf_packetstream_create();
	parentStub->_on_exit = JS_UNDEFINED;
	parentStub->_on_error = JS_UNDEFINED;
	parentStub->_on_print = JS_UNDEFINED;

	JS_SetOpaque(parentObject, parentStub);
	parentStub->_owner = task;
	parentStub->_id = k_task_parent_id;
	parentStub->_object = parentObject;

	if (uv_pipe_init(tf_task_get_loop(task), tf_packetstream_get_pipe(parentStub->_stream), 1) != 0)
	{
		tf_printf("uv_pipe_init failed\n");
	}
	tf_packetstream_set_on_receive(parentStub->_stream, tf_task_on_receive_packet, parentStub);
	int result = uv_pipe_open(tf_packetstream_get_pipe(parentStub->_stream), file);
	if (result != 0)
	{
		tf_printf("uv_pipe_open failed: %s\n", uv_strerror(result));
	}
	tf_packetstream_start(parentStub->_stream);
	return parentStub;
}

static void _taskstub_cleanup(tf_taskstub_t* stub)
{
	if (!stub->_process.data && JS_IsUndefined(stub->_object) && stub->_finalized)
	{
		tf_free(stub);
	}
}

static void _taskstub_finalizer(JSRuntime* runtime, JSValue value)
{
	tf_taskstub_t* stub = JS_GetOpaque(value, _classId);
	stub->_object = JS_UNDEFINED;
	JSContext* context = tf_task_get_context(stub->_owner);
	if (!JS_IsUndefined(stub->_on_exit))
	{
		JS_FreeValue(context, stub->_on_exit);
		stub->_on_exit = JS_UNDEFINED;
	}
	if (!JS_IsUndefined(stub->_on_error))
	{
		JS_FreeValue(context, stub->_on_error);
		stub->_on_error = JS_UNDEFINED;
	}
	if (!JS_IsUndefined(stub->_on_print))
	{
		JS_FreeValue(context, stub->_on_print);
		stub->_on_print = JS_UNDEFINED;
	}
	if (stub->_stream)
	{
		tf_packetstream_destroy(stub->_stream);
		stub->_stream = NULL;
	}
	stub->_finalized = true;
	tf_task_remove_child(stub->_owner, stub);
	_taskstub_cleanup(stub);
}

static void _taskstub_on_handle_close(uv_handle_t* handle)
{
	tf_taskstub_t* stub = handle->data;
	tf_task_remove_child(stub->_owner, stub);
	handle->data = NULL;
	_taskstub_cleanup(stub);
}

static void _taskstub_on_process_exit(uv_process_t* process, int64_t status, int terminationSignal)
{
	tf_taskstub_t* stub = process->data;
	JSContext* context = tf_task_get_context(stub->_owner);
	if (!JS_IsUndefined(stub->_on_exit))
	{
		JSValue ref = JS_DupValue(context, stub->_on_exit);
		JSValue argv[] = { JS_NewInt32(context, status), JS_NewInt32(context, terminationSignal) };
		JSValue result = JS_Call(context, stub->_on_exit, JS_NULL, 2, argv);
		tf_util_report_error(context, result);
		JS_FreeValue(context, result);
		JS_FreeValue(context, argv[0]);
		JS_FreeValue(context, argv[1]);
		JS_FreeValue(context, ref);
	}
	if (stub->_stream)
	{
		tf_packetstream_destroy(stub->_stream);
		stub->_stream = NULL;
	}
	uv_close((uv_handle_t*)process, _taskstub_on_handle_close);
}

static JSValue _taskstub_getExports(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv)
{
	tf_taskstub_t* stub = JS_GetOpaque(this_val, _classId);
	promiseid_t promise = -1;
	JSValue result = tf_task_allocate_promise(stub->_owner, &promise);
	tf_task_send_promise_message(stub->_owner, stub, kGetExports, promise, JS_UNDEFINED);
	return result;
}

static JSValue _taskstub_setImports(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv)
{
	tf_taskstub_t* stub = JS_GetOpaque(this_val, _classId);
	void* buffer;
	size_t size;
	tf_serialize_store(tf_task_get(context), stub, &buffer, &size, argv[0]);
	tf_packetstream_send(stub->_stream, kSetImports, (char*)buffer, size);
	tf_free(buffer);
	return JS_UNDEFINED;
}

static JSValue _taskstub_loadFile(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv)
{
	tf_taskstub_t* stub = JS_GetOpaque(this_val, _classId);
	void* buffer;
	size_t size;
	tf_serialize_store(tf_task_get(context), stub, &buffer, &size, argv[0]);
	tf_packetstream_send(stub->_stream, kLoadFile, (char*)buffer, size);
	tf_free(buffer);
	return JS_UNDEFINED;
}

static JSValue _taskstub_get_on_exit(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv)
{
	tf_taskstub_t* stub = JS_GetOpaque(this_val, _classId);
	return JS_DupValue(context, stub->_on_exit);
}

static JSValue _taskstub_set_on_exit(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv)
{
	tf_taskstub_t* stub = JS_GetOpaque(this_val, _classId);
	if (!JS_IsUndefined(stub->_on_exit))
	{
		JS_FreeValue(context, stub->_on_exit);
	}
	stub->_on_exit = JS_DupValue(context, argv[0]);
	return JS_UNDEFINED;
}

static JSValue _taskstub_get_on_error(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv)
{
	tf_taskstub_t* stub = JS_GetOpaque(this_val, _classId);
	return JS_DupValue(context, stub->_on_error);
}

static JSValue _taskstub_set_on_error(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv)
{
	tf_taskstub_t* stub = JS_GetOpaque(this_val, _classId);
	if (!JS_IsUndefined(stub->_on_error))
	{
		JS_FreeValue(context, stub->_on_error);
	}
	stub->_on_error = JS_DupValue(context, argv[0]);
	return JS_UNDEFINED;
}

static JSValue _taskstub_get_on_print(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv)
{
	tf_taskstub_t* stub = JS_GetOpaque(this_val, _classId);
	return JS_DupValue(context, stub->_on_print);
}

static JSValue _taskstub_set_on_print(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv)
{
	tf_taskstub_t* stub = JS_GetOpaque(this_val, _classId);
	if (!JS_IsUndefined(stub->_on_print))
	{
		JS_FreeValue(context, stub->_on_print);
	}
	stub->_on_print = JS_DupValue(context, argv[0]);
	return JS_UNDEFINED;
}

static JSValue _taskstub_activate(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv)
{
	tf_taskstub_t* stub = JS_GetOpaque(this_val, _classId);
	if (stub)
	{
		tf_packetstream_send(stub->_stream, kActivate, 0, 0);
	}
	return JS_NULL;
}

static JSValue _taskstub_execute(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv)
{
	tf_taskstub_t* stub = JS_GetOpaque(this_val, _classId);
	promiseid_t promise = -1;
	JSValue result = tf_task_allocate_promise(stub->_owner, &promise);
	tf_task_send_promise_message(stub->_owner, stub, kExecute, promise, argv[0]);
	return result;
}

JSValue tf_taskstub_kill(tf_taskstub_t* stub)
{
	JSValue result = JS_UNDEFINED;
	if (!tf_task_get_one_proc(stub->_owner))
	{
		uv_process_kill(&stub->_process, SIGKILL);
	}
	else
	{
		promiseid_t promise = -1;
		result = tf_task_allocate_promise(stub->_owner, &promise);
		tf_task_send_promise_message(stub->_owner, stub, kKill, promise, JS_UNDEFINED);
	}
	return result;
}

static JSValue _taskstub_kill(JSContext* context, JSValueConst this_val, int argc, JSValueConst* argv)
{
	tf_taskstub_t* stub = JS_GetOpaque(this_val, _classId);
	return tf_taskstub_kill(stub);
}

void tf_taskstub_on_error(tf_taskstub_t* stub, JSValue error)
{
	JSContext* context = tf_task_get_context(stub->_owner);
	if (!JS_IsUndefined(stub->_on_error))
	{
		JSValue result = JS_Call(context, stub->_on_error, JS_NULL, 1, &error);
		tf_util_report_error(context, result);
		JS_FreeValue(context, result);
	}
}

void tf_taskstub_on_print(tf_taskstub_t* stub, JSValue arguments)
{
	JSContext* context = tf_task_get_context(stub->_owner);
	if (!JS_IsUndefined(stub->_on_print))
	{
		JSValue result = JS_Call(context, stub->_on_print, JS_NULL, 1, &arguments);
		tf_util_report_error(context, result);
		JS_FreeValue(context, result);
	}
}
