use std::{
    cell::{Cell, RefCell},
    collections::BTreeMap,
};
use turkos_os_mapping::mapping::{
    Backend, Change, Error as MappingError, Mapping, Owner, PAGE, Value, Window,
};
use turkos_os_page_tables::descriptor::Permissions;
use turkos_os_user_copy::user_copy::*;

const USER: u64 = 0x8000_0000;
const OWNER: Owner = Owner {
    id: 7,
    generation: 3,
};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct Ref {
    slot: usize,
    generation: u64,
}
struct Memory {
    root: BTreeMap<u64, Value<Ref>>,
    data: Vec<RefCell<Vec<u8>>>,
    generation: Cell<u64>,
    reads: RefCell<Vec<usize>>,
    writes: RefCell<Vec<usize>>,
}
impl Memory {
    fn new() -> Self {
        Self {
            root: BTreeMap::new(),
            data: (0..20)
                .map(|slot| {
                    RefCell::new(
                        (0..4096)
                            .map(|offset| ((slot * 17 + offset) % 251 + 1) as u8)
                            .collect(),
                    )
                })
                .collect(),
            generation: Cell::new(1),
            reads: RefCell::new(vec![0; 20]),
            writes: RefCell::new(vec![0; 20]),
        }
    }
}
impl Backend for Memory {
    type FrameRef = Ref;
    type Prepared = BTreeMap<u64, Value<Ref>>;
    fn owner(&self) -> Owner {
        OWNER
    }
    fn validate_frame(&self, owner: Owner, frame: Ref) -> Result<(), MappingError> {
        if owner == OWNER
            && frame.slot < self.data.len()
            && frame.generation == self.generation.get()
        {
            Ok(())
        } else {
            Err(MappingError::Backing)
        }
    }
    fn current(&self, address: u64) -> Result<Option<Value<Ref>>, MappingError> {
        Ok(self.root.get(&address).copied())
    }
    fn prepare(&mut self, changes: &[Change<Ref>]) -> Result<Self::Prepared, MappingError> {
        let mut draft = self.root.clone();
        for change in changes {
            assert_eq!(draft.get(&change.virtual_address).copied(), change.before);
            if let Some(value) = change.after {
                draft.insert(change.virtual_address, value);
            } else {
                draft.remove(&change.virtual_address);
            }
        }
        Ok(draft)
    }
    fn commit(&mut self, draft: Self::Prepared) {
        self.root = draft;
    }
}
// SAFETY: actual owned host buffers are retained by Memory, bounded accesses
// cannot fail, and no backend callback modifies frame identity or root state.
unsafe impl PinnedMemory for Memory {
    unsafe fn read_ram(&self, frame: Ref, offset: usize, destination: &mut [u8]) {
        assert_eq!(frame.generation, self.generation.get());
        destination
            .copy_from_slice(&self.data[frame.slot].borrow()[offset..offset + destination.len()]);
        self.reads.borrow_mut()[frame.slot] += destination.len();
    }
    unsafe fn write_ram(&self, frame: Ref, offset: usize, source: &[u8]) {
        assert_eq!(frame.generation, self.generation.get());
        self.data[frame.slot].borrow_mut()[offset..offset + source.len()].copy_from_slice(source);
        self.writes.borrow_mut()[frame.slot] += source.len();
    }
}
fn rw() -> Permissions {
    Permissions {
        user: true,
        write: true,
        execute: false,
    }
}
fn rx() -> Permissions {
    Permissions {
        user: true,
        write: false,
        execute: true,
    }
}
fn window() -> UserWindow {
    UserWindow::new(USER, USER + 20 * PAGE, 39).unwrap()
}
fn mapper(pages: usize) -> Mapping<Memory, 20> {
    let mut mapping = Mapping::new(
        Memory::new(),
        861,
        Window::new(USER, USER + 20 * PAGE, 39).unwrap(),
    )
    .unwrap();
    let frames: Vec<_> = (0..pages)
        .map(|slot| Ref {
            slot,
            generation: 1,
        })
        .collect();
    if pages != 0 {
        mapping.map(OWNER, USER, &frames, rw()).unwrap();
    }
    mapping
}
fn buffers(mapping: &Mapping<Memory, 20>) -> Vec<Vec<u8>> {
    mapping.with_backend_observation(|memory| {
        memory
            .data
            .iter()
            .map(|data| data.borrow().clone())
            .collect()
    })
}
fn accesses(mapping: &Mapping<Memory, 20>) -> (usize, usize) {
    mapping.with_backend_observation(|memory| {
        (
            memory.reads.borrow().iter().sum(),
            memory.writes.borrow().iter().sum(),
        )
    })
}

