use std::collections::BTreeMap;
use turkos_os_mapping::mapping::*;
use turkos_os_page_tables::descriptor::Permissions;

const USER: u64 = 0x8000_0000;
const OWNER: Owner = Owner {
    id: 1,
    generation: 7,
};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct Ref {
    slot: u8,
    generation: u64,
}
struct Memory {
    root: BTreeMap<u64, Value<Ref>>,
    data: Vec<Vec<u8>>,
    prepares: usize,
    commits: usize,
    fail_prepare: bool,
}
impl Memory {
    fn new() -> Self {
        Self {
            root: BTreeMap::new(),
            data: (0..8).map(|slot| vec![0x31 + slot; 4096]).collect(),
            prepares: 0,
            commits: 0,
            fail_prepare: false,
        }
    }
}
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<(), Error> {
        if owner != OWNER || usize::from(frame.slot) >= self.data.len() || frame.generation != 1 {
            Err(Error::Backing)
        } else {
            Ok(())
        }
    }
    fn current(&self, address: u64) -> Result<Option<Value<Ref>>, Error> {
        Ok(self.root.get(&address).copied())
    }
    fn prepare(&mut self, changes: &[Change<Ref>]) -> Result<Self::Prepared, Error> {
        self.prepares += 1;
        let mut draft = self.root.clone();
        for (index, change) in changes.iter().enumerate() {
            assert_eq!(draft.get(&change.virtual_address).copied(), change.before);
            if let Some(after) = change.after {
                draft.insert(change.virtual_address, after);
            } else {
                draft.remove(&change.virtual_address);
            }
            if self.fail_prepare && index == changes.len() - 1 {
                return Err(Error::NoMemory);
            }
        }
        Ok(draft)
    }
    fn commit(&mut self, draft: Self::Prepared) {
        self.commits += 1;
        self.root = draft;
    }
}
fn rw() -> Permissions {
    Permissions {
        user: true,
        write: true,
        execute: false,
    }
}
fn ro() -> Permissions {
    Permissions {
        user: true,
        write: false,
        execute: false,
    }
}
fn rx() -> Permissions {
    Permissions {
        user: true,
        write: false,
        execute: true,
    }
}
fn refs() -> [Ref; 3] {
    [
        Ref {
            slot: 0,
            generation: 1,
        },
        Ref {
            slot: 1,
            generation: 1,
        },
        Ref {
            slot: 2,
            generation: 1,
        },
    ]
}
fn mapper() -> Mapping<Memory, 16> {
    Mapping::new(
        Memory::new(),
        841,
        Window::new(USER, USER + 16 * PAGE, 39).unwrap(),
    )
    .unwrap()
}
fn snapshot(mapper: &Mapping<Memory, 16>) -> (BTreeMap<u64, Value<Ref>>, usize, usize) {
    mapper.with_backend_observation(|memory| (memory.root.clone(), memory.prepares, memory.commits))
}

