// SPDX-FileCopyrightText: edef <edef@unfathomable.blue>
// SPDX-License-Identifier: OSL-3.0

use {
	nix::{
		libc,
		sys::{
			ptrace,
			wait::{waitpid, WaitPidFlag, WaitStatus},
		},
		unistd::Pid,
	},
	spawn_ptrace::CommandPtraceSpawn,
	std::{env, io, process::Command},
};

// TODO(edef): consider implementing this in terms of TID?
// tgids are a strict subset of tids
#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
struct Tgid(pub libc::pid_t);

impl Tgid {
	fn as_pid(&self) -> Pid {
		Pid::from_raw(self.0)
	}
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
struct Tid(pub libc::pid_t);

impl Tid {
	fn as_pid(&self) -> Pid {
		Pid::from_raw(self.0)
	}
}

#[derive(Debug)]
struct Process {
	tgid: Tgid,
}

impl Process {
	fn spawn(cmd: &mut Command) -> io::Result<Process> {
		let child = cmd.spawn_ptrace()?;

		// the thread group leader's TID is equal to the TGID
		let tgid = Tgid(child.id() as _);

		Ok(Process { tgid })
	}
}

#[derive(Debug, Copy, Clone)]
struct SyscallEntry {
	number: u64,
	// rdi, rsi, rdx, rcx, r8, r9
	args: [u64; 6],
}

impl SyscallEntry {
	fn from_regs(regs: libc::user_regs_struct) -> SyscallEntry {
		SyscallEntry {
			number: regs.orig_rax,
			args: [regs.rdi, regs.rsi, regs.rdx, regs.rcx, regs.r8, regs.r9],
		}
	}
}

#[derive(Debug, Copy, Clone)]
enum EntryExit {
	/// Process is about to enter a syscall
	Entry(SyscallEntry),
	/// Process is about to exit a syscall
	Exit(SyscallEntry, i64),
}

fn main() -> anyhow::Result<()> {
	let process = Process::spawn(&mut {
		let mut args = env::args();

		// drop argv[0]
		args.next();

		let mut cmd = Command::new(args.next().unwrap());
		for arg in args {
			cmd.arg(arg);
		}

		cmd
	})?;

	let options = ptrace::Options::PTRACE_O_TRACESYSGOOD | ptrace::Options::PTRACE_O_TRACECLONE;
	ptrace::setoptions(process.tgid.as_pid(), options)?;

	// this is always equal to tgid for now,
	// but I'm keeping this separate so it's obvious what has to be tgid
	let tid = Tid(process.tgid.0);

	let mut syscall_state: Option<EntryExit> = None;

	loop {
		ptrace::syscall(tid.as_pid(), None)?;
		if let Some(EntryExit::Exit(..)) = syscall_state {
			// syscall has completed now
			syscall_state = None;
		}

		let status = waitpid(tid.as_pid(), Some(WaitPidFlag::__WALL))?;

		match (syscall_state, status) {
			(None, WaitStatus::PtraceSyscall(event_tid)) => {
				let event_tid = Tid(event_tid.as_raw());
				assert_eq!(tid, event_tid);

				let regs = ptrace::getregs(event_tid.as_pid())?;
				let entry = SyscallEntry::from_regs(regs);

				syscall_state = Some(EntryExit::Entry(entry));

				if !check_syscall(entry) {
					ptrace::kill(event_tid.as_pid())?;
					panic!("unsupported syscall {:?}", entry);
				}
			}
			(Some(EntryExit::Entry(entry)), WaitStatus::PtraceSyscall(event_tid)) => {
				let event_tid = Tid(event_tid.as_raw());
				assert_eq!(tid, event_tid);

				let regs = ptrace::getregs(event_tid.as_pid())?;
				let ret = regs.rax as i64;
				syscall_state = Some(EntryExit::Exit(entry, ret));
			}
			(_, WaitStatus::Exited(event_tid, _)) => {
				let event_tid = Tid(event_tid.as_raw());
				assert_eq!(tid, event_tid);

				// TODO(edef): this only works for main thread
				break;
			}
			_ => panic!(
				"unknown status {:?} with syscall_state = {:?}",
				status, syscall_state
			),
		}
	}

	Ok(())
}

fn check_syscall(entry: SyscallEntry) -> bool {
	match entry.number {
		// read
		0 => {}

		// write
		1 => {}

		// close
		3 => {}

		// mmap
		9 => {}

		// mprotect
		10 => {}

		// brk
		12 => {}

		// rt_sigaction
		13 => {}

		// ioctl
		16 => {}

		// pread64
		17 => {}

		// access
		21 => {}

		// getcwd
		79 => {}

		// readlink
		89 => {}

		// sysinfo
		99 => {}

		// times
		100 => {}

		// arch_prctl
		158 => {}

		// exit_group
		231 => {}

		// openat
		257 => {}

		// newfstatat
		262 => {}

		// prlimit64
		302 => {}
		_ => return false,
	}
	true
}