#include <obs/obs-module.h>
#include <obs/obs-interaction.h>
#include <obs/util/platform.h>

#include <errno.h>
#include <fcntl.h>
#include <pthread.h>
#include <poll.h>
#include <signal.h>
#include <spawn.h>
#include <stdatomic.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/mman.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <time.h>
#include <unistd.h>

#include "protocol.h"

OBS_DECLARE_MODULE()
OBS_MODULE_USE_DEFAULT_LOCALE("obs-webkit-browser", "en-US")
MODULE_EXPORT const char *obs_module_description(void)
{
	return "WebKitGTK browser source for OBS (no CEF)";
}

void owb_docks_initialize(void);
void owb_docks_shutdown(void);

const char *owb_module_text(const char *key)
{
	return obs_module_text(key);
}

extern char **environ;

struct browser_source;

struct renderer_process {
	pid_t pid;
	int control_fd;
	int frame_fd;
	int log_fd;
	int shared_fd;
	struct owb_shared_frames *shared;
	size_t shared_size;
	pthread_t frame_thread;
	pthread_t log_thread;
	bool frame_thread_started;
	bool log_thread_started;
	atomic_bool stopping;
	uint32_t width;
	uint32_t height;
	struct browser_source *owner;
};

struct browser_source {
	obs_source_t *source;
	pthread_mutex_t mutex;
	struct renderer_process *renderer;
	char *url;
	char *css;
	uint32_t width;
	uint32_t height;
	uint32_t fps;
	bool transparent;
	bool shutdown_when_hidden;
	bool visible;
};

static bool write_all_socket(int fd, const void *data, size_t size)
{
	const uint8_t *cursor = data;
	while (size > 0) {
		ssize_t written = send(fd, cursor, size, MSG_NOSIGNAL);
		if (written < 0 && errno == EINTR)
			continue;
		if (written <= 0)
			return false;
		cursor += (size_t)written;
		size -= (size_t)written;
	}
	return true;
}

static void set_close_on_exec(int fd)
{
	int flags = fcntl(fd, F_GETFD);
	if (flags >= 0)
		fcntl(fd, F_SETFD, flags | FD_CLOEXEC);
}

static uint32_t setting_u32(obs_data_t *settings, const char *name, uint32_t minimum,
			    uint32_t maximum)
{
	long long value = obs_data_get_int(settings, name);
	if (value < (long long)minimum)
		return minimum;
	if (value > (long long)maximum)
		return maximum;
	return (uint32_t)value;
}

static void *frame_reader(void *opaque)
{
	struct renderer_process *renderer = opaque;
	size_t stride = (size_t)renderer->width * 4;
	size_t frame_size = owb_frame_size(renderer->width, renderer->height);
	uint8_t *pixels = malloc(frame_size);
	if (!pixels) {
		blog(LOG_ERROR, "[obs-webkit-browser] Unable to allocate a %zu-byte frame", frame_size);
		return NULL;
	}

	while (!atomic_load(&renderer->stopping)) {
		struct pollfd notification = {.fd = renderer->frame_fd, .events = POLLIN};
		int ready = poll(&notification, 1, 200);
		if (ready < 0 && errno == EINTR)
			continue;
		if (ready <= 0)
			continue;
		uint8_t slot;
		ssize_t count = read(renderer->frame_fd, &slot, sizeof(slot));
		if (count < 0 && errno == EINTR)
			continue;
		if (count <= 0)
			break;
		if (slot >= OWB_SHARED_FRAME_SLOTS)
			continue;
		bool stable = false;
		for (unsigned attempt = 0; attempt < 3 && !stable; attempt++) {
			uint64_t before = __atomic_load_n(&renderer->shared->sequence[slot], __ATOMIC_ACQUIRE);
			if (before & 1u)
				continue;
			memcpy(pixels, owb_shared_slot(renderer->shared, slot), frame_size);
			uint64_t after = __atomic_load_n(&renderer->shared->sequence[slot], __ATOMIC_ACQUIRE);
			stable = before == after && !(after & 1u);
		}
		if (!stable)
			continue;
		struct obs_source_frame frame = {0};
		frame.data[0] = pixels;
		frame.linesize[0] = (uint32_t)stride;
		frame.width = renderer->width;
		frame.height = renderer->height;
		frame.timestamp = os_gettime_ns();
		frame.format = VIDEO_FORMAT_BGRA;
		frame.full_range = true;
		obs_source_output_video(renderer->owner->source, &frame);
	}

	free(pixels);
	return NULL;
}

