//! Read-only window into a *running* CS2 server's memory — the runtime oracle that verifies the
//! offline derivations against ground truth. No injection, no debugger: just `/proc/<pid>/mem` (needs
//! ptrace access — same-user with `yama/ptrace_scope=0`, or `CAP_SYS_PTRACE`).
//!
//! Offline we resolve `.rela.dyn` by hand to recover as-loaded pointer values; the running process is
//! the authority on what those values actually are. So reading the same structures live and comparing
//! confirms both our relocation logic and the struct layout — and, because runtime-populated fields
//! (e.g. `m_pSchemaBinding`) are non-null live but zero on disk, proves we are reading live state.

use anyhow::{Context, Result};
use std::collections::HashMap;
use std::fs::File;
use std::os::unix::fs::FileExt;

pub struct LiveProcess {
    mem: File,
    bases: HashMap<String, u64>, // library filename -> load base (lowest mapping address)
    paths: HashMap<String, String>, // library filename -> the FULL path the process actually mapped
    writable: Vec<(u64, u64)>,   // rw anonymous regions (heap etc.) — where live objects live
    executable: Vec<(u64, u64)>, // r-x regions — where valid code/vtable-slot targets must land
}

/// The pathname field of a `/proc/<pid>/maps` line — everything after the fifth whitespace-delimited field.
/// It must NOT be read as "whatever follows the last whitespace": a mapped path may contain spaces (Steam's
/// default install directory is `.../common/Counter-Strike Global Offensive/...`), and taking the last field
/// yields a fragment that fails the leading-`/` test, so the library is silently dropped from the map and
/// every lookup for it then reports it as not mapped. Empty for an anonymous mapping.
fn maps_path(line: &str) -> &str {
    let mut rest = line;
    for _ in 0..5 {
        rest = rest.trim_start();
        match rest.find(char::is_whitespace) {
            Some(i) => rest = &rest[i..],
            None => return "",
        }
    }
    rest.trim_start()
}

impl LiveProcess {
    pub fn attach(pid: u32) -> Result<Self> {
        let maps = std::fs::read_to_string(format!("/proc/{pid}/maps"))
            .with_context(|| format!("read /proc/{pid}/maps (is pid {pid} running?)"))?;
        let mut bases: HashMap<String, u64> = HashMap::new();
        let mut paths: HashMap<String, String> = HashMap::new();
        let mut writable: Vec<(u64, u64)> = Vec::new();
        let mut executable: Vec<(u64, u64)> = Vec::new();
        for line in maps.lines() {
            // format: START-END perms offset dev inode  path
            let (range, rest) = match line.split_once(' ') {
                Some(x) => x,
                None => continue,
            };
            let perms = rest.split(' ').next().unwrap_or("");
            let path = maps_path(line);
            let Some((start, end)) = range.split_once('-').and_then(|(a, b)| {
                Some((
                    u64::from_str_radix(a, 16).ok()?,
                    u64::from_str_radix(b, 16).ok()?,
                ))
            }) else {
                continue;
            };
            if path.ends_with(".so") && path.starts_with('/') {
                let fname = path.rsplit('/').next().unwrap_or(path).to_string();
                paths
                    .entry(fname.clone())
                    .or_insert_with(|| path.to_string());
                bases
                    .entry(fname)
                    .and_modify(|b| *b = (*b).min(start))
                    .or_insert(start);
            }
            // writable anonymous memory = the heap where runtime objects (entities) are allocated
            if perms.starts_with("rw") && (path.is_empty() || path == "[heap]") {
                writable.push((start, end));
            }
            if perms.starts_with('r') && perms.contains('x') {
                executable.push((start, end));
            }
        }
        executable.sort_unstable();
        let mem = File::open(format!("/proc/{pid}/mem")).with_context(|| {
            format!(
                "open /proc/{pid}/mem — needs ptrace access (yama ptrace_scope=0 or run as root)"
            )
        })?;
        Ok(Self {
            mem,
            bases,
            paths,
            writable,
            executable,
        })
    }