#[test]
fn os_user_copy_read_positive() {
    let mut mapping = mapper(3);
    let expected: Vec<_> = buffers(&mapping)[..3]
        .iter()
        .flatten()
        .copied()
        .skip(4095)
        .take(8193)
        .collect();
    let mut destination = vec![0xcc; 8193];
    copy_from_user(&mut mapping, OWNER, window(), USER + 4095, &mut destination).unwrap();
    assert_eq!(destination, expected);
    assert_eq!(accesses(&mapping), (8193, 0));
    assert!(mapping.records().all(|record| record.3 == 0));
}
#[test]
fn os_user_copy_read_negative() {
    let mut mapping = mapper(2);
    let mut destination = vec![0xcc; 8193];
    assert_eq!(
        copy_from_user(&mut mapping, OWNER, window(), USER + 4095, &mut destination),
        Err(Error::Mapping(MappingError::Unmapped))
    );
    assert!(destination.iter().all(|byte| *byte == 0xcc));
    assert_eq!(accesses(&mapping), (0, 0));
    assert!(mapping.records().all(|record| record.3 == 0));
}
#[test]
fn os_user_copy_write_positive() {
    let mut mapping = mapper(2);
    let before = buffers(&mapping);
    let source = vec![0x5a; 7000];
    copy_to_user(&mut mapping, OWNER, window(), USER + 333, &source).unwrap();
    let after = buffers(&mapping);
    let flattened: Vec<_> = after[..2].iter().flatten().copied().collect();
    assert_eq!(&flattened[333..7333], &source);
    assert_eq!(&after[0][..333], &before[0][..333]);
    assert_eq!(&after[1][3237..], &before[1][3237..]);
    assert_eq!(accesses(&mapping), (0, 7000));
    assert!(mapping.records().all(|record| record.3 == 0));
}
#[test]
fn os_user_copy_write_negative() {
    let mut mapping = mapper(2);
    mapping.protect(OWNER, USER + PAGE, 1, rx()).unwrap();
    let before = buffers(&mapping);
    let source = vec![0x5a; 7000];
    assert_eq!(
        copy_to_user(&mut mapping, OWNER, window(), USER + 333, &source),
        Err(Error::Mapping(MappingError::Permissions))
    );
    assert_eq!(buffers(&mapping), before);
    assert_eq!(accesses(&mapping), (0, 0));
    assert!(mapping.records().all(|record| record.3 == 0));
}
#[test]
fn os_user_copy_strings_positive() {
    let mut mapping = mapper(2);
    copy_to_user(&mut mapping, OWNER, window(), USER + 4093, b"turkos\0").unwrap();
    let mut destination = [0xcc; 16];
    let mut workspace = StringWorkspace::<32>::new();
    assert_eq!(
        copy_cstr(
            &mut mapping,
            OWNER,
            window(),
            USER + 4093,
            32,
            &mut destination,
            &mut workspace
        ),
        Ok(6)
    );
    assert_eq!(&destination[..6], b"turkos");
    assert!(destination[6..].iter().all(|byte| *byte == 0xcc));
    assert!(mapping.records().all(|record| record.3 == 0));
}
#[test]
fn os_user_copy_strings_negative() {
    let mut mapping = mapper(17);
    mapping.with_backend_observation(|memory| {
        for page in &memory.data[..16] {
            page.borrow_mut().fill(0x7e);
        }
        memory.data[16].borrow_mut()[0] = 0;
    });
    let mut destination = [0xcc; 32];
    let mut workspace = StringWorkspace::<65536>::new();
    assert_eq!(
        copy_cstr(
            &mut mapping,
            OWNER,
            window(),
            USER,
            65536,
            &mut destination,
            &mut workspace
        ),
        Err(Error::NameTooLong)
    );
    assert_eq!(destination, [0xcc; 32]);
    assert_eq!(accesses(&mapping), (65536, 0));
    mapping.with_backend_observation(|memory| assert_eq!(memory.reads.borrow()[16], 0));
    assert!(mapping.records().all(|record| record.3 == 0));
}
#[test]
fn os_user_copy_race_positive() {
    let mut mapping = mapper(3);
    let expected = buffers(&mapping)[0].clone();
    let mut destination = vec![0; 4096];
    {
        let mut range =
            PinnedRange::acquire(&mut mapping, OWNER, window(), USER, 4096, false).unwrap();
        assert_eq!(range.attempt_unmap(OWNER, USER, 1), Err(MappingError::Busy));
        assert_eq!(
            range.attempt_protect(OWNER, USER, 1, rx()),
            Err(MappingError::Busy)
        );
        range.read_into(&mut destination).unwrap();
    }
    assert_eq!(destination, expected);
    mapping.unmap(OWNER, USER, 1).unwrap();
}
#[test]
fn os_user_copy_race_negative() {
    let mut mapping = mapper(2);
    let before = buffers(&mapping);
    {
        let mut range =
            PinnedRange::acquire(&mut mapping, OWNER, window(), USER, 8192, true).unwrap();
        assert_eq!(range.attempt_unmap(OWNER, USER, 2), Err(MappingError::Busy));
        range.write_from(&vec![0x63; 8192]).unwrap();
    }
    let after = buffers(&mapping);
    assert_ne!(after[..2], before[..2]);
    assert!(after[..2].iter().flatten().all(|byte| *byte == 0x63));
    assert!(mapping.records().all(|record| record.3 == 0));
    mapping.unmap(OWNER, USER, 2).unwrap();
}
#[test]
fn foreign_owner_cannot_read_or_write_any_byte() {
    let mut mapping = mapper(2);
    let before = buffers(&mapping);
    let mut destination = [0xcc; 8];
    let foreign = Owner { id: 8, ..OWNER };
    assert_eq!(
        copy_from_user(&mut mapping, foreign, window(), USER, &mut destination),
        Err(Error::Mapping(MappingError::Foreign))
    );
    assert_eq!(
        copy_to_user(&mut mapping, foreign, window(), USER, &[0x5a; 8]),
        Err(Error::Mapping(MappingError::Foreign))
    );
    assert_eq!(destination, [0xcc; 8]);
    assert_eq!(buffers(&mapping), before);
    assert_eq!(accesses(&mapping), (0, 0));
}
#[test]
fn kernel_mmio_and_high_half_addresses_never_reach_ram_adapter() {
    let mut mapping = mapper(1);
    for address in [
        0x0900_0000,
        0x4000_0000,
        (1u64 << 39) - 8,
        0xffff_ffff_ffff_0000,
    ] {
        let mut destination = [0xcc; 8];
        assert_eq!(
            copy_from_user(&mut mapping, OWNER, window(), address, &mut destination),
            Err(Error::OutsideWindow)
        );
        assert_eq!(destination, [0xcc; 8]);
    }
    assert_eq!(accesses(&mapping), (0, 0));
}
#[test]
fn pointer_overflow_is_rejected_before_any_pin_or_copy() {
    let mut mapping = mapper(1);
    let mut destination = [0xcc; 8];
    assert_eq!(
        copy_from_user(
            &mut mapping,
            OWNER,
            window(),
            u64::MAX - 3,
            &mut destination
        ),
        Err(Error::Overflow)
    );
    assert_eq!(destination, [0xcc; 8]);
    assert_eq!(accesses(&mapping), (0, 0));
}
#[test]
fn overlimit_copy_and_small_string_workspace_are_preflight_errors() {
    let mut mapping = mapper(17);
    let mut destination = vec![0xcc; MAX_COPY + 1];
    assert_eq!(
        copy_from_user(&mut mapping, OWNER, window(), USER, &mut destination),
        Err(Error::TooLong)
    );
    let mut workspace = StringWorkspace::<16>::new();
    assert_eq!(
        copy_cstr(
            &mut mapping,
            OWNER,
            window(),
            USER,
            17,
            &mut destination,
            &mut workspace
        ),
        Err(Error::Capacity)
    );
    assert!(destination.iter().all(|byte| *byte == 0xcc));
    assert_eq!(accesses(&mapping), (0, 0));
}
#[test]
fn zero_length_copy_at_window_boundary_performs_no_access() {
    let mut mapping = mapper(0);
    copy_from_user(&mut mapping, OWNER, window(), USER + 20 * PAGE, &mut []).unwrap();
    copy_to_user(&mut mapping, OWNER, window(), USER, &[]).unwrap();
    assert_eq!(accesses(&mapping), (0, 0));
    assert_eq!(mapping.records().count(), 0);
}
#[test]
fn wrong_length_and_read_pin_write_leave_both_destinations_unchanged() {
    let mut mapping = mapper(1);
    let before = buffers(&mapping);
    {
        let range = PinnedRange::acquire(&mut mapping, OWNER, window(), USER, 8, false).unwrap();
        let mut destination = [0xcc; 7];
        assert_eq!(range.read_into(&mut destination), Err(Error::WrongLength));
        assert_eq!(destination, [0xcc; 7]);
        assert_eq!(range.write_from(&[0x5a; 8]), Err(Error::ReadOnly));
    }
    assert_eq!(buffers(&mapping), before);
    assert_eq!(accesses(&mapping), (0, 0));
    mapping.unmap(OWNER, USER, 1).unwrap();
}
#[test]
fn stale_generation_never_reaches_the_ram_alias() {
    let mut mapping = mapper(1);
    mapping.with_backend_observation(|memory| memory.generation.set(2));
    let mut destination = [0xcc; 8];
    assert_eq!(
        copy_from_user(&mut mapping, OWNER, window(), USER, &mut destination),
        Err(Error::Mapping(MappingError::Backing))
    );
    assert_eq!(destination, [0xcc; 8]);
    assert_eq!(accesses(&mapping), (0, 0));
    assert!(mapping.records().all(|record| record.3 == 0));
}
#[test]
fn unterminated_string_with_missing_later_page_preserves_kernel_output() {
    let mut mapping = mapper(1);
    let mut destination = [0xcc; 32];
    let mut workspace = StringWorkspace::<32>::new();
    assert_eq!(
        copy_cstr(
            &mut mapping,
            OWNER,
            window(),
            USER + 4090,
            32,
            &mut destination,
            &mut workspace
        ),
        Err(Error::Mapping(MappingError::Unmapped))
    );
    assert_eq!(destination, [0xcc; 32]);
    assert_eq!(accesses(&mapping), (6, 0));
    assert!(mapping.records().all(|record| record.3 == 0));
}
#[test]
fn terminated_string_does_not_require_unmapped_later_pages() {
    let mut mapping = mapper(1);
    copy_to_user(&mut mapping, OWNER, window(), USER + 4093, b"x\0").unwrap();
    let mut destination = [0xcc; 8];
    let mut workspace = StringWorkspace::<32>::new();
    assert_eq!(
        copy_cstr(
            &mut mapping,
            OWNER,
            window(),
            USER + 4093,
            32,
            &mut destination,
            &mut workspace
        ),
        Ok(1)
    );
    assert_eq!(destination[0], b'x');
    assert_eq!(&destination[1..], &[0xcc; 7]);
}
#[test]
fn insufficient_string_output_and_zero_limit_publish_no_prefix() {
    let mut mapping = mapper(1);
    copy_to_user(&mut mapping, OWNER, window(), USER, b"turkos\0").unwrap();
    let mut destination = [0xcc; 5];
    let mut workspace = StringWorkspace::<16>::new();
    assert_eq!(
        copy_cstr(
            &mut mapping,
            OWNER,
            window(),
            USER,
            16,
            &mut destination,
            &mut workspace
        ),
        Err(Error::Capacity)
    );
    assert_eq!(destination, [0xcc; 5]);
    let before = accesses(&mapping);
    assert_eq!(
        copy_cstr(
            &mut mapping,
            OWNER,
            window(),
            USER,
            0,
            &mut destination,
            &mut workspace
        ),
        Err(Error::NameTooLong)
    );
    assert_eq!(destination, [0xcc; 5]);
    assert_eq!(accesses(&mapping), before);
}
#[test]
fn guard_page_copy_never_acquires_or_reads_backing() {
    let mut mapping = mapper(0);
    mapping.guard(OWNER, USER, 1).unwrap();
    let mut destination = [0xcc; 8];
    assert_eq!(
        copy_from_user(&mut mapping, OWNER, window(), USER, &mut destination),
        Err(Error::Mapping(MappingError::Guard))
    );
    assert_eq!(destination, [0xcc; 8]);
    assert_eq!(accesses(&mapping), (0, 0));
}
#[test]
fn invalid_user_windows_are_rejected_before_a_copy_exists() {
    for (start, end, bits) in [
        (USER, USER, 39),
        (USER + 1, USER + PAGE, 39),
        (USER, USER + PAGE + 1, 39),
        (USER, USER + PAGE, 42),
        (USER, 1 << 40, 39),
    ] {
        assert_eq!(UserWindow::new(start, end, bits), Err(Error::InvalidWindow));
    }
}