//! One-page, at-most-256-byte EL0 copy boundary for an owned UserSpace.
//! Arbitrary user VAs are never cast to kernel pointers. The caller supplies
//! physical frames owned and identity mapped for the lifetime of this view.

use crate::memory::PAGE_SIZE;
use crate::mmu::{USER_CODE_VA, USER_DATA_VA, USER_STACK_VA};

pub const MAX_COPY: usize = 256;
pub const EFAULT: i64 = -14;
pub const EINVAL: i64 = -22;

/// UART's text record accepts only printable ASCII. Validate the entire
/// kernel-owned copy before emitting even the prefix of a user message.
pub fn serial_text(bytes: &[u8]) -> Result<&str, i64> {
    if bytes.iter().any(|byte| !matches!(byte, 0x20..=0x7e)) {
        return Err(EINVAL);
    }
    core::str::from_utf8(bytes).map_err(|_| EINVAL)
}

#[derive(Clone, Copy)]
pub struct UserPages {
    pub code_frame: u64,
    pub data_frame: u64,
    pub stack_frame: u64,
}

impl UserPages {
    /// Resolve the *entire* range before any byte is read or written. A zero
    /// length operation succeeds without evaluating its untrusted address.
    pub fn resolve(&self, address: u64, length: usize, write: bool) -> Result<Option<u64>, i64> {
        if length == 0 {
            return Ok(None);
        }
        if length > MAX_COPY {
            return Err(EINVAL);
        }
        let last = address.checked_add(length as u64 - 1).ok_or(EFAULT)?;
        for (virtual_base, physical_base, writable) in [
            (USER_CODE_VA, self.code_frame, false),
            (USER_DATA_VA, self.data_frame, true),
            (USER_STACK_VA, self.stack_frame, true),
        ] {
            if address >= virtual_base && last < virtual_base + PAGE_SIZE {
                if write && !writable {
                    return Err(EFAULT);
                }
                return physical_base
                    .checked_add(address - virtual_base)
                    .map(Some)
                    .ok_or(EFAULT);
            }
        }
        Err(EFAULT)
    }

    /// # Safety
    /// `code_frame`, `data_frame`, and `stack_frame` must be distinct, live,
    /// identity-mapped 4 KiB frames exclusively owned by the current user
    /// space. The output slice must not alias any of those frames.
    pub unsafe fn copy_from(&self, address: u64, output: &mut [u8]) -> Result<usize, i64> {
        let Some(physical) = self.resolve(address, output.len(), false)? else {
            return Ok(0);
        };
        // SAFETY: resolve bounded the whole range to one live owned frame;
        // the caller guarantees the physical mapping and no output alias.
        unsafe {
            core::ptr::copy_nonoverlapping(physical as *const u8, output.as_mut_ptr(), output.len())
        };
        Ok(output.len())
    }

    /// # Safety
    /// Same live-frame guarantee as `copy_from`. `input` must not alias a
    /// destination frame. A rejected range writes no prefix of the input.
    pub unsafe fn copy_to(&self, address: u64, input: &[u8]) -> Result<usize, i64> {
        let Some(physical) = self.resolve(address, input.len(), true)? else {
            return Ok(0);
        };
        // SAFETY: permission and the entire range are checked before copying;
        // the caller guarantees live backing and no input alias.
        unsafe { core::ptr::copy_nonoverlapping(input.as_ptr(), physical as *mut u8, input.len()) };
        Ok(input.len())
    }
}

#[cfg(test)]
mod tests {
    extern crate std;
    use super::*;
    use sha2::{Digest, Sha256};

    #[test]
    fn zero_overflow_guard_and_read_only_ranges() {
        let pages = UserPages {
            code_frame: 0x1000,
            data_frame: 0x2000,
            stack_frame: 0x3000,
        };
        assert_eq!(pages.resolve(u64::MAX, 0, false), Ok(None));
        assert_eq!(pages.resolve(u64::MAX - 1, 4, false), Err(EFAULT));
        assert_eq!(
            pages.resolve(USER_DATA_VA + PAGE_SIZE - 2, 4, false),
            Err(EFAULT)
        );
        assert_eq!(pages.resolve(USER_CODE_VA, 4, true), Err(EFAULT));
        assert_eq!(
            pages.resolve(USER_DATA_VA + PAGE_SIZE, 4, false),
            Err(EFAULT)
        );
        assert_eq!(
            pages.resolve(USER_DATA_VA, MAX_COPY + 1, false),
            Err(EINVAL)
        );
        assert_eq!(
            pages.resolve(USER_STACK_VA + PAGE_SIZE - 4, 4, true),
            Ok(Some(0x3ffc))
        );
    }