    /// Is `addr` inside an executable mapping? A valid function pointer / vtable slot target must be.
    pub fn is_exec(&self, addr: u64) -> bool {
        self.executable
            .binary_search_by(|&(s, e)| {
                if addr < s {
                    std::cmp::Ordering::Greater
                } else if addr >= e {
                    std::cmp::Ordering::Less
                } else {
                    std::cmp::Ordering::Equal
                }
            })
            .is_ok()
    }

    /// Best-effort read of `n` bytes at runtime `addr` (short/empty on an unmapped page).
    pub fn read_bytes(&self, addr: u64, n: usize) -> Vec<u8> {
        let mut buf = vec![0u8; n];
        let got = self.mem.read_at(&mut buf, addr).unwrap_or(0);
        buf.truncate(got);
        buf
    }

    /// Scan the writable/heap regions for object instances whose vtable pointer is `vtable` — i.e.
    /// live instances of the class that owns that vtable. Returns the object base addresses (an
    /// object's first qword is its vtable pointer). Stops at `max` hits.
    pub fn find_instances(&self, vtable: u64, max: usize) -> Vec<u64> {
        let mut hits = Vec::new();
        let mut buf = vec![0u8; 1 << 20]; // 1 MiB window
        let needle = vtable.to_le_bytes();
        'outer: for &(start, end) in &self.writable {
            let mut addr = start;
            while addr < end {
                let n = ((end - addr) as usize).min(buf.len());
                // `buf` is reused across windows, so scanning past a SHORT read matches stale bytes from the
                // previous window and reports addresses that hold nothing of the sort. Bind both the scan and
                // the advance to what was actually read.
                let got = self.mem.read_at(&mut buf[..n], addr).unwrap_or(0);
                if got < 8 {
                    addr += n as u64;
                    continue;
                }
                // objects are pointer-aligned, so only 8-aligned positions can be a vtable slot
                let mut i = 0;
                while i + 8 <= got {
                    if buf[i..i + 8] == needle {
                        hits.push(addr + i as u64);
                        if hits.len() >= max {
                            break 'outer;
                        }
                    }
                    i += 8;
                }
                // advance by the 8-aligned prefix consumed: a persistently short-reading region still makes
                // progress (never re-reads the same bytes), and the skipped tail is retried on the next pass.
                addr += (got & !7) as u64;
            }
        }
        hits
    }

    /// Load base (slide) of a library — its lowest mapping address. Since CS2 `.so` files link at
    /// vaddr 0, the runtime address of a file vaddr `v` is simply `base + v`.
    pub fn base(&self, lib: &str) -> Option<u64> {
        self.bases.get(lib).copied()
    }

    /// The full path the process actually mapped for `lib` (a basename). The authority on WHICH file of a
    /// given name is loaded when several exist on disk — a game tree can hold the engine's own
    /// `libserver.so` and a loader shim of the same name several directories away.
    pub fn mapped_path(&self, lib: &str) -> Option<&str> {
        self.paths.get(lib).map(String::as_str)
    }

    fn read(&self, addr: u64, buf: &mut [u8]) -> Result<()> {
        self.mem
            .read_exact_at(buf, addr)
            .with_context(|| format!("read {} bytes at {addr:#x}", buf.len()))
    }

    pub fn read_u64(&self, addr: u64) -> Result<u64> {
        let mut b = [0u8; 8];
        self.read(addr, &mut b)?;
        Ok(u64::from_le_bytes(b))
    }

    pub fn read_i32(&self, addr: u64) -> Result<i32> {
        let mut b = [0u8; 4];
        self.read(addr, &mut b)?;
        Ok(i32::from_le_bytes(b))
    }

    pub fn read_u16(&self, addr: u64) -> Result<u16> {
        let mut b = [0u8; 2];
        self.read(addr, &mut b)?;
        Ok(u16::from_le_bytes(b))
    }

    /// NUL-terminated string at runtime `addr` (bounded). Reads may land near an unmapped page, so a
    /// short read is fine — we take whatever came back up to the terminator.
    pub fn read_cstr(&self, addr: u64) -> Result<String> {
        let mut buf = [0u8; 256];
        let n = self.mem.read_at(&mut buf, addr).unwrap_or(0);
        let end = buf[..n].iter().position(|&c| c == 0).unwrap_or(n);
        Ok(String::from_utf8_lossy(&buf[..end]).into_owned())
    }
}