static void *log_reader(void *opaque)
{
	struct renderer_process *renderer = opaque;
	char line[2048];
	size_t used = 0;
	while (!atomic_load(&renderer->stopping)) {
		struct pollfd item = {.fd = renderer->log_fd, .events = POLLIN};
		int ready = poll(&item, 1, 200);
		if (ready < 0 && errno == EINTR)
			continue;
		if (ready <= 0)
			continue;
		char chunk[512];
		ssize_t count = read(renderer->log_fd, chunk, sizeof(chunk));
		if (count <= 0)
			break;
		for (ssize_t index = 0; index < count; index++) {
			char byte = chunk[index];
			if (byte == '\n') {
				line[used] = '\0';
				if (used)
					blog(LOG_INFO, "[obs-webkit-browser renderer] %s", line);
				used = 0;
			} else if (byte != '\r' && used + 1 < sizeof(line)) {
				line[used++] = byte;
			}
		}
	}
	if (used) {
		line[used] = '\0';
		blog(LOG_INFO, "[obs-webkit-browser renderer] %s", line);
	}
	close(renderer->log_fd);
	renderer->log_fd = -1;
	return NULL;
}

static void terminate_renderer_group(pid_t pid)
{
	if (pid <= 0)
		return;
	kill(-pid, SIGTERM);
	for (unsigned attempt = 0; attempt < 200; attempt++) {
		pid_t result = waitpid(pid, NULL, WNOHANG);
		if (result == pid || (result < 0 && errno == ECHILD))
			return;
		struct timespec pause = {.tv_nsec = 10000000L};
		nanosleep(&pause, NULL);
	}
	kill(-pid, SIGKILL);
	while (waitpid(pid, NULL, 0) < 0 && errno == EINTR) {
	}
}

static void stop_renderer_locked(struct browser_source *browser)
{
	struct renderer_process *renderer = browser->renderer;
	if (!renderer)
		return;
	browser->renderer = NULL;
	atomic_store(&renderer->stopping, true);

	if (renderer->control_fd >= 0) {
		shutdown(renderer->control_fd, SHUT_RDWR);
		close(renderer->control_fd);
		renderer->control_fd = -1;
	}
	terminate_renderer_group(renderer->pid);
	if (renderer->frame_fd >= 0) {
		close(renderer->frame_fd);
		renderer->frame_fd = -1;
	}
	if (renderer->frame_thread_started)
		pthread_join(renderer->frame_thread, NULL);
	if (renderer->log_thread_started)
		pthread_join(renderer->log_thread, NULL);
	else if (renderer->log_fd >= 0)
		close(renderer->log_fd);
	if (renderer->shared && renderer->shared != MAP_FAILED)
		munmap(renderer->shared, renderer->shared_size);
	if (renderer->shared_fd >= 0)
		close(renderer->shared_fd);
	free(renderer);
}

static bool send_initial_config(struct browser_source *browser, struct renderer_process *renderer)
{
	size_t url_length = strlen(browser->url);
	size_t css_length = strlen(browser->css);
	if (url_length > OWB_MAX_STRING || css_length > OWB_MAX_STRING)
		return false;
	struct owb_config_header header = {
		.magic = OWB_CONFIG_MAGIC,
		.version = OWB_PROTOCOL_VERSION,
		.width = browser->width,
		.height = browser->height,
		.fps = browser->fps,
		.transparent = browser->transparent,
		.hardware_acceleration = true,
		.url_length = (uint32_t)url_length,
		.css_length = (uint32_t)css_length,
	};
	return write_all_socket(renderer->control_fd, &header, sizeof(header)) &&
	       write_all_socket(renderer->control_fd, browser->url, url_length) &&
	       write_all_socket(renderer->control_fd, browser->css, css_length);
}

