#include <sys/types.h>
#include <sys/personality.h>
#include <sys/ptrace.h>
#include <sys/wait.h>

#if defined(__aarch64__) || defined(__loongarch__)
#  include <sys/uio.h>
#  include <elf.h>
#endif

#include <dirent.h>
#include <sched.h>
#include <stdio.h>
#include <string.h>
#include <unordered_map>

#include "ptrace_sys.hh"
#include "utils.hh"

enum
{
	i386_EIP = 12, x86_64_RIP = 16, ppc_NIP = 32, arm_PC = 15, aarch64_PC = 32, // See Linux arch/arm64/include/asm/ptrace.h
	riscv_EPC = 0, loongarch_ERA = 33
};

static void arch_adjustPcAfterBreakpoint(unsigned long *regs);
static unsigned long arch_getPcFromRegs(unsigned long *regs);
static int attachLwp(int lwpid);
static unsigned long getPcFromRegs(unsigned long *regs);
static long getRegs(pid_t pid, void *addr, void *regs, size_t len);
static int kill_lwp(unsigned long lwpid, int signo);
static int linux_proc_get_int(pid_t lwpid, const char *field);
static int linux_proc_pid_has_state(pid_t pid, const char *state);
static int linux_proc_pid_is_stopped(pid_t pid);
static int linux_proc_get_tgid(pid_t lwpid);
static long setRegs(pid_t pid, void *addr, void *regs, size_t len);

static void arch_adjustPcAfterBreakpoint(unsigned long *regs)
{
#if defined(__i386__)
	regs[i386_EIP]--;
#elif defined(__x86_64__)
	regs[x86_64_RIP]--;
#elif defined(__powerpc__) || defined(__arm__) || defined(__aarch64__) || defined(__riscv) || defined(__loongarch__)
	// Do nothing
#else
# error Unsupported architecture
#endif
}

static unsigned long arch_getPcFromRegs(unsigned long *regs)
{
	unsigned long out;

#if defined(__i386__)
	out = regs[i386_EIP] - 1;
#elif defined(__x86_64__)
	out = regs[x86_64_RIP] - 1;
#elif defined(__arm__)
	out = regs[arm_PC];
#elif defined(__aarch64__)
	out = regs[aarch64_PC];
#elif defined(__powerpc__)
	out = regs[ppc_NIP];
#elif defined(__riscv)
	out = regs[riscv_EPC];
#elif defined(__loongarch__)
	out = regs[loongarch_ERA];
#else
# error Unsupported architecture
#endif

	return out;
}

int ptrace_sys::attachAll(pid_t pid)
{
	int err;

	/* Attach to PID.  We will check for other threads soon. */
	err = attachLwp(pid);
	if (err != 0)
		error("Cannot attach to process %d\n", pid);

	if (linux_proc_get_tgid(pid) != pid)
	{
		return 0;
	}

	DIR *dir;
	char pathname[128];

	sprintf(pathname, "/proc/%d/task", pid);

	dir = opendir(pathname);

	if (!dir)
	{
		error("Could not open /proc/%d/task.\n", pid);

		return 0;
	}

	/* At this point we attached to the tgid.  Scan the task for
	 * existing threads. */
	int new_threads_found;
	int iterations = 0;

	std::unordered_map<unsigned long, bool> threads;
	threads[pid] = true;

	while (iterations < 2)
	{
		struct dirent *dp;

		new_threads_found = 0;
		/* Add all the other threads.  While we go through the
		 * threads, new threads may be spawned.  Cycle through
		 * the list of threads until we have done two iterations without
		 * finding new threads.  */
		while ((dp = readdir(dir)) != NULL)
		{
			int lwp;

			/* Fetch one lwp.  */
			lwp = strtoul(dp->d_name, NULL, 10);

			/* Is this a new thread?  */
			if (lwp != 0 && threads.find(lwp) == threads.end())
			{
				int err;
				threads[lwp] = true;

				err = attachLwp(lwp);
				if (err != 0)
					warning("Cannot attach to lwp %d\n", lwp);

				new_threads_found++;
			}
		}

		if (!new_threads_found)
			iterations++;
		else
			iterations = 0;

		rewinddir(dir);
	}
	closedir(dir);

	return 0;
}