/// Result of a remote call: the return value (RAX) and whether the function returned cleanly to our
/// trap (vs faulting internally on a bad argument).
pub struct CallResult {
    pub rax: u64,
    pub clean_return: bool,
    /// Where the scratch blob was placed, so the caller can read back what the callee wrote into it.
    /// Zero when the call carried no scratch.
    pub scratch_base: u64,
}

/// One argument to a remote call.
///
/// [`Arg::Scratch`] exists because a callee that takes a POINTER needs a structure to point at, and the
/// address of that structure is not known until the call frame is laid out. Naming it relative to the
/// scratch base lets the caller describe "argument 5 points at my blob" without knowing where the blob
/// will land.
#[derive(Clone, Copy)]
pub enum Arg {
    Val(u64),
    /// `scratch_base + addend`.
    Scratch(i64),
}

/// A blob placed in the target's stack scratch before the call.
pub struct Scratch<'a> {
    pub bytes: &'a [u8],
    /// `(offset, addend)` — write `scratch_base + addend` as a little-endian u64 at `offset` in the blob.
    /// This is how a pointer INSIDE the blob becomes absolute; an array-of-pointers argument is otherwise
    /// impossible to build, since every element has to name an address that does not exist yet.
    pub relocs: &'a [(usize, i64)],
}

/// Write `data` into the target at `addr`, a word at a time.
///
/// A trailing partial word is read back and merged rather than zero-filled: `PTRACE_POKEDATA` writes a
/// whole word, so writing the tail without preserving the bytes past it would clobber memory the caller
/// never asked to touch.
unsafe fn poke_bytes(pid: i32, addr: u64, data: &[u8]) -> Result<()> {
    use anyhow::bail;
    let mut i = 0usize;
    while i < data.len() {
        let at = addr + i as u64;
        let n = (data.len() - i).min(8);
        let mut word = if n == 8 {
            [0u8; 8]
        } else {
            // PEEKDATA returns -1 both for an error and for a word whose value IS -1, so errno is the
            // only way to tell them apart and it must be cleared first.
            unsafe { *libc::__errno_location() = 0 };
            let cur = unsafe { libc::ptrace(libc::PTRACE_PEEKDATA, pid, at as usize, 0usize) };
            if cur == -1 && errno() != 0 {
                bail!("PEEKDATA at {at:#x} failed (errno {})", errno());
            }
            (cur as u64).to_le_bytes()
        };
        word[..n].copy_from_slice(&data[i..i + n]);
        let w = u64::from_le_bytes(word) as usize;
        if unsafe { libc::ptrace(libc::PTRACE_POKEDATA, pid, at as usize, w) } < 0 {
            bail!("POKEDATA at {at:#x} failed (errno {})", errno());
        }
        i += n;
    }
    Ok(())
}

/// Call the function at runtime address `func` inside process `pid` with `args` (SysV: up to 6 in
/// registers), via ptrace. Attaches, saves the main thread's registers, sets up a call frame whose
/// return address is 0 (so the function traps on return, where we read RAX), runs it, then restores
/// the thread exactly — the SIGSEGV from the return trap is suppressed. Needs ptrace permission
/// (owned child, or same-user with ptrace_scope=0). UNSAFE: only call leaf-ish functions with valid args.
pub fn call_remote(pid: i32, func: u64, args: &[u64]) -> Result<CallResult> {
    let regs: Vec<Arg> = args.iter().map(|&v| Arg::Val(v)).collect();
    call_remote_ex(pid, func, &regs, &[], None)
}

