use core::{fmt, num::NonZero, slice};
use nix::sys::mman::{mmap_anonymous, munmap, MapFlags, ProtFlags};
use proptest::prelude::*;
use std::ptr::NonNull;
use vernos_buddy_allocator::BuddyAllocator;
use vernos_physmem_free_list::FreeListAllocator;

proptest! {
    #![proptest_config(proptest::test_runner::Config {
        failure_persistence: None,
        ..Default::default()
    })]

    #[test]
    fn test_scenario(scenario in Scenario::any()) {
        scenario.run(true)
    }
}

#[test]
fn test_simple_scenario() {
    let scenario = Scenario {
        range_sizes: vec![7],
        actions: vec![
            Action::Debug,
            Action::Alloc {
                sentinel_value: 0xaa,
                size_class: 0,
            },
            Action::Debug,
        ],
    };
    scenario.run(false)
}

const PAGE_SIZE: usize = 64;
const PAGE_SIZE_BITS: usize = PAGE_SIZE.trailing_zeros() as usize;
const SIZE_CLASS_COUNT: usize = 4;

/// A single action the property test should perform.
#[derive(Debug)]
enum Action {
    /// Makes an allocation.
    Alloc {
        /// A sentinel value, which is expected to be present when the allocation is freed.
        sentinel_value: u8,

        /// The size class, which is constrained to be less than `SIZE_CLASS_COUNT`.
        size_class: usize,
    },

    /// Deallocates an allocation.
    Dealloc {
        /// The index of the allocation in the set of live allocations, which is taken to be a
        /// circular list (so this is always in-bounds if there are any allocations).
        index: usize,
    },

    /// Prints the allocator and all live allocations, for debugging purposes. This is never
    /// automatically generated.
    Debug,

    /// Overwrites an allocation, changing its sentinel.
    Overwrite {
        /// The index of the allocation in the set of live allocations, which is taken to be a
        /// circular list (so this is always in-bounds if there are any allocations).
        index: usize,

        /// The new sentinel value.
        sentinel_value: u8,
    },
}

impl Action {
    /// Generates a random action.
    pub fn any() -> impl Strategy<Value = Action> {
        prop_oneof![
            (any::<u8>(), 0..SIZE_CLASS_COUNT).prop_map(|(sentinel_value, size_class)| {
                Action::Alloc {
                    sentinel_value,
                    size_class,
                }
            }),
            any::<usize>().prop_map(|index| { Action::Dealloc { index } }),
            (any::<usize>(), any::<u8>()).prop_map(|(index, sentinel_value)| {
                Action::Overwrite {
                    index,
                    sentinel_value,
                }
            })
        ]
    }

    /// Runs an action.
    pub fn run(
        &self,
        buddy: &mut BuddyAllocator<PAGE_SIZE, PAGE_SIZE_BITS, SIZE_CLASS_COUNT>,
        allocs: &mut Vec<Alloc>,
        allow_errors: bool,
    ) {
        match *self {
            Action::Alloc {
                sentinel_value,
                size_class,
            } => match buddy.alloc(size_class) {
                Ok(ptr) => unsafe {
                    let slice = slice::from_raw_parts_mut(ptr.as_ptr(), PAGE_SIZE << size_class);
                    slice.fill(sentinel_value);
                    allocs.push(Alloc {
                        slice,
                        sentinel_value,
                    });
                },
                Err(err) => {
                    if !allow_errors {
                        Err(err).expect("failed to perform alloc action")
                    }
                }
            },
            Action::Dealloc { index } => {
                if allow_errors && allocs.is_empty() {
                    return;
                }

                let index = index % allocs.len();
                let alloc = allocs.remove(index);
                alloc.check_sentinel();
                unsafe {
                    buddy.dealloc(
                        NonNull::from(alloc.slice.as_mut()).cast(),
                        alloc.slice.len().trailing_zeros() as usize - PAGE_SIZE_BITS,
                    );
                }
            }
            Action::Debug => {
                dbg!(buddy);
                dbg!(allocs);
            }
            Action::Overwrite {
                index,
                sentinel_value,
            } => {
                if allow_errors && allocs.is_empty() {
                    return;
                }

                let index = index % allocs.len();
                let alloc = &mut allocs[index];
                alloc.slice.fill(sentinel_value);
                alloc.sentinel_value = sentinel_value;
            }
        }
    }
}