static bool start_renderer_locked(struct browser_source *browser)
{
	if (browser->renderer || (browser->shutdown_when_hidden && !browser->visible))
		return true;

	char *helper = obs_module_file("obs-webkit-renderer");
	if (!helper) {
		blog(LOG_ERROR, "[obs-webkit-browser] Renderer helper is missing from the plugin data directory");
		return false;
	}
	if (access(helper, X_OK) != 0) {
		blog(LOG_ERROR, "[obs-webkit-browser] Renderer is not executable: %s", helper);
		bfree(helper);
		return false;
	}

	int control[2] = {-1, -1};
	int frames[2] = {-1, -1};
	int logs[2] = {-1, -1};
	int shared_fd = -1;
	struct owb_shared_frames *shared = MAP_FAILED;
	size_t shared_size = owb_shared_size(browser->width, browser->height);
	shared_fd = memfd_create("obs-webkit-browser-frames", MFD_CLOEXEC);
	if (shared_fd < 0 || ftruncate(shared_fd, (off_t)shared_size) != 0) {
		blog(LOG_ERROR, "[obs-webkit-browser] Shared frame allocation failed: %s", strerror(errno));
		goto fail;
	}
	shared = mmap(NULL, shared_size, PROT_READ | PROT_WRITE, MAP_SHARED, shared_fd, 0);
	if (shared == MAP_FAILED) {
		blog(LOG_ERROR, "[obs-webkit-browser] Shared frame mapping failed: %s", strerror(errno));
		goto fail;
	}
	memset(shared, 0, sizeof(*shared));
	shared->magic = OWB_SHARED_MAGIC;
	shared->version = OWB_PROTOCOL_VERSION;
	shared->width = browser->width;
	shared->height = browser->height;
	shared->stride = browser->width * 4;
	shared->slot_count = OWB_SHARED_FRAME_SLOTS;
	if (socketpair(AF_UNIX, SOCK_STREAM, 0, control) != 0 || pipe(frames) != 0 || pipe(logs) != 0) {
		blog(LOG_ERROR, "[obs-webkit-browser] IPC setup failed: %s", strerror(errno));
		goto fail;
	}
	for (size_t i = 0; i < 2; i++) {
		set_close_on_exec(control[i]);
		set_close_on_exec(frames[i]);
		set_close_on_exec(logs[i]);
	}

	posix_spawn_file_actions_t actions;
	posix_spawnattr_t attributes;
	posix_spawn_file_actions_init(&actions);
	posix_spawnattr_init(&attributes);
	sigset_t defaults;
	sigemptyset(&defaults);
	sigaddset(&defaults, SIGTERM);
	sigaddset(&defaults, SIGINT);
	sigaddset(&defaults, SIGPIPE);
	posix_spawnattr_setsigdefault(&attributes, &defaults);
	posix_spawnattr_setpgroup(&attributes, 0);
	posix_spawnattr_setflags(&attributes, POSIX_SPAWN_SETSIGDEF | POSIX_SPAWN_SETPGROUP);
	posix_spawn_file_actions_adddup2(&actions, control[1], STDIN_FILENO);
	posix_spawn_file_actions_adddup2(&actions, frames[1], STDOUT_FILENO);
	posix_spawn_file_actions_adddup2(&actions, logs[1], STDERR_FILENO);
	posix_spawn_file_actions_adddup2(&actions, shared_fd, OWB_SHARED_FRAME_FD);
	posix_spawn_file_actions_addclose(&actions, control[0]);
	posix_spawn_file_actions_addclose(&actions, frames[0]);
	posix_spawn_file_actions_addclose(&actions, logs[0]);
	char *argv[] = {helper, NULL};
	pid_t pid = 0;
	int spawn_error = posix_spawn(&pid, helper, &actions, &attributes, argv, environ);
	posix_spawn_file_actions_destroy(&actions);
	posix_spawnattr_destroy(&attributes);
	if (spawn_error != 0) {
		blog(LOG_ERROR, "[obs-webkit-browser] Cannot start renderer: %s", strerror(spawn_error));
		goto fail;
	}

	close(control[1]);
	close(frames[1]);
	close(logs[1]);
	control[1] = frames[1] = logs[1] = -1;
	struct renderer_process *renderer = calloc(1, sizeof(*renderer));
	if (!renderer) {
		terminate_renderer_group(pid);
		goto fail;
	}
	renderer->pid = pid;
	renderer->control_fd = control[0];
	renderer->frame_fd = frames[0];
	renderer->log_fd = logs[0];
	renderer->shared_fd = shared_fd;
	renderer->shared = shared;
	renderer->shared_size = shared_size;
	renderer->width = browser->width;
	renderer->height = browser->height;
	renderer->owner = browser;
	atomic_init(&renderer->stopping, false);
	control[0] = frames[0] = logs[0] = -1;
	shared_fd = -1;
	shared = MAP_FAILED;

	if (!send_initial_config(browser, renderer)) {
		blog(LOG_ERROR, "[obs-webkit-browser] Cannot send configuration to renderer");
		terminate_renderer_group(renderer->pid);
		close(renderer->control_fd);
		close(renderer->frame_fd);
		close(renderer->log_fd);
		munmap(renderer->shared, renderer->shared_size);
		close(renderer->shared_fd);
		free(renderer);
		goto fail;
	}
	if (pthread_create(&renderer->frame_thread, NULL, frame_reader, renderer) == 0)
		renderer->frame_thread_started = true;
	if (pthread_create(&renderer->log_thread, NULL, log_reader, renderer) == 0)
		renderer->log_thread_started = true;
	if (!renderer->frame_thread_started) {
		blog(LOG_ERROR, "[obs-webkit-browser] Cannot start frame reader thread");
		browser->renderer = renderer;
		stop_renderer_locked(browser);
		goto fail;
	}
	browser->renderer = renderer;
	blog(LOG_INFO, "[obs-webkit-browser] Started %ux%u@%u renderer", browser->width,
	     browser->height, browser->fps);
	bfree(helper);
	return true;

fail:
	for (size_t i = 0; i < 2; i++) {
		if (control[i] >= 0)
			close(control[i]);
		if (frames[i] >= 0)
			close(frames[i]);
		if (logs[i] >= 0)
			close(logs[i]);
	}
	bfree(helper);
	if (shared != MAP_FAILED)
		munmap(shared, shared_size);
	if (shared_fd >= 0)
		close(shared_fd);
	return false;
}

