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

use {
	crate::{
		pidfd::PidFd,
		syscall_abi::{AtFlags, DirFd, MapFlags, SyscallEntry},
	},
	anyhow::{bail, Context, Result},
	bitflags::bitflags,
	nix::{
		libc,
		sys::{
			personality::{self, Persona},
			ptrace,
			resource::{self, Resource as HostResource},
			signal::Signal as HostSignal,
			wait::{waitpid, WaitPidFlag, WaitStatus},
		},
		unistd::Pid,
	},
	std::{
		env,
		ffi::CString,
		fmt::Debug,
		fs::{self, File},
		io::{self, BufRead, Seek, SeekFrom},
		os::unix::process::CommandExt,
		process::Command,
	},
};

mod maps_file;
mod pidfd;
mod syscall_abi;

// 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,
	mem: File,
	pidfd: PidFd,
}

impl Drop for Process {
	fn drop(&mut self) {
		if let Err(err) = self.terminate() {
			eprintln!("{err}");
		}
	}
}

impl Process {
	fn spawn(cmd: &mut Command) -> Result<Process> {
		unsafe {
			cmd.pre_exec(|| {
				// disable ASLR
				let mut persona = personality::get()?;
				persona.insert(Persona::ADDR_NO_RANDOMIZE);
				personality::set(persona)?;

				// set stack limit (to guarantee top-down address space layout)
				resource::setrlimit(
					HostResource::RLIMIT_STACK,
					Some(
						// not a particularly meaningful number,
						// just plucked from the author's system
						8 * 1024 * 1024,
					),
					Some(
						128 * 1024 * 1024
						// guard page
						- 4096,
					),
				)?;

				ptrace::traceme()?;
				Ok(())
			});
		}

		let child = cmd.spawn()?;

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

		match waitpid(tgid.as_pid(), None).context("Couldn't waitpid on fresh child")? {
			WaitStatus::Stopped(_, HostSignal::SIGTRAP) => {}
			status => bail!("unexpected child state: {status:?}"),
		}

		let pidfd = PidFd::open(tgid.as_pid()).context("Couldn't open child pidfd")?;

		Ok(Process {
			tgid,
			mem: File::open(format!("/proc/{}/mem", tgid.0))
				.context("Couldn't open child memory")?,
			pidfd,
		})
	}

	fn terminate(&self) -> Result<()> {
		match self.pidfd.kill(HostSignal::SIGKILL) {
			Ok(()) | Err(nix::Error::ESRCH) => Ok(()),
			Err(err) => Err(anyhow::Error::from(err).context("Couldn't terminate child")),
		}
	}

	fn read_mem_cstr(&self, ptr: u64) -> Result<CString> {
		let mut mem = io::BufReader::new(&self.mem);
		mem.seek(SeekFrom::Start(ptr))?;
		let mut buf = vec![];
		mem.read_until(0, &mut buf)?;

		Ok(CString::from_vec_with_nul(buf).expect("logic error"))
	}

	fn read_mappings(&self) -> Result<Vec<maps_file::Mapping>> {
		let pid = self.tgid.as_pid();
		let contents = fs::read_to_string(format!("/proc/{pid}/maps"))?;

		let mut mappings = contents
			.lines()
			.map(maps_file::parse_mapping_line)
			.collect::<Result<_, _>>()?;

		for &mut maps_file::Mapping {
			start,
			end,
			inode,
			ref mut pathname,
			..
		} in &mut mappings
		{
			if inode.is_none() {
				let is_special = pathname.starts_with('[') && pathname.ends_with(']');
				assert!(is_special || pathname.is_empty());
				// these won't exist in map_files
				continue;
			}

			let map_path = format!("/proc/{pid}/map_files/{start:x}-{end:x}");
			let target = fs::read_link(&map_path)
				.with_context(|| {
					format!("Cannot readlink({map_path:?}) (expected target: {pathname:?})")
				})?
				.into_os_string()
				.into_string()
				.expect("path is not valid UTF-8");
			assert_eq!(*pathname, maps_file::escape_path(&target), "escaping bug?");
			*pathname = target;
		}

		Ok(mappings)
	}

	fn dump_mappings(&self) -> Result<()> {
		let mappings = self.read_mappings()?;
		let mut mappings = mappings.iter().peekable();
		let mut segments = vec![];

		while let Some(mut last) = mappings.next() {
			let mut segment = vec![];
			segment.push(last);
			while let Some(&next) = mappings.peek() {
				if last.inode != next.inode {
					// not the same file
					break;
				}
				if last.end != next.start || last.offset + last.end - last.start != next.offset {
					// not contiguous
					break;
				}
				last = mappings.next().unwrap();
				segment.push(last);
			}
			segments.push(segment);
		}

		println!("{:#?}", segments);
		Ok(())
	}
}

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

fn main() -> 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.env_clear();
		cmd
	})?;

	process.dump_mappings()?;

	let options = ptrace::Options::PTRACE_O_TRACESYSGOOD
		| ptrace::Options::PTRACE_O_TRACECLONE
		| ptrace::Options::PTRACE_O_EXITKILL;
	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 = match SyscallEntry::from_regs(&process, regs) {
					Ok(entry) => entry,
					Err(err) => panic!("{err}"),
				};

				if !check_syscall(&entry) {
					panic!("invalid syscall {entry:?}");
				}

				match entry {
					SyscallEntry::times { .. } => {}
					_ => println!("{entry:?}"),
				}

				syscall_state = Some(EntryExit::Entry(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;
			}
			(syscall_state, status) => {
				panic!("unknown status {status:?} with syscall_state = {syscall_state:?}")
			}
		}
	}

	Ok(())
}

fn check_syscall(entry: &SyscallEntry) -> bool {
	match *entry {
		SyscallEntry::mmap {
			addr,
			len: _,
			prot: _,
			flags,
			fd,
			off,
		} => {
			if addr % 4096 != 0 {
				return false;
			}

			match fd {
				None => flags.contains(MapFlags::ANONYMOUS) && off == 0,
				Some(_) => {
					flags.intersection(MapFlags::PRIVATE | MapFlags::ANONYMOUS) == MapFlags::PRIVATE
				}
			}
		}
		SyscallEntry::mprotect { addr, len, prot: _ } => addr % 4096 == 0 && len % 4096 == 0,
		SyscallEntry::openat {
			dfd,
			filename: _,
			flags: _,
			mode: _,
		} => dfd == DirFd::Cwd,
		SyscallEntry::newfstatat {
			dfd,
			ref filename,
			statbuf: _,
			flags,
		} => {
			match (dfd, filename.as_bytes()) {
				(_, b"") if !flags.contains(AtFlags::EMPTY_PATH) => {
					// empty path without AT_EMPTY_PATH
					false
				}
				(DirFd::Cwd, _) | (_, b"") => true,
				_ => false,
			}
		}
		_ => true,
	}
}