// SPDX-License-Identifier: GPL-3.0-or-later
// Copyright (C) 2025 Andrea Mazzoleni

#include "os/portable.h"

#include "snapraid.h"

#define MODE_DEFAULT 0
#define MODE_SPINDOWN 1

static int build_spindown_argv(int argc, char* argv[], char* spindown_argv[5])
{
	char* conf = 0;
	int count = 0;
	int i;

	for (i = 1; i < argc; ++i) {
		if (strcmp(argv[i], "-c") == 0 || strcmp(argv[i], "--conf") == 0) {
			if (i + 1 < argc) {
				conf = argv[++i];
			}
		} else if (strncmp(argv[i], "-c", 2) == 0 && argv[i][2] != 0) {
			conf = argv[i] + 2;
		} else if (strncmp(argv[i], "--conf=", 7) == 0) {
			conf = argv[i] + 7;
		}
	}

	spindown_argv[count++] = argv[0];
	if (conf) {
		spindown_argv[count++] = "-c";
		spindown_argv[count++] = conf;
	}
	spindown_argv[count++] = "down";
	spindown_argv[count] = 0;

	return count;
}

#ifdef _WIN32
#define COMMAND_LINE_MAX 32767

static int needs_quote(const WCHAR* arg)
{
	if (*arg == 0)
		return 1;

	while (*arg) {
		if (*arg == L' ' || *arg == L'\t' || *arg == L'"')
			return 1;
		++arg;
	}

	return 0;
}

#define charcat(c) \
	do { \
		if (pos + 1 >= size) { \
			errno = E2BIG; \
			return -1; \
		} \
		cmd[pos++] = (c); \
	} while (0)

/*
 * Append one argument using the Microsoft C runtime command-line
 * quoting rules. Use this when building a normal CreateProcessW()
 * command line that the target process will parse back into argv[].
 *
 * Do not use this for cmd.exe /c or batch-script arguments: cmd.exe
 * applies an additional command language with different metacharacter
 * parsing rules.
 */
static int argcat(WCHAR* cmd, int size, int pos, const WCHAR* arg)
{
	int has_quote;

	/* space separator */
	if (pos != 0)
		charcat(L' ');

	has_quote = needs_quote(arg);

	if (!has_quote) {
		while (*arg)
			charcat(*arg++);
	} else {
		/* starting quote */
		charcat(L'"');

		while (*arg) {
			int bl = 0;
			while (*arg == L'\\') {
				++arg;
				++bl;
			}

			if (*arg == 0) {
				/* double backslashes before closing quote */
				bl = bl * 2;
				while (bl--)
					charcat(L'\\');
			} else if (*arg == '"') {
				/* double backslashes + escape the quote */
				bl = bl * 2 + 1;
				while (bl--)
					charcat(L'\\');
				charcat(L'"');
				++arg;
			} else {
				/* normal backslashes */
				while (bl--)
					charcat(L'\\');
				charcat(*arg);
				++arg;
			}
		}

		/* ending quote */
		charcat(L'"');
	}

	return pos;
}

static char* argutf8(const WCHAR* arg)
{
	size_t len = wcslen(arg);
	char* utf8_arg;
	DWORD utf8_len;
	DWORD res;

	utf8_len = WideCharToMultiByte(CP_UTF8, 0, arg, len + 1 /* include final null */, 0, 0, 0, 0);

	utf8_arg = malloc(utf8_len);
	if (!utf8_arg)
		return 0;

	res = WideCharToMultiByte(CP_UTF8, 0, arg, len + 1 /* include final null */, utf8_arg, utf8_len, 0, 0);
	if (res != utf8_len) {
		free(utf8_arg);
		return 0;
	}

	return utf8_arg;
}

/* Global variable to store child process handle */
static void* volatile child_process = NULL;

/*
 * Safe atomic access helpers for child_process.
 *
 * The console_handler runs concurrently on an OS-managed thread created by
 * the Windows kernel. Because plain writes and standard volatile access do not
 * guarantee CPU cache visibility (flushing from store buffers) or prevent CPU/compiler
 * out-of-order reordering (due to lack of hardware memory barriers/fences on MinGW),
 * we must use Win32 Interlocked APIs to ensure safe, atomic, fence-protected reads
 * and writes of the child process handle.
 */