static void restart_renderer(struct browser_source *browser)
{
	pthread_mutex_lock(&browser->mutex);
	stop_renderer_locked(browser);
	start_renderer_locked(browser);
	pthread_mutex_unlock(&browser->mutex);
}

static void browser_update(void *data, obs_data_t *settings)
{
	struct browser_source *browser = data;
	const char *url = obs_data_get_string(settings, "url");
	const char *css = obs_data_get_string(settings, "css");
	char *new_url = bstrdup(url && *url ? url : "about:blank");
	char *new_css = bstrdup(css ? css : "");
	pthread_mutex_lock(&browser->mutex);
	bfree(browser->url);
	bfree(browser->css);
	browser->url = new_url;
	browser->css = new_css;
	browser->width = setting_u32(settings, "width", 2, 8192);
	browser->height = setting_u32(settings, "height", 2, 8192);
	browser->fps = setting_u32(settings, "fps", 1, 60);
	browser->transparent = obs_data_get_bool(settings, "transparent");
	browser->shutdown_when_hidden = obs_data_get_bool(settings, "shutdown_when_hidden");
	stop_renderer_locked(browser);
	start_renderer_locked(browser);
	pthread_mutex_unlock(&browser->mutex);
}

static void *browser_create(obs_data_t *settings, obs_source_t *source)
{
	struct browser_source *browser = bzalloc(sizeof(*browser));
	browser->source = source;
	browser->visible = true;
	pthread_mutex_init(&browser->mutex, NULL);
	browser_update(browser, settings);
	return browser;
}

static void browser_destroy(void *data)
{
	struct browser_source *browser = data;
	pthread_mutex_lock(&browser->mutex);
	stop_renderer_locked(browser);
	pthread_mutex_unlock(&browser->mutex);
	pthread_mutex_destroy(&browser->mutex);
	bfree(browser->url);
	bfree(browser->css);
	bfree(browser);
}

static uint32_t browser_width(void *data)
{
	return ((struct browser_source *)data)->width;
}

static uint32_t browser_height(void *data)
{
	return ((struct browser_source *)data)->height;
}

static void browser_defaults(obs_data_t *settings)
{
	obs_data_set_default_string(settings, "url", "https://example.com");
	obs_data_set_default_int(settings, "width", 1280);
	obs_data_set_default_int(settings, "height", 720);
	obs_data_set_default_int(settings, "fps", 30);
	obs_data_set_default_bool(settings, "transparent", true);
	obs_data_set_default_bool(settings, "shutdown_when_hidden", false);
}

static bool reload_clicked(obs_properties_t *props, obs_property_t *property, void *data)
{
	UNUSED_PARAMETER(props);
	UNUSED_PARAMETER(property);
	restart_renderer(data);
	return false;
}