/// [`call_remote`] plus stack arguments and a scratch blob placed in the target.
///
/// Needed for callees that take more than six integer arguments or a pointer to a structure the caller has
/// to build — neither of which the register-only form can express.
///
/// **Stack geometry**, descending from the interrupted `rsp`, chosen so three regions cannot collide:
/// the 128-byte red zone is left alone (the interrupted frame lives there); the scratch blob sits at
/// `rsp-1024`; the call frame starts at `rsp-2048`, so the callee's own stack — which grows DOWN from
/// there — can never reach the scratch ABOVE it. Entry keeps SysV's `rsp % 16 == 8`, with the return
/// address at `[rsp]` and stack argument *i* at `[rsp + 8 + 8i]`.
pub fn call_remote_ex(
    pid: i32,
    func: u64,
    regs_in: &[Arg],
    stack_in: &[Arg],
    scratch: Option<Scratch<'_>>,
) -> Result<CallResult> {
    use anyhow::bail;
    let dbg = std::env::var("SOURCE2ROSETTA_DBG").is_ok();
    if regs_in.len() > 6 {
        bail!("{} register arguments; SysV has 6", regs_in.len());
    }
    if let Some(s) = &scratch {
        // The blob lives in the 1 KiB between the frame and the red zone. Refuse rather than silently
        // overlap the call frame, which would corrupt the return address mid-call.
        if s.bytes.len() > 768 {
            bail!(
                "scratch blob is {} bytes; the reserved window is 768",
                s.bytes.len()
            );
        }
    }
    unsafe {
        if libc::ptrace(libc::PTRACE_ATTACH, pid, 0usize, 0usize) < 0 {
            bail!(
                "PTRACE_ATTACH {pid} failed (errno {}) — need ptrace permission",
                errno()
            );
        }
        let mut status = 0i32;
        if libc::waitpid(pid, &mut status, 0) < 0 {
            libc::ptrace(libc::PTRACE_DETACH, pid, 0usize, 0usize);
            bail!("waitpid(attach) failed");
        }
        if dbg {
            eprintln!(
                "[call] attached; stop status {status:#x} (stopped={})",
                libc::WIFSTOPPED(status)
            );
        }
        let mut saved: libc::user_regs_struct = std::mem::zeroed();
        if libc::ptrace(libc::PTRACE_GETREGS, pid, 0usize, &mut saved as *mut _) < 0 {
            libc::ptrace(libc::PTRACE_DETACH, pid, 0usize, 0usize);
            bail!("PTRACE_GETREGS failed");
        }
        let restore = |saved: &libc::user_regs_struct| {
            libc::ptrace(libc::PTRACE_SETREGS, pid, 0usize, saved as *const _);
            libc::ptrace(libc::PTRACE_DETACH, pid, 0usize, 0usize);
        };

        let mut regs = saved;
        // If we attached mid-syscall, orig_rax holds the syscall number and the kernel would run its
        // syscall-restart logic on our injected rip. Setting it to -1 says "no syscall in progress".
        regs.orig_rax = u64::MAX;

        // Place the scratch blob first: every Arg::Scratch resolves against its base.
        let scratch_base = (saved.rsp - 1024) & !0xfu64;
        if let Some(s) = &scratch {
            let mut blob = s.bytes.to_vec();
            for &(off, addend) in s.relocs {
                let Some(dst) = blob.get_mut(off..off + 8) else {
                    restore(&saved);
                    bail!(
                        "scratch reloc at {off} runs past the {}-byte blob",
                        s.bytes.len()
                    );
                };
                dst.copy_from_slice(&scratch_base.wrapping_add(addend as u64).to_le_bytes());
            }
            if let Err(e) = poke_bytes(pid, scratch_base, &blob) {
                restore(&saved);
                return Err(e.context(format!("placing scratch at {scratch_base:#x}")));
            }
        }
        let resolve = |a: Arg| match a {
            Arg::Val(v) => v,
            Arg::Scratch(addend) => scratch_base.wrapping_add(addend as u64),
        };

        let slots = [
            &mut regs.rdi as *mut u64,
            &mut regs.rsi,
            &mut regs.rdx,
            &mut regs.rcx,
            &mut regs.r8,
            &mut regs.r9,
        ];
        for (i, &a) in regs_in.iter().enumerate() {
            *slots[i] = resolve(a);
        }
        // Call frame well below the scratch, so the callee's downward stack growth cannot reach it. Write a
        // return address of 0 and keep SysV's `rsp % 16 == 8` at function entry.
        let mut sp = (saved.rsp - 2048) & !0xfu64;
        sp -= 8;
        if libc::ptrace(libc::PTRACE_POKEDATA, pid, sp as usize, 0usize) < 0 {
            restore(&saved);
            bail!(
                "POKEDATA(return addr) at {sp:#x} failed (errno {})",
                errno()
            );
        }
        // Stack arguments sit immediately above the return address, which is where the callee reads them.
        for (i, &a) in stack_in.iter().enumerate() {
            let at = sp + 8 + 8 * i as u64;
            if libc::ptrace(libc::PTRACE_POKEDATA, pid, at as usize, resolve(a) as usize) < 0 {
                restore(&saved);
                bail!(
                    "POKEDATA(stack arg {i}) at {at:#x} failed (errno {})",
                    errno()
                );
            }
        }
        let wrote = libc::ptrace(libc::PTRACE_PEEKDATA, pid, sp as usize, 0usize);
        regs.rsp = sp;
        regs.rip = func;
        if libc::ptrace(libc::PTRACE_SETREGS, pid, 0usize, &regs as *const _) < 0 {
            restore(&saved);
            bail!("PTRACE_SETREGS failed");
        }
        if dbg {
            eprintln!(
                "[call] rip={func:#x} rsp={sp:#x} rdi={:#x} retaddr-slot={wrote:#x} (want 0)",
                regs.rdi
            );
        }

        // Run, absorbing any spurious signals, until the function returns into our null trap.
        loop {
            libc::ptrace(libc::PTRACE_CONT, pid, 0usize, 0usize);
            if libc::waitpid(pid, &mut status, 0) < 0 || !libc::WIFSTOPPED(status) {
                restore(&saved);
                bail!("target vanished mid-call (status {status:#x})");
            }
            let sig = libc::WSTOPSIG(status);
            let mut cur: libc::user_regs_struct = std::mem::zeroed();
            libc::ptrace(libc::PTRACE_GETREGS, pid, 0usize, &mut cur as *mut _);
            if dbg {
                eprintln!(
                    "[call] stop sig={sig} rip={:#x} rax={:#x}",
                    cur.rip, cur.rax
                );
            }
            if cur.rip == 0 {
                let r = CallResult {
                    rax: cur.rax,
                    clean_return: true,
                    scratch_base,
                };
                restore(&saved);
                return Ok(r);
            }
            if sig == libc::SIGSEGV || sig == libc::SIGILL || sig == libc::SIGBUS {
                let r = CallResult {
                    rax: cur.rax,
                    clean_return: false,
                    scratch_base,
                };
                restore(&saved);
                return Ok(r);
            }
            // any other signal (SIGSTOP/timer/…): swallow it and keep running the call
        }
    }
}

