//! The two enforcement passes: ProtoSSL cert gates (once per pid) and the
//! CardsDLL store patches (every tick).
//!
//! Both are written against the [`Memory`] trait rather than `/proc` directly, so
//! the log lines and their order — which the launcher reads — are unit-testable
//! without a live FIFA client.

use std::collections::HashSet;
use std::io;

use crate::{
    guard_state_after, guarded_action, hex, live_addr, GuardState, GuardedAction,
    EMPTY_MYPACKS_RESOLVER_CAPABILITY, EMPTY_MYPACKS_RESOLVER_VERSION, GATE1, GATE1_ORIG,
    GATE1_PATCH, GATE2, GATE2_ORIG, GATE2_PATCH, MAX_PATCH_LEN, RESOLVER_GUARD_VA, STORE_PATCHES,
    STORE_PATCHES_GUARDED,
};

/// Byte-level access to one client process's address space.
pub trait Memory {
    /// The pid being patched; it appears in every log line.
    fn pid(&self) -> u32;
    /// Fill `buf` from virtual address `va`. An error means "not mapped (yet)".
    fn read(&self, va: u64, buf: &mut [u8]) -> io::Result<()>;
    /// Write `data` at virtual address `va`.
    fn write(&self, va: u64, data: &[u8]) -> io::Result<()>;
}

/// Apply the two ProtoSSL cert gates, once per pid.
///
/// Returns `false` when the gates could not be read — the packer has not mapped
/// that code yet, which is the Python's `continue`, not an error to report.
pub fn cert_pass<M: Memory>(
    mem: &M,
    log: &mut impl FnMut(&str),
    patched: &mut HashSet<u32>,
) -> bool {
    let pid = mem.pid();
    // Sized from the patterns themselves; the initial contents are overwritten by
    // the reads and are never compared unless both reads succeed.
    let mut g2 = GATE2_ORIG;
    let mut g1 = GATE1_ORIG;
    if mem.read(GATE2, &mut g2).is_err() || mem.read(GATE1, &mut g1).is_err() {
        return false;
    }

    if g2 == GATE2_PATCH && g1 == GATE1_PATCH {
        log(&format!("pid {pid}: cert gates already patched"));
        patched.insert(pid);
    } else if g2 == GATE2_ORIG && g1 == GATE1_ORIG {
        match mem
            .write(GATE2, &GATE2_PATCH)
            .and_then(|()| mem.write(GATE1, &GATE1_PATCH))
        {
            Ok(()) => {
                log(&format!("pid {pid}: PATCHED cert gates"));
                patched.insert(pid);
            }
            Err(e) => log(&format!("pid {pid}: cert patch write failed: {e}")),
        }
    }
    // Anything else is a build we do not recognise: left alone, as in the Python.
    true
}