static obs_properties_t *browser_properties(void *data)
{
	obs_properties_t *props = obs_properties_create();
	obs_properties_add_text(props, "url", obs_module_text("URL"), OBS_TEXT_DEFAULT);
	obs_properties_add_int(props, "width", obs_module_text("Width"), 2, 8192, 1);
	obs_properties_add_int(props, "height", obs_module_text("Height"), 2, 8192, 1);
	obs_properties_add_int(props, "fps", obs_module_text("FPS"), 1, 60, 1);
	obs_properties_add_bool(props, "transparent", obs_module_text("Transparent"));
	obs_properties_add_text(props, "css", obs_module_text("CustomCSS"), OBS_TEXT_MULTILINE);
	obs_properties_add_bool(props, "shutdown_when_hidden", obs_module_text("ShutdownWhenHidden"));
	if (data)
		obs_properties_add_button2(props, "reload", obs_module_text("Reload"), reload_clicked, data);
	return props;
}

static void send_command(struct browser_source *browser, struct owb_command *command)
{
	command->magic = OWB_COMMAND_MAGIC;
	pthread_mutex_lock(&browser->mutex);
	if (browser->renderer && browser->renderer->control_fd >= 0)
		write_all_socket(browser->renderer->control_fd, command, sizeof(*command));
	pthread_mutex_unlock(&browser->mutex);
}

static void browser_mouse_move(void *data, const struct obs_mouse_event *event, bool leave)
{
	struct owb_command command = {.type = OWB_MOUSE_MOVE, .x = event->x, .y = event->y,
				      .value1 = leave, .modifiers = event->modifiers};
	send_command(data, &command);
}

static void browser_mouse_click(void *data, const struct obs_mouse_event *event, int32_t type,
				bool mouse_up, uint32_t click_count)
{
	struct owb_command command = {.type = OWB_MOUSE_CLICK, .x = event->x, .y = event->y,
				      .value1 = type, .value2 = mouse_up,
				      .modifiers = event->modifiers};
	command.text_length = click_count;
	send_command(data, &command);
}

static void browser_mouse_wheel(void *data, const struct obs_mouse_event *event, int x_delta, int y_delta)
{
	struct owb_command command = {.type = OWB_MOUSE_WHEEL, .x = event->x, .y = event->y,
				      .value1 = x_delta, .value2 = y_delta,
				      .modifiers = event->modifiers};
	send_command(data, &command);
}

static void browser_focus(void *data, bool focus)
{
	struct owb_command command = {.type = OWB_FOCUS, .value1 = focus};
	send_command(data, &command);
}

static void browser_key(void *data, const struct obs_key_event *event, bool key_up)
{
	struct owb_command command = {.type = OWB_KEY, .value1 = key_up,
				      .value2 = (int32_t)event->native_vkey,
				      .modifiers = event->modifiers};
	if (event->text) {
		size_t length = strlen(event->text);
		if (length > sizeof(command.text))
			length = sizeof(command.text);
		memcpy(command.text, event->text, length);
		command.text_length = (uint32_t)length;
	}
	send_command(data, &command);
}

static void browser_show(void *data)
{
	struct browser_source *browser = data;
	pthread_mutex_lock(&browser->mutex);
	browser->visible = true;
	start_renderer_locked(browser);
	pthread_mutex_unlock(&browser->mutex);
}

static void browser_hide(void *data)
{
	struct browser_source *browser = data;
	pthread_mutex_lock(&browser->mutex);
	browser->visible = false;
	if (browser->shutdown_when_hidden)
		stop_renderer_locked(browser);
	pthread_mutex_unlock(&browser->mutex);
}

static const char *browser_name(void *unused)
{
	UNUSED_PARAMETER(unused);
	return obs_module_text("WebKitBrowserSource");
}

static struct obs_source_info browser_info = {
	.id = "webkit_browser_source",
	.type = OBS_SOURCE_TYPE_INPUT,
	.output_flags = OBS_SOURCE_ASYNC_VIDEO | OBS_SOURCE_INTERACTION,
	.get_name = browser_name,
	.create = browser_create,
	.destroy = browser_destroy,
	.update = browser_update,
	.get_defaults = browser_defaults,
	.get_properties = browser_properties,
	.get_width = browser_width,
	.get_height = browser_height,
	.show = browser_show,
	.hide = browser_hide,
	.mouse_click = browser_mouse_click,
	.mouse_move = browser_mouse_move,
	.mouse_wheel = browser_mouse_wheel,
	.focus = browser_focus,
	.key_click = browser_key,
};

bool obs_module_load(void)
{
	obs_register_source(&browser_info);
	owb_docks_initialize();
	blog(LOG_INFO, "[obs-webkit-browser] Loaded WebKitGTK browser source");
	return true;
}

void obs_module_unload(void)
{
	owb_docks_shutdown();
}