static HANDLE child_process_read(void)
{
	return InterlockedCompareExchangePointer(&child_process, NULL, NULL);
}

static void child_process_write(HANDLE handle)
{
	InterlockedExchangePointer(&child_process, handle);
}

/* Console control handler - forwards Ctrl+C, Ctrl+Break to child */
static BOOL WINAPI console_handler(DWORD ctrl_type)
{
	/* if no child, default behavior */
	if (child_process_read() == NULL)
		return FALSE;

	switch (ctrl_type) {
	case CTRL_C_EVENT :
	case CTRL_BREAK_EVENT :
		/*
		 * Return TRUE to prevent parent termination. The child process
		 * will receive these events automatically because it's attached
		 * to the same console, so we don't need to forward them.
		 */
		return TRUE; /* signal handled, don't terminate parent */
	case CTRL_CLOSE_EVENT :
	case CTRL_LOGOFF_EVENT :
	case CTRL_SHUTDOWN_EVENT :
		/*
		 * Return TRUE to prevent our termination while child handles shutdown.
		 * The child receives these events automatically (same console).
		 * Note: Windows will forcibly kill us after a timeout regardless
		 * of returning TRUE: ~5 seconds for CLOSE_EVENT and LOGOFF_EVENT,
		 * ~5-20 seconds for SHUTDOWN_EVENT (configurable in registry).
		 */
		return TRUE; /* signal handled, but Windows will kill us after timeout */
	default :
		return FALSE;
	}
}

int main(int argc, char* argv[])
{
	int wide_argc;
	WCHAR** wide_argv;
	int utf8_argc;
	char** utf8_argv;
	WCHAR app_buffer[COMMAND_LINE_MAX];
	WCHAR cmd_buffer[COMMAND_LINE_MAX];
	DWORD res;
	int i;
	int pos;
	int mode;
	int lock_options;
	int ret;

	(void)argc;
	(void)argv;

	res = GetModuleFileNameW(NULL, app_buffer, COMMAND_LINE_MAX);
	if (res == 0 || res >= COMMAND_LINE_MAX) {
		exit(EXIT_FAILURE);
	}

	wide_argv = CommandLineToArgvW(GetCommandLineW(), &wide_argc);
	if (!wide_argv) {
		exit(EXIT_FAILURE);
	}

	utf8_argc = 0;
	utf8_argv = malloc((wide_argc + 1) * sizeof(void*));
	if (!utf8_argv) {
		exit(EXIT_FAILURE);
	}

	pos = argcat(cmd_buffer, COMMAND_LINE_MAX, 0, app_buffer);
	if (pos < 0) {
		exit(EXIT_FAILURE);
	}

	utf8_argv[0] = argutf8(app_buffer);
	if (!utf8_argv[0]) {
		exit(EXIT_FAILURE);
	}
	utf8_argc = 1;

	mode = MODE_DEFAULT;
	lock_options = 0;
	for (i = 1; i < wide_argc; i++) {
		if (wcscmp(wide_argv[i], L"--lock-options") == 0)
			lock_options = 1;

		if (wcscmp(wide_argv[i], L"-s") == 0 || wcscmp(wide_argv[i], L"--spin-down-on-error") == 0) {
			mode = MODE_SPINDOWN;
		} else {
			pos = argcat(cmd_buffer, COMMAND_LINE_MAX, pos, wide_argv[i]);
			if (pos < 0) {
				exit(EXIT_FAILURE);
			}
			utf8_argv[utf8_argc] = argutf8(wide_argv[i]);
			if (!utf8_argv[utf8_argc]) {
				exit(EXIT_FAILURE);
			}
			++utf8_argc;
		}
	}
	utf8_argv[utf8_argc] = 0;
	cmd_buffer[pos] = 0;

	if (mode == MODE_SPINDOWN && lock_options) {
		fprintf(stderr, "The --spin-down-on-error option cannot be used with --lock-options.\n");
		exit(EXIT_FAILURE);
	}

	LocalFree(wide_argv);

	if (mode == MODE_DEFAULT) {
		ret = snapraid_main(utf8_argc, utf8_argv);
	} else {
		STARTUPINFOW si;
		PROCESS_INFORMATION pi;
		DWORD wait;

		/* install console control handler */
		if (!SetConsoleCtrlHandler(console_handler, TRUE)) {
			fprintf(stderr, "Failed to set console handler\n");
			exit(EXIT_FAILURE);
		}

		ZeroMemory(&si, sizeof(si));
		si.cb = sizeof(si);
		ZeroMemory(&pi, sizeof(pi));

		if (!CreateProcessW(app_buffer, cmd_buffer, NULL, NULL, FALSE, 0, NULL, NULL, &si, &pi)) {
			fprintf(stderr, "Failed to exec SnapRAID\n");
			exit(EXIT_FAILURE);
		}

		/* store child process handle for signal handler */
		child_process_write(pi.hProcess);

		wait = WaitForSingleObject(pi.hProcess, INFINITE);
		if (wait != WAIT_OBJECT_0) {
			fprintf(stderr, "WaitForSingleObject failed: %lu\n", (unsigned long)GetLastError());
			TerminateProcess(pi.hProcess, 1);
			CloseHandle(pi.hThread);
			CloseHandle(pi.hProcess);
			return 1;
		}

		/* clear child process handle */
		child_process_write(NULL);

		if (!GetExitCodeProcess(pi.hProcess, &res)) {
			fprintf(stderr, "GetExitCodeProcess failed: %lu\n", (unsigned long)GetLastError());
			CloseHandle(pi.hThread);
			CloseHandle(pi.hProcess);
			return 1;
		}

		CloseHandle(pi.hThread);
		CloseHandle(pi.hProcess);

		ret = res;

		if (ret != 0) {
			char* spindown_argv[5];
			int spindown_argc = build_spindown_argv(utf8_argc, utf8_argv, spindown_argv);

			/* sync all disks */
			printf("Flush...\n");
			if (sync() != 0)
				printf("WARNING! Failed to flush disks!\n");

			snapraid_main(spindown_argc, spindown_argv);
		}
	}

	for (i = 0; i < utf8_argc; ++i)
		free(utf8_argv[i]);
	free(utf8_argv);

	return ret;
}
#else