fn errno() -> i32 {
    unsafe { *libc::__errno_location() }
}

#[cfg(test)]
mod tests {
    use super::maps_path;

    #[test]
    fn maps_path_survives_spaces_in_the_mapped_path() {
        // Steam's default install directory contains spaces; taking the last whitespace-delimited field
        // yields "Offensive/..." which fails the leading-`/` test, so the library silently vanishes from
        // the map and the live oracle reports it as not mapped.
        let spaced = "7f1a2b000000-7f1a2c000000 r-xp 00000000 08:01 12345      \
                      /home/cs2/.steam/SteamApps/common/Counter-Strike Global Offensive/game/csgo/bin/linuxsteamrt64/libserver.so";
        assert_eq!(
            maps_path(spaced),
            "/home/cs2/.steam/SteamApps/common/Counter-Strike Global Offensive/game/csgo/bin/linuxsteamrt64/libserver.so"
        );

        // The no-space case must be unchanged.
        let plain = "7f1a2b000000-7f1a2c000000 r-xp 00000000 08:01 12345      /home/snake/game/csgo/bin/linuxsteamrt64/libserver.so";
        assert_eq!(
            maps_path(plain),
            "/home/snake/game/csgo/bin/linuxsteamrt64/libserver.so"
        );

        // An anonymous mapping has no pathname — it must read as EMPTY, since that is what marks the
        // writable heap regions the instance scan walks.
        assert_eq!(
            maps_path("24557800000-24597800000 rw-p 00000000 00:00 0 "),
            ""
        );
        assert_eq!(
            maps_path("29618000-29639000 rw-p 00000000 00:00 0     [heap]"),
            "[heap]"
        );
    }
}