/// Re-apply every store patch whose live bytes have drifted, then the guarded
/// patch, then report the resolver-guard capability once per pid.
///
/// An `Err` is a read or write that failed outside the guarded site's own
/// handling; it aborts the rest of this pid's pass for this tick, exactly like
/// the Python's enclosing `try`.
pub fn enforce_store_patches<M: Memory>(
    mem: &M,
    cbase: u64,
    log: &mut impl FnMut(&str),
    guard_reported: &mut HashSet<u32>,
) -> io::Result<()> {
    let pid = mem.pid();

    for (va, data) in STORE_PATCHES {
        let live = live_addr(cbase, va);
        let mut buf = [0u8; MAX_PATCH_LEN];
        let cur = &mut buf[..data.len()];
        mem.read(live, cur)?;
        if cur != data {
            mem.write(live, data)?;
            log(&format!("pid {pid}: ENFORCED store patch @ {live:#x}"));
        }
    }

    for (va, orig, patch) in STORE_PATCHES_GUARDED {
        let live = live_addr(cbase, va);
        let mut before = [0u8; MAX_PATCH_LEN];
        mem.read(live, &mut before[..patch.len()])?;
        let cur = &before[..patch.len()];

        let mut wrote_ok = true;
        // The Python starts with `cur_after = cur`, which only matters on the
        // branches that never re-read.
        let mut after = [0u8; MAX_PATCH_LEN];
        after[..patch.len()].copy_from_slice(cur);
        let mut after_len = patch.len();

        match guarded_action(cur, orig, patch) {
            GuardedAction::Patch => {
                match mem.write(live, patch) {
                    Ok(()) => log(&format!(
                        "pid {pid}: ENFORCED guarded store patch @ {live:#x} (JNZ->JG, empty My Packs)"
                    )),
                    Err(e) => {
                        wrote_ok = false;
                        log(&format!(
                            "pid {pid}: guarded patch write failed @ {live:#x}: {e}"
                        ));
                    }
                }
                if wrote_ok && mem.read(live, &mut after[..patch.len()]).is_err() {
                    // The Python's `cur_after = b""`: unverifiable, so not verified.
                    after_len = 0;
                }
            }
            GuardedAction::Skip => log(&format!(
                "pid {pid}: SKIP guarded patch @ {live:#x}: unexpected {} (build mismatch)",
                hex(cur)
            )),
            // Already patched; nothing to write.
            GuardedAction::Noop => {}
        }

        if va == RESOLVER_GUARD_VA && !guard_reported.contains(&pid) {
            let state = guard_state_after(cur, orig, patch, wrote_ok, &after[..after_len]);
            if state == GuardState::Verified {
                // PARSED BY THE LAUNCHER (fifa17_capability::parse_capability_line):
                // this line must keep both `verified capability` and the
                // `fifa17.empty_mypacks_resolver=<version>` token verbatim.
                log(&format!(
                    "[store-guard] verified capability {EMPTY_MYPACKS_RESOLVER_CAPABILITY}={EMPTY_MYPACKS_RESOLVER_VERSION} fifa_pid={pid}"
                ));
            } else {
                log(&format!(
                    "[store-guard] guard status={state} fifa_pid={pid} (no capability advertised)"
                ));
            }
            guard_reported.insert(pid);
        }
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::cell::RefCell;
    use std::collections::BTreeMap;

    const CBASE: u64 = 0x7f2a11c00000;
    const GUARD_LIVE: u64 = CBASE + 0x14858;

    /// Sparse fake address space: an unmapped byte reads as `NotFound`, mirroring
    /// `/proc/<pid>/mem` refusing an address the packer has not produced yet.
    struct FakeMemory {
        bytes: RefCell<BTreeMap<u64, u8>>,
        /// Writes to these addresses fail.
        fail_writes: Vec<u64>,
        /// Writes to these addresses report success but change nothing (the
        /// VERIFY_FAILED shape).
        swallow_writes: Vec<u64>,
        /// Reads of these addresses fail even when mapped.
        fail_reads: Vec<u64>,
    }

    impl FakeMemory {
        fn new() -> Self {
            Self {
                bytes: RefCell::new(BTreeMap::new()),
                fail_writes: Vec::new(),
                swallow_writes: Vec::new(),
                fail_reads: Vec::new(),
            }
        }

        fn map(self, va: u64, bytes: &[u8]) -> Self {
            {
                let mut mem = self.bytes.borrow_mut();
                for (i, b) in bytes.iter().enumerate() {
                    mem.insert(va + i as u64, *b);
                }
            }
            self
        }

        /// Every store-patch site mapped with filler that is neither the patch
        /// nor (for the guarded site) the original.
        fn with_store_sites(mut self, filler: u8) -> Self {
            for (va, data) in STORE_PATCHES {
                self = self.map(live_addr(CBASE, va), &vec![filler; data.len()]);
            }
            for (va, _, patch) in STORE_PATCHES_GUARDED {
                self = self.map(live_addr(CBASE, va), &vec![filler; patch.len()]);
            }
            self
        }

        fn at(&self, va: u64, len: usize) -> Vec<u8> {
            let mem = self.bytes.borrow();
            (0..len as u64).map(|i| mem[&(va + i)]).collect()
        }
    }

    impl Memory for FakeMemory {
        fn pid(&self) -> u32 {
            4242
        }

        fn read(&self, va: u64, buf: &mut [u8]) -> io::Result<()> {
            if self.fail_reads.contains(&va) {
                return Err(io::Error::from(io::ErrorKind::PermissionDenied));
            }
            let mem = self.bytes.borrow();
            for (i, slot) in buf.iter_mut().enumerate() {
                *slot = *mem
                    .get(&(va + i as u64))
                    .ok_or_else(|| io::Error::from(io::ErrorKind::NotFound))?;
            }
            Ok(())
        }

        fn write(&self, va: u64, data: &[u8]) -> io::Result<()> {
            if self.fail_writes.contains(&va) {
                return Err(io::Error::from(io::ErrorKind::PermissionDenied));
            }
            if self.swallow_writes.contains(&va) {
                return Ok(());
            }
            let mut mem = self.bytes.borrow_mut();
            for (i, b) in data.iter().enumerate() {
                mem.insert(va + i as u64, *b);
            }
            Ok(())
        }
    }

    /// Collects log lines so the contract strings can be asserted verbatim.
    #[derive(Default)]
    struct Lines(Vec<String>);

    impl Lines {
        fn sink(&mut self) -> impl FnMut(&str) + '_ {
            |line: &str| self.0.push(line.to_string())
        }
    }

    #[test]
    fn cert_pass_defers_while_the_code_is_not_mapped() {
        let mem = FakeMemory::new();
        let mut lines = Lines::default();
        let mut patched = HashSet::new();
        assert!(!cert_pass(&mem, &mut lines.sink(), &mut patched));
        assert!(lines.0.is_empty(), "{:?}", lines.0);
        assert!(patched.is_empty());
    }

    #[test]
    fn cert_pass_defers_when_only_the_first_gate_is_mapped() {
        let mem = FakeMemory::new().map(GATE2, &GATE2_ORIG);
        let mut lines = Lines::default();
        let mut patched = HashSet::new();
        assert!(!cert_pass(&mem, &mut lines.sink(), &mut patched));
        assert!(lines.0.is_empty());
        assert!(patched.is_empty());
    }

    #[test]
    fn cert_pass_writes_both_gates_once() {
        let mem = FakeMemory::new()
            .map(GATE2, &GATE2_ORIG)
            .map(GATE1, &GATE1_ORIG);
        let mut lines = Lines::default();
        let mut patched = HashSet::new();
        assert!(cert_pass(&mem, &mut lines.sink(), &mut patched));
        assert_eq!(lines.0, vec!["pid 4242: PATCHED cert gates"]);
        assert_eq!(mem.at(GATE2, 3), GATE2_PATCH);
        assert_eq!(mem.at(GATE1, 6), GATE1_PATCH);
        assert!(patched.contains(&4242));
    }

    #[test]
    fn cert_pass_recognises_an_already_patched_client() {
        let mem = FakeMemory::new()
            .map(GATE2, &GATE2_PATCH)
            .map(GATE1, &GATE1_PATCH);
        let mut lines = Lines::default();
        let mut patched = HashSet::new();
        assert!(cert_pass(&mem, &mut lines.sink(), &mut patched));
        assert_eq!(lines.0, vec!["pid 4242: cert gates already patched"]);
        assert!(patched.contains(&4242));
    }

    #[test]
    fn cert_pass_reports_a_write_failure_and_stays_unpatched() {
        let mut mem = FakeMemory::new()
            .map(GATE2, &GATE2_ORIG)
            .map(GATE1, &GATE1_ORIG);
        mem.fail_writes.push(GATE2);
        let mut lines = Lines::default();
        let mut patched = HashSet::new();
        assert!(cert_pass(&mem, &mut lines.sink(), &mut patched));
        assert_eq!(lines.0.len(), 1);
        assert!(
            lines.0[0].starts_with("pid 4242: cert patch write failed: "),
            "{}",
            lines.0[0]
        );
        // Not recorded as patched, so the next tick tries again.
        assert!(patched.is_empty());
        assert_eq!(mem.at(GATE2, 3), GATE2_ORIG);
    }

    #[test]
    fn cert_pass_leaves_an_unrecognised_build_alone() {
        let mem = FakeMemory::new()
            .map(GATE2, &[0x55, 0x48, 0x89])
            .map(GATE1, &[0x0f, 0x84, 0x76, 0x01, 0x00, 0x00]);
        let mut lines = Lines::default();
        let mut patched = HashSet::new();
        assert!(cert_pass(&mem, &mut lines.sink(), &mut patched));
        assert!(lines.0.is_empty(), "{:?}", lines.0);
        assert!(patched.is_empty());
        assert_eq!(mem.at(GATE2, 3), [0x55, 0x48, 0x89]);
    }

    #[test]
    fn store_pass_enforces_every_site_in_table_order_then_advertises() {
        let mem = FakeMemory::new().with_store_sites(0xcc);
        let mut lines = Lines::default();
        let mut reported = HashSet::new();
        enforce_store_patches(&mem, CBASE, &mut lines.sink(), &mut reported).unwrap();

        let mut expected: Vec<String> = STORE_PATCHES
            .iter()
            .map(|(va, _)| {
                let live = live_addr(CBASE, *va);
                format!("pid 4242: ENFORCED store patch @ {live:#x}")
            })
            .collect();
        // 0xcc is neither the original nor the patch: fail-closed SKIP, and the
        // capability is withheld.
        expected.push(format!(
            "pid 4242: SKIP guarded patch @ {GUARD_LIVE:#x}: unexpected cccc (build mismatch)"
        ));
        expected.push(
            "[store-guard] guard status=UNSUPPORTED_BUILD fifa_pid=4242 (no capability advertised)"
                .to_string(),
        );
        assert_eq!(lines.0, expected);
        assert!(reported.contains(&4242));

        for (va, data) in STORE_PATCHES {
            assert_eq!(mem.at(live_addr(CBASE, va), data.len()), data);
        }
        // The guarded site was NOT overwritten.
        assert_eq!(mem.at(GUARD_LIVE, 2), [0xcc, 0xcc]);
    }

    #[test]
    fn store_pass_patches_the_guard_and_advertises_the_capability() {
        let mem = FakeMemory::new()
            .with_store_sites(0xcc)
            .map(GUARD_LIVE, STORE_PATCHES_GUARDED[0].1);
        let mut lines = Lines::default();
        let mut reported = HashSet::new();
        enforce_store_patches(&mem, CBASE, &mut lines.sink(), &mut reported).unwrap();

        assert_eq!(
            lines.0[lines.0.len() - 2],
            format!(
                "pid 4242: ENFORCED guarded store patch @ {GUARD_LIVE:#x} (JNZ->JG, empty My Packs)"
            )
        );
        assert_eq!(
            lines.0[lines.0.len() - 1],
            "[store-guard] verified capability fifa17.empty_mypacks_resolver=1 fifa_pid=4242"
        );
        assert_eq!(mem.at(GUARD_LIVE, 2), [0x7f, 0x0f]);

        // Second tick: everything already enforced, and the capability is not
        // re-advertised.
        let mut lines = Lines::default();
        enforce_store_patches(&mem, CBASE, &mut lines.sink(), &mut reported).unwrap();
        assert!(lines.0.is_empty(), "{:?}", lines.0);
    }

    #[test]
    fn store_pass_verifies_an_already_patched_guard() {
        let mem = FakeMemory::new()
            .with_store_sites(0xcc)
            .map(GUARD_LIVE, STORE_PATCHES_GUARDED[0].2);
        let mut lines = Lines::default();
        let mut reported = HashSet::new();
        enforce_store_patches(&mem, CBASE, &mut lines.sink(), &mut reported).unwrap();
        assert_eq!(
            lines.0.last().unwrap(),
            "[store-guard] verified capability fifa17.empty_mypacks_resolver=1 fifa_pid=4242"
        );
    }

    #[test]
    fn store_pass_reports_a_guarded_write_failure_without_aborting_the_tick() {
        let mut mem = FakeMemory::new()
            .with_store_sites(0xcc)
            .map(GUARD_LIVE, STORE_PATCHES_GUARDED[0].1);
        mem.fail_writes.push(GUARD_LIVE);
        let mut lines = Lines::default();
        let mut reported = HashSet::new();
        // The guarded write failure is handled inline, so the pass still succeeds.
        enforce_store_patches(&mem, CBASE, &mut lines.sink(), &mut reported).unwrap();
        assert!(
            lines.0[lines.0.len() - 2]
                .starts_with(&format!("pid 4242: guarded patch write failed @ {GUARD_LIVE:#x}: ")),
            "{}",
            lines.0[lines.0.len() - 2]
        );
        assert_eq!(
            lines.0[lines.0.len() - 1],
            "[store-guard] guard status=WRITE_FAILED fifa_pid=4242 (no capability advertised)"
        );
    }

    #[test]
    fn store_pass_withholds_the_capability_when_verification_fails() {
        let mut mem = FakeMemory::new()
            .with_store_sites(0xcc)
            .map(GUARD_LIVE, STORE_PATCHES_GUARDED[0].1);
        // Write reported OK, memory unchanged: the re-read still shows the original.
        mem.swallow_writes.push(GUARD_LIVE);
        let mut lines = Lines::default();
        let mut reported = HashSet::new();
        enforce_store_patches(&mem, CBASE, &mut lines.sink(), &mut reported).unwrap();
        assert_eq!(
            lines.0[lines.0.len() - 1],
            "[store-guard] guard status=VERIFY_FAILED fifa_pid=4242 (no capability advertised)"
        );
    }

    #[test]
    fn store_pass_aborts_the_tick_when_a_site_is_not_mapped() {
        // CardsDLL is mapped but this tick catches a site mid-unpack.
        let mem = FakeMemory::new();
        let mut lines = Lines::default();
        let mut reported = HashSet::new();
        let err = enforce_store_patches(&mem, CBASE, &mut lines.sink(), &mut reported).unwrap_err();
        assert_eq!(err.kind(), io::ErrorKind::NotFound);
        assert!(lines.0.is_empty());
        // Nothing advertised, so the next tick re-evaluates the guard.
        assert!(reported.is_empty());
    }

    #[test]
    fn store_pass_skips_writing_sites_that_already_hold_the_patch() {
        let mut mem = FakeMemory::new().with_store_sites(0xcc);
        for (va, data) in STORE_PATCHES {
            mem = mem.map(live_addr(CBASE, va), data);
        }
        mem = mem.map(GUARD_LIVE, STORE_PATCHES_GUARDED[0].2);
        // Any write at all would fail these sites, proving none is attempted.
        mem.fail_writes
            .extend(STORE_PATCHES.iter().map(|(va, _)| live_addr(CBASE, *va)));
        mem.fail_writes.push(GUARD_LIVE);

        let mut lines = Lines::default();
        let mut reported = HashSet::new();
        enforce_store_patches(&mem, CBASE, &mut lines.sink(), &mut reported).unwrap();
        assert_eq!(
            lines.0,
            vec!["[store-guard] verified capability fifa17.empty_mypacks_resolver=1 fifa_pid=4242"]
        );
    }
}