/* Attach to an inferior process. */
static int attachLwp(int lwpid)
{
	int rv;

	rv = ptrace(PTRACE_ATTACH, lwpid, 0, 0);

	if (rv < 0)
		return errno;
	ptrace(PTRACE_SETOPTIONS, lwpid, 0, PTRACE_O_TRACECLONE | PTRACE_O_TRACEFORK | PTRACE_O_TRACEVFORK);

	if (!linux_proc_pid_is_stopped(lwpid))
	{
		/*
		 * First make sure there is a pending SIGSTOP.  Since we are
		 * already attached, the process can not transition from stopped
		 * to running without a PTRACE_CONT; so we know this signal will
		 * go into the queue.  The SIGSTOP generated by PTRACE_ATTACH is
		 * probably already in the queue (unless this kernel is old
		 * enough to use TASK_STOPPED for ptrace stops); but since
		 * SIGSTOP is not an RT signal, it can only be queued once.  */
		kill_lwp(lwpid, SIGSTOP);
	}

	return 0;
}

int ptrace_sys::cont(pid_t pid, int signal)
{
	return ptrace(PTRACE_CONT, pid, 0, signal);
}

void ptrace_sys::detach(pid_t pid)
{
	ptrace(PTRACE_DETACH, pid, 0, 0);
}

bool ptrace_sys::disable_aslr(void)
{
	int persona = personality(0xffffffff);
	if (persona < 0)
	{
		perror("Can't get personality");
		return false;
	}
	persona |= 0x0040000; /* ADDR_NO_RANDOMIZE */
	if (personality(persona) < 0)
	{
		perror("Can't set personality");
		return false;
	}
	return true;
}

bool ptrace_sys::eventIsForky(int signal, int event)
{
	return (signal == SIGTRAP && (event == PTRACE_EVENT_CLONE || event == PTRACE_EVENT_FORK || event == PTRACE_EVENT_VFORK));
}

bool ptrace_sys::eventIsNewChild(int signal, int event)
{
	return false;
}

int ptrace_sys::follow_child(pid_t pid)
{
	return 0;
}

int ptrace_sys::follow_fork(pid_t pid)
{
	return ptrace(PTRACE_SETOPTIONS, pid, 0, PTRACE_O_TRACECLONE | PTRACE_O_TRACEFORK | PTRACE_O_TRACEVFORK);
}

/* Return non-zero if 'State' of /proc/PID/status contains STATE.  */
static int linux_proc_get_int(pid_t lwpid, const char *field)
{
	size_t field_len = strlen(field);
	FILE *status_file;
	char buf[100];
	int retval = -1;

	snprintf(buf, sizeof(buf), "/proc/%d/status", (int) lwpid);
	status_file = fopen(buf, "r");
	if (status_file == NULL)
	{
		warning("unable to open /proc file '%s'", buf);
		return -1;
	}

	while (fgets(buf, sizeof(buf), status_file))
	{
		if (strncmp(buf, field, field_len) == 0 && buf[field_len] == ':')
		{
			retval = (int) strtol(&buf[field_len + 1], NULL, 10);
			break;
		}
	}

	fclose(status_file);

	return retval;
}

static int linux_proc_pid_has_state(pid_t pid, const char *state)
{
	char buffer[100];
	FILE *procfile;
	int retval;
	int have_state;

	xsnprintf(buffer, sizeof(buffer), "/proc/%d/status", (int ) pid);
	procfile = fopen(buffer, "r");
	if (procfile == NULL)
	{
		warning("unable to open /proc file '%s'", buffer);
		return 0;
	}

	have_state = 0;
	while (fgets(buffer, sizeof(buffer), procfile) != NULL)
	{
		if (strncmp(buffer, "State:", 6) == 0)
		{
			have_state = 1;
			break;
		}
	}
	retval = (have_state && strstr(buffer, state) != NULL);
	fclose(procfile);
	return retval;
}

/* Detect `T (stopped)' in `/proc/PID/status'.
 * Other states including `T (tracing stop)' are reported as false.
 */
static int linux_proc_pid_is_stopped(pid_t pid)
{
	return linux_proc_pid_has_state(pid, "T (stopped)");
}

static int linux_proc_get_tgid(pid_t lwpid)
{
	return linux_proc_get_int(lwpid, "Tgid");
}