#[test]
fn os_mapping_insert_positive() {
    let mut map = mapper();
    map.map(OWNER, USER, &refs(), rw()).unwrap();
    assert_eq!(snapshot(&map).1, 1);
    assert_eq!(snapshot(&map).2, 1);
    for (index, reference) in refs().iter().enumerate() {
        assert_eq!(
            map.lookup(OWNER, USER + index as u64 * PAGE),
            Ok(Some(Value {
                frame: *reference,
                permissions: rw()
            }))
        );
    }
    let actual = map.with_backend_observation(|memory| {
        memory
            .root
            .values()
            .map(|value| memory.data[usize::from(value.frame.slot)][4095])
            .collect::<Vec<_>>()
    });
    assert_eq!(actual, vec![0x31, 0x32, 0x33]);
}
#[test]
fn os_mapping_insert_negative() {
    let mut map = mapper();
    map.map(OWNER, USER + 2 * PAGE, &refs()[2..], rw()).unwrap();
    let before = snapshot(&map);
    assert_eq!(map.map(OWNER, USER, &refs(), rw()), Err(Error::Conflict));
    assert_eq!(snapshot(&map), before);
    assert_eq!(map.lookup(OWNER, USER), Err(Error::Unmapped));
    assert_eq!(map.lookup(OWNER, USER + PAGE), Err(Error::Unmapped));
}
#[test]
fn preparation_failure_commits_no_root_prefix() {
    let mut memory = Memory::new();
    memory.fail_prepare = true;
    let mut map = Mapping::<_, 16>::new(
        memory,
        841,
        Window::new(USER, USER + 16 * PAGE, 39).unwrap(),
    )
    .unwrap();
    assert_eq!(map.map(OWNER, USER, &refs(), rw()), Err(Error::NoMemory));
    assert!(snapshot(&map).0.is_empty());
    assert_eq!(snapshot(&map).2, 0);
    assert_eq!(map.records().count(), 0);
}
#[test]
fn os_mapping_remove_positive() {
    let mut map = mapper();
    map.map(OWNER, USER, &refs(), rw()).unwrap();
    map.unmap(OWNER, USER + PAGE, 1).unwrap();
    assert_eq!(map.lookup(OWNER, USER + PAGE), Err(Error::Unmapped));
    for offset in [0, 2] {
        assert!(map.lookup(OWNER, USER + offset * PAGE).unwrap().is_some());
    }
    assert_eq!(snapshot(&map).0.len(), 2);
}
#[test]
fn os_mapping_remove_negative() {
    let mut map = mapper();
    map.map(OWNER, USER, &refs(), rw()).unwrap();
    let before = snapshot(&map);
    assert_eq!(
        map.unmap(Owner { id: 2, ..OWNER }, USER, 3),
        Err(Error::Foreign)
    );
    assert_eq!(snapshot(&map), before);
}
#[test]
fn os_mapping_protect_positive() {
    let mut map = mapper();
    map.map(OWNER, USER, &refs(), rw()).unwrap();
    map.protect(OWNER, USER, 3, ro()).unwrap();
    map.protect(OWNER, USER, 3, rx()).unwrap();
    for offset in 0..3 {
        assert_eq!(
            map.lookup(OWNER, USER + offset * PAGE)
                .unwrap()
                .unwrap()
                .permissions,
            rx()
        );
    }
    assert_eq!(snapshot(&map).2, 3);
}
#[test]
fn os_mapping_protect_negative() {
    let mut map = mapper();
    map.map(OWNER, USER, &refs(), rw()).unwrap();
    let before = snapshot(&map);
    assert_eq!(
        map.protect(
            OWNER,
            USER,
            3,
            Permissions {
                user: true,
                write: true,
                execute: true
            }
        ),
        Err(Error::Permissions)
    );
    assert_eq!(snapshot(&map), before);
}
#[test]
fn aliases_cannot_combine_writable_and_executable_user_backing() {
    let mut map = mapper();
    map.map(OWNER, USER, &refs()[..1], rw()).unwrap();
    let before = snapshot(&map);
    assert_eq!(
        map.map(OWNER, USER + PAGE, &refs()[..1], rx()),
        Err(Error::Permissions)
    );
    assert_eq!(snapshot(&map), before);
    map.map(OWNER, USER + PAGE, &refs()[..1], ro()).unwrap();
    let before = snapshot(&map);
    assert_eq!(
        map.protect(OWNER, USER + PAGE, 1, rx()),
        Err(Error::Permissions)
    );
    assert_eq!(snapshot(&map), before);
    map.protect(OWNER, USER, 2, ro()).unwrap();
    map.protect(OWNER, USER, 2, rx()).unwrap();
}
#[test]
fn os_mapping_guard_positive() {
    let mut map = mapper();
    map.guard(OWNER, USER, 1).unwrap();
    assert_eq!(map.lookup(OWNER, USER), Ok(None));
    assert_eq!(snapshot(&map).1, 0);
    assert_eq!(snapshot(&map).2, 0);
    assert_eq!(
        map.map(OWNER, USER, &refs()[..1], rw()),
        Err(Error::Conflict)
    );
    map.unmap(OWNER, USER, 1).unwrap();
    assert_eq!(snapshot(&map).1, 0);
    assert_eq!(map.records().count(), 0);
}
#[test]
fn os_mapping_guard_negative() {
    let mut map = mapper();
    map.guard(OWNER, USER, 1).unwrap();
    let before = snapshot(&map);
    assert_eq!(map.protect(OWNER, USER, 1, rx()), Err(Error::Guard));
    assert_eq!(map.pin(OWNER, USER, 1, false).err(), Some(Error::Guard));
    assert_eq!(snapshot(&map), before);
}
#[test]
fn copy_pin_blocks_unmap_protect_and_backend_escape() {
    let mut map = mapper();
    map.map(OWNER, USER, &refs(), rw()).unwrap();
    let mut pin = map.pin(OWNER, USER, 3, false).unwrap();
    let before = snapshot(&map);
    assert_eq!(map.unmap(OWNER, USER + PAGE, 1), Err(Error::Busy));
    assert_eq!(map.protect(OWNER, USER, 3, ro()), Err(Error::Busy));
    assert_eq!(snapshot(&map), before);
    let mut map = match map.into_backend() {
        Err((Error::Busy, map)) => map,
        _ => panic!("pinned backing escaped"),
    };
    map.unpin(&mut pin).unwrap();
    assert_eq!(map.unpin(&mut pin), Err(Error::StalePin));
    map.unmap(OWNER, USER, 3).unwrap();
    assert!(map.into_backend().is_ok());
}
#[test]
fn pin_preflight_last_page_fault_changes_no_prefix_pins() {
    let mut map = mapper();
    map.map(OWNER, USER, &refs()[..2], rw()).unwrap();
    assert_eq!(map.pin(OWNER, USER, 3, false).err(), Some(Error::Unmapped));
    assert!(map.records().all(|row| row.3 == 0));
    map.unmap(OWNER, USER, 2).unwrap();
}
#[test]
fn same_caller_scope_ids_cannot_unpin_an_independent_mapping() {
    let mut a = mapper();
    let mut b = mapper();
    a.map(OWNER, USER, &refs(), rw()).unwrap();
    b.map(OWNER, USER, &refs(), rw()).unwrap();
    let mut a_pin = a.pin(OWNER, USER, 3, false).unwrap();
    let mut b_pin = b.pin(OWNER, USER, 3, false).unwrap();
    assert_eq!(b.unpin(&mut a_pin), Err(Error::StalePin));
    assert!(b.records().all(|row| row.3 == 1));
    a.unpin(&mut a_pin).unwrap();
    b.unpin(&mut b_pin).unwrap();
}
#[test]
fn readonly_or_stale_frame_pin_never_grants_write() {
    let mut map = mapper();
    map.map(OWNER, USER, &refs(), ro()).unwrap();
    assert_eq!(
        map.pin(OWNER, USER, 3, true).err(),
        Some(Error::Permissions)
    );
    assert!(map.records().all(|row| row.3 == 0));
    let before = snapshot(&map);
    assert_eq!(
        map.map(
            OWNER,
            USER + 4 * PAGE,
            &[Ref {
                slot: 4,
                generation: 2
            }],
            rw()
        ),
        Err(Error::Backing)
    );
    assert_eq!(snapshot(&map), before);
}
#[test]
fn overflow_alignment_window_and_metadata_limits_fail_before_prepare() {
    let mut map = mapper();
    assert_eq!(map.guard(OWNER, u64::MAX - 4095, 2), Err(Error::Overflow));
    assert_eq!(map.guard(OWNER, USER + 1, 1), Err(Error::Invalid));
    assert_eq!(map.guard(OWNER, 0x09000000, 1), Err(Error::OutsideWindow));
    assert_eq!(map.guard(OWNER, USER, 0), Err(Error::Invalid));
    map.guard(OWNER, USER, 16).unwrap();
    assert_eq!(map.map(OWNER, USER, &refs(), rw()), Err(Error::Capacity));
    assert_eq!(snapshot(&map).1, 0);
}