char full_argv0[PATH_MAX];

const char* get_argv0(const char* argv0)
{
	ssize_t len = readlink("/proc/self/exe", full_argv0, sizeof(full_argv0) - 1);
	if (len != -1) {
		full_argv0[len] = '\0';
		return full_argv0;
	} else {
#ifdef __APPLE__
		uint32_t size = sizeof(full_argv0);
		if (_NSGetExecutablePath(full_argv0, &size) == 0)
			return full_argv0;
#endif
	}
	return argv0;
}

/* global variable to store child PID for signal handler */
static volatile sig_atomic_t child_pid = 0;

/* signal handler that forwards signals to child */
/* LCOV_EXCL_START */
static void forward_signal(int sig)
{
	if (child_pid > 0) {
		kill((pid_t)child_pid, sig);
	}
}
/* LCOV_EXCL_STOP */

int main(int argc, char* argv[])
{
	int mode;
	int lock_options;
	int i, j;
	int ret;

	mode = MODE_DEFAULT;
	lock_options = 0;
	j = 1;
	for (i = 1; i < argc; ++i) {
		if (strcmp(argv[i], "--lock-options") == 0)
			lock_options = 1;

		if (strcmp(argv[i], "-s") == 0 || strcmp(argv[i], "--spin-down-on-error") == 0) {
			mode = MODE_SPINDOWN;
		} else {
			argv[j] = argv[i];
			++j;
		}
	}
	argc = j;
	argv[argc] = 0;

	if (mode == MODE_SPINDOWN && lock_options) {
		fprintf(stderr, "The --spin-down-on-error option cannot be used with --lock-options.\n");
		exit(EXIT_FAILURE);
	}

	if (mode == MODE_DEFAULT) {
		ret = snapraid_main(argc, argv);
	} else {
		struct sigaction sa;

		/* set up signal handler to ignore signals */
		sa.sa_handler = forward_signal;
		sigemptyset(&sa.sa_mask);
		sa.sa_flags = SA_RESTART; /* restart interrupted system calls */

		/* install forwarding handler for signals */
		sigaction(SIGINT, &sa, NULL); /* Ctrl+C */
		sigaction(SIGQUIT, &sa, NULL); /* Ctrl+\ */
		sigaction(SIGTERM, &sa, NULL); /* termination */
		sigaction(SIGHUP, &sa, NULL); /* hangup */
		sigaction(SIGUSR1, &sa, NULL); /* user-defined 1 */
		sigaction(SIGUSR2, &sa, NULL); /* user-defined 2 */

		sigset_t set, oldset;
		sigemptyset(&set);
		sigaddset(&set, SIGINT);
		sigaddset(&set, SIGQUIT);
		sigaddset(&set, SIGTERM);
		sigaddset(&set, SIGHUP);
		sigaddset(&set, SIGUSR1);
		sigaddset(&set, SIGUSR2);

		/* block signals to prevent signal race before child_pid is set */
		sigprocmask(SIG_BLOCK, &set, &oldset);

		pid_t pid = fork();
		if (pid == -1) {
			/* LCOV_EXCL_START */
			perror("Failed to fork SnapRAID");
			exit(EXIT_FAILURE);
			/* LCOV_EXCL_STOP */
		}

		if (pid == 0) {
			/* child process */

			/* isolate child in its own process group to prevent double signal delivery from terminal */
			setpgid(0, 0);

			/* restore default signal handlers */
			signal(SIGINT, SIG_DFL);
			signal(SIGQUIT, SIG_DFL);
			signal(SIGTERM, SIG_DFL);
			signal(SIGHUP, SIG_DFL);
			signal(SIGUSR1, SIG_DFL);
			signal(SIGUSR2, SIG_DFL);

			/* unblock signals in child */
			sigprocmask(SIG_SETMASK, &oldset, NULL);

			execvp(get_argv0(argv[0]), argv);

			/* here it's an error */
			/* LCOV_EXCL_START */
			perror("Failed to exec SnapRAID");
			exit(EXIT_FAILURE);
			/* LCOV_EXCL_STOP */
		} else {
			/* parent process */
			int status;

			/* set child process group in parent to prevent race condition */
			setpgid(pid, pid);

			/* store child PID so signal handler can forward signals */
			child_pid = pid;

			/* unblock signals in parent */
			sigprocmask(SIG_SETMASK, &oldset, NULL);

			/*
			 * First wait for child termination without reaping (WNOWAIT).
			 * The child enters zombie state, holding its PID reserved in the kernel
			 * so no newly created process can receive the same PID. We clear
			 * child_pid before the final waitpid() reap so the signal forwarder
			 * never sends signals to a recycled PID.
			 */
			siginfo_t info;
			int wret;
			do {
				memset(&info, 0, sizeof(info));
				wret = waitid(P_PID, (id_t)pid, &info, WEXITED | WNOWAIT);
			} while (wret == -1 && errno == EINTR);

			if (wret == -1) {
				/* LCOV_EXCL_START */
				perror("Failed to wait for SnapRAID");
				exit(EXIT_FAILURE);
				/* LCOV_EXCL_STOP */
			}

			/* clear child PID while PID is still reserved in zombie state */
			child_pid = 0;

			do {
				ret = waitpid(pid, &status, 0);
			} while (ret == -1 && errno == EINTR); /* retry if interrupted by signal */

			if (ret == -1) {
				/* LCOV_EXCL_START */
				perror("Failed to reap SnapRAID");
				exit(EXIT_FAILURE);
				/* LCOV_EXCL_STOP */
			}

			/* restore default signal handlers */
			signal(SIGINT, SIG_DFL);
			signal(SIGQUIT, SIG_DFL);
			signal(SIGTERM, SIG_DFL);
			signal(SIGHUP, SIG_DFL);
			signal(SIGUSR1, SIG_DFL);
			signal(SIGUSR2, SIG_DFL);

			if (WIFEXITED(status)) {
				ret = WEXITSTATUS(status);
			} else if (WIFSIGNALED(status)) {
				/* terminated by signal */
				ret = 128 + WTERMSIG(status);
			} else {
				ret = -1;
			}

			if (ret != 0) {
				char* spindown_argv[5];
				int spindown_argc = build_spindown_argv(argc, argv, spindown_argv);

				/* ignore sigpipe to allow spindown even if the terminal is closed */
				signal(SIGPIPE, SIG_IGN);

				/* in *nix DO NOT CALL sync() because we use syncfs() in all single disks befode the down command */

				snapraid_main(spindown_argc, spindown_argv);
			}
		}
	}

	return ret;
}
#endif