/// The entire series of actions to be performed.
#[derive(Debug)]
struct Scenario {
    /// The sizes of the ranges the buddy allocator should be initialized with.
    range_sizes: Vec<usize>,

    /// Actions to perform after initializing the buddy allocator.
    actions: Vec<Action>,
}

impl Scenario {
    /// Generates a random scenario.
    pub fn any() -> impl Strategy<Value = Scenario> {
        (
            prop::collection::vec(1usize..1 << (SIZE_CLASS_COUNT + 2), 0..4),
            prop::collection::vec(Action::any(), 0..64),
        )
            .prop_map(|(range_sizes, actions)| Scenario {
                range_sizes,
                actions,
            })
    }

    /// Runs the scenario.
    pub fn run(&self, allow_errors: bool) {
        // Allocate each of the page ranges.
        let backing_mem = self
            .range_sizes
            .iter()
            .map(|&size| {
                Mmap::new(NonZero::new(size * PAGE_SIZE).unwrap())
                    .expect("failed to allocate memory")
            })
            .collect::<Vec<_>>();

        // Create the free list allocator and move the pages there.
        let mut free_list: FreeListAllocator<PAGE_SIZE> = FreeListAllocator::new();
        for (&size, mmap) in self.range_sizes.iter().zip(backing_mem.iter()) {
            unsafe {
                free_list.add(mmap.ptr(), size);
            }
        }

        // Create the buddy allocator and move the pages from the free list to there.
        match BuddyAllocator::<PAGE_SIZE, PAGE_SIZE_BITS, SIZE_CLASS_COUNT>::new(free_list) {
            Ok(mut buddy) => {
                let mut allocs = Vec::new();

                // Run each of the actions.
                for action in &self.actions {
                    action.run(&mut buddy, &mut allocs, allow_errors);

                    // Check each allocation.
                    for alloc in &allocs {
                        alloc.check_sentinel();
                    }
                }
            }
            Err(err) => {
                if !allow_errors {
                    Err(err).expect("failed to make buddy allocator")
                }
            }
        }
    }
}

/// Information about an allocation we've made from the buddy allocator.
struct Alloc<'alloc> {
    slice: &'alloc mut [u8],
    sentinel_value: u8,
}

impl<'alloc> Alloc<'alloc> {
    pub fn check_sentinel(&self) {
        let s = self.sentinel_value;
        for (i, &b) in self.slice.iter().enumerate() {
            assert_eq!(b, s, "at index {i}, {b} != {s}",);
        }
    }
}

impl<'alloc> fmt::Debug for Alloc<'alloc> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Alloc")
            .field("ptr", &self.slice.as_ptr())
            .field("len", &self.slice.len())
            .field("sentinel_value", &self.sentinel_value)
            .finish()
    }
}

/// A mmap-allocated region of memory.
struct Mmap {
    ptr: NonNull<u8>,
    len: NonZero<usize>,
}

impl Mmap {
    pub fn new(len: NonZero<usize>) -> Result<Mmap, nix::Error> {
        let ptr = unsafe {
            mmap_anonymous(
                None,
                len,
                ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
                MapFlags::MAP_PRIVATE,
            )?
            .cast()
        };
        Ok(Mmap { ptr, len })
    }

    pub fn ptr<T>(&self) -> NonNull<T> {
        self.ptr.cast()
    }
}

impl Drop for Mmap {
    fn drop(&mut self) {
        unsafe {
            munmap(self.ptr(), self.len.get()).expect("failed to free memory");
        }
    }
}