    #[test]
    fn exact_copies_and_failure_atomicity() {
        let code = [0x11u8; 4096];
        let mut data = [0u8; 4096];
        let mut stack = [0u8; 4096];
        data[0..4].copy_from_slice(b"ABCD");
        data[4094] = 0xa5;
        data[4095] = 0x5a;
        let pages = UserPages {
            code_frame: code.as_ptr() as u64,
            data_frame: data.as_mut_ptr() as u64,
            stack_frame: stack.as_mut_ptr() as u64,
        };
        let mut output = [0xccu8; 4];
        assert_eq!(unsafe { pages.copy_from(USER_DATA_VA, &mut output) }, Ok(4));
        assert_eq!(&output, b"ABCD");
        assert_eq!(unsafe { pages.copy_to(USER_DATA_VA + 16, b"WXYZ") }, Ok(4));
        assert_eq!(&data[16..20], b"WXYZ");
        assert_eq!(unsafe { pages.copy_from(u64::MAX, &mut []) }, Ok(0));
        assert_eq!(
            unsafe { pages.copy_from(USER_DATA_VA + PAGE_SIZE - 2, &mut output) },
            Err(EFAULT)
        );
        assert_eq!(&output, b"ABCD");
        assert_eq!(
            unsafe { pages.copy_to(USER_DATA_VA + PAGE_SIZE - 2, b"FAIL") },
            Err(EFAULT)
        );
        assert_eq!(&data[4094..4096], &[0xa5, 0x5a]);
        assert_eq!(unsafe { pages.copy_to(USER_CODE_VA, b"FAIL") }, Err(EFAULT));
        assert_eq!(&code[0..4], &[0x11; 4]);
        assert_eq!(
            unsafe { pages.copy_from(USER_DATA_VA + PAGE_SIZE, &mut output) },
            Err(EFAULT)
        );
    }

    #[test]
    fn serial_text_rejects_control_and_non_ascii_before_output() {
        assert_eq!(serial_text(b"ABCD"), Ok("ABCD"));
        assert_eq!(serial_text(b""), Ok(""));
        assert_eq!(serial_text(b"A\n"), Err(EINVAL));
        assert_eq!(serial_text(&[b'A', 0x7f]), Err(EINVAL));
        assert_eq!(serial_text(&[b'A', 0xc3, 0xa9]), Err(EINVAL));
    }