int ptrace_sys::get_current_cpu(void)
{
	return sched_getcpu();
}

int ptrace_sys::getEvent(pid_t pid, int status)
{
	return (status >> 16);
}

unsigned long ptrace_sys::getPc(int pid)
{
	unsigned long regs[1024];

	memset(regs, 0, sizeof regs);
	getRegs(pid, NULL, regs, sizeof regs);
	return getPcFromRegs(regs);
}

static unsigned long getPcFromRegs(unsigned long *regs)
{
	return arch_getPcFromRegs(regs);
}

static long getRegs(pid_t pid, void *addr, void *regs, size_t len)
{
#if defined(__aarch64__) || defined(__loongarch__)
	struct iovec iov =
	{	regs, len};
	return ptrace(PTRACE_GETREGSET, pid, (void *)NT_PRSTATUS, &iov);
#else
	(void) len;
	return ptrace(PTRACE_GETREGS, pid, NULL, regs);
#endif
}

static int kill_lwp(unsigned long lwpid, int signo)
{
	/* Use tkill, if possible, in case we are using nptl threads.  If tkill
	 * fails, then we are not using nptl threads and we should be using kill. */

#ifdef __NR_tkill
	{
		static int tkill_failed;

		if (!tkill_failed)
		{
			int ret;

			errno = 0;
			ret = syscall (__NR_tkill, lwpid, signo);
			if (errno != ENOSYS)
			return ret;
			tkill_failed = 1;
		}
	}
#endif

	return kill(lwpid, signo);
}

unsigned long ptrace_sys::peekWord(pid_t pid, unsigned long aligned_addr)
{

	return ptrace(PTRACE_PEEKTEXT, pid, aligned_addr, 0);
}

void ptrace_sys::pokeWord(pid_t pid, unsigned long aligned_addr, unsigned long value)
{
	ptrace(PTRACE_POKETEXT, pid, aligned_addr, value);
}

static long setRegs(pid_t pid, void *addr, void *regs, size_t len)
{
#if defined(__aarch64__) || defined(__loongarch__)
	struct iovec iov =
	{	regs, len};
	return ptrace(PTRACE_SETREGSET, pid, (void *)NT_PRSTATUS, &iov);
#else
	(void) len;
	return ptrace(PTRACE_SETREGS, pid, NULL, regs);
#endif
}

void ptrace_sys::singleStep(pid_t pid)
{
	// Step back one instruction
	unsigned long regs[1024];
	getRegs(pid, NULL, regs, sizeof regs);
	arch_adjustPcAfterBreakpoint(regs);
	setRegs(pid, NULL, regs, sizeof regs);
}

// Skip over this instruction
void ptrace_sys::skipInstruction(pid_t pid)
{
	// Nop on x86, op on PowerPC/ARM
#if defined(__powerpc__) || defined(__arm__) || defined(__aarch64__) || defined(__loongarch__)
	unsigned long regs[1024];

	getRegs(pid, NULL, regs, sizeof regs);

# if defined(__powerpc__)
	regs[ppc_NIP] += 4;
# elif defined(__aarch64__)
	regs[aarch64_PC] += 4;
# elif defined(__arm__)
	regs[arm_PC] += 4;
# else
	regs[loongarch_ERA] += 4;
# endif
	setRegs(pid, NULL, regs, sizeof regs);
#endif
}

void ptrace_sys::tie_process_to_cpu(pid_t pid, int cpu)
{
	// Switching CPU while running will cause icache
	// conflicts. So let's just forbid that.

	int max_cpu = 4096;
	cpu_set_t *set = CPU_ALLOC(max_cpu);
	panic_if(!set, "Can't allocate CPU set!\n");
	CPU_ZERO_S(CPU_ALLOC_SIZE(max_cpu), set);
	CPU_SET(cpu, set);
	panic_if(sched_setaffinity(pid, CPU_ALLOC_SIZE(max_cpu), set) < 0, "Can't set CPU affinity. Coincident won't work");
	CPU_FREE(set);
}

int ptrace_sys::trace_me(void)
{
	return ptrace(PTRACE_TRACEME, 0, 0, 0);
}

pid_t ptrace_sys::wait_all(int *status)
{
	return waitpid(-1, status, __WALL);
}