    #[test]
    fn k444_deterministic_syscall_length_fuzz_rejects_wrap_and_partial_copy() {
        let pages = UserPages {
            code_frame: 0x10000,
            data_frame: 0x20000,
            stack_frame: 0x30000,
        };
        let oracle = |address: u64, length: usize, write: bool| {
            if length == 0 {
                return Ok(None);
            }
            if length > MAX_COPY {
                return Err(EINVAL);
            }
            // Independent half-open u128 interval arithmetic avoids the
            // production helper's checked_add(last-byte) formulation.
            let end = u128::from(address) + length as u128;
            for (base, frame, writable) in [
                (USER_CODE_VA, 0x10000, false),
                (USER_DATA_VA, 0x20000, true),
                (USER_STACK_VA, 0x30000, true),
            ] {
                if address >= base && end <= u128::from(base + PAGE_SIZE) {
                    return if write && !writable {
                        Err(EFAULT)
                    } else {
                        Ok(Some(frame + address - base))
                    };
                }
            }
            Err(EFAULT)
        };
        let anchors = [
            0,
            USER_CODE_VA - 1,
            USER_CODE_VA,
            USER_CODE_VA + PAGE_SIZE - 1,
            USER_DATA_VA,
            USER_DATA_VA + PAGE_SIZE - 1,
            USER_STACK_VA,
            USER_STACK_VA + PAGE_SIZE - 1,
            u64::MAX - 256,
            u64::MAX,
        ];
        let lengths = [0usize, 1, 2, 4, 255, 256, 257, usize::MAX];
        let deltas = [-257i64, -256, -2, -1, 0, 1, 2, 255, 256, 257];
        let mut digest = Sha256::new();
        let mut cases = 0usize;
        let mut accepted = 0usize;
        let mut rejected_fault = 0usize;
        let mut rejected_length = 0usize;
        let mut empty = 0usize;
        let mut check = |address: u64, length: usize, write: bool| {
            let expected = oracle(address, length, write);
            let actual = pages.resolve(address, length, write);
            assert_eq!(
                actual, expected,
                "address={address:#x} length={length} write={write}"
            );
            digest.update(address.to_le_bytes());
            digest.update((length as u64).to_le_bytes());
            digest.update([u8::from(write)]);
            digest.update([match actual {
                Ok(Some(_)) => 1,
                Ok(None) => 2,
                Err(EFAULT) => 3,
                Err(EINVAL) => 4,
                _ => unreachable!(),
            }]);
            cases += 1;
            match actual {
                Ok(Some(_)) => accepted += 1,
                Ok(None) => empty += 1,
                Err(EFAULT) => rejected_fault += 1,
                Err(EINVAL) => rejected_length += 1,
                _ => unreachable!(),
            }
        };
        for anchor in anchors {
            for delta in deltas {
                let address = if delta < 0 {
                    anchor.checked_sub((-delta) as u64)
                } else {
                    anchor.checked_add(delta as u64)
                };
                if let Some(address) = address {
                    for length in lengths {
                        for write in [false, true] {
                            check(address, length, write);
                        }
                    }
                }
            }
        }
        let mut state = 0x4430_2026_d15c_0a11u64;
        for index in 0..32_768 {
            state ^= state << 13;
            state ^= state >> 7;
            state ^= state << 17;
            let anchor = anchors[state as usize % anchors.len()];
            let delta = ((state >> 16) % 1025) as i64 - 512;
            let address = if index % 4 == 0 {
                state
            } else if delta < 0 {
                anchor.saturating_sub((-delta) as u64)
            } else {
                anchor.saturating_add(delta as u64)
            };
            let length = if index % 4096 == 0 {
                usize::MAX
            } else {
                ((state >> 32) % 514) as usize
            };
            check(address, length, index & 1 == 1);
        }
        assert!(cases > 34_000);
        assert!(accepted > 0 && rejected_fault > 0 && rejected_length > 0 && empty > 0);

        let mut code = [0x11u8; PAGE_SIZE as usize];
        let mut data = [0xa5u8; PAGE_SIZE as usize];
        let mut stack = [0x5au8; PAGE_SIZE as usize];
        let live = UserPages {
            code_frame: code.as_mut_ptr() as u64,
            data_frame: data.as_mut_ptr() as u64,
            stack_frame: stack.as_mut_ptr() as u64,
        };
        let input = [0x77u8; MAX_COPY + 1];
        let mut output = [0xccu8; MAX_COPY + 1];
        for (address, length, error) in [
            (u64::MAX - 1, 4, EFAULT),
            (USER_DATA_VA + PAGE_SIZE - 1, 2, EFAULT),
            (USER_STACK_VA + PAGE_SIZE - 1, 2, EFAULT),
            (USER_CODE_VA, 4, EFAULT),
            (USER_DATA_VA, MAX_COPY + 1, EINVAL),
        ] {
            assert_eq!(
                unsafe { live.copy_to(address, &input[..length]) },
                Err(error)
            );
            assert_eq!(&code[..], &[0x11; PAGE_SIZE as usize]);
            assert_eq!(&data[..], &[0xa5; PAGE_SIZE as usize]);
            assert_eq!(&stack[..], &[0x5a; PAGE_SIZE as usize]);
            if address != USER_CODE_VA {
                assert_eq!(
                    unsafe { live.copy_from(address, &mut output[..length]) },
                    Err(error)
                );
                assert_eq!(&output[..], &[0xcc; MAX_COPY + 1]);
            }
        }
        assert_eq!(
            unsafe { live.copy_to(USER_DATA_VA + PAGE_SIZE - 256, &input[..256]) },
            Ok(256)
        );
        assert_eq!(&data[PAGE_SIZE as usize - 256..], &input[..256]);
        assert_eq!(&stack[..], &[0x5a; PAGE_SIZE as usize]);
        let hash = digest
            .finalize()
            .iter()
            .map(|byte| std::format!("{byte:02x}"))
            .collect::<std::string::String>();
        std::eprintln!(
            "K444 CASES={cases} OK={accepted} EMPTY={empty} EFAULT={rejected_fault} EINVAL={rejected_length} SHA256={hash}"
        );
    }
}