#[cfg(test)]
mod test;

mod entry;
mod into_iter;
mod iter;
mod iter_mut;
mod keys;
mod values;
mod values_mut;

pub use entry::{Entry, OccupiedEntry, VacantEntry};
pub use into_iter::IntoIter;
pub use iter::Iter;
pub use iter_mut::IterMut;
pub use keys::Keys;
pub use values::Values;
pub use values_mut::ValuesMut;

use core::{
    borrow::Borrow,
    cmp::Ord,
    fmt,
    hash::{Hash, Hasher},
    mem::{self, MaybeUninit},
    ops::Index,
    ptr,
};

/// A heapless, ord-based array map. Operates very similarly to [`BTreeMap`],
/// but maintains its keys and values as single arrays stored inside itself.
/// Sorted insertions are optimized and will be processed in constant time.
///
/// [`BTreeMap`]: https://doc.rust-lang.org/std/collections/struct.BTreeMap.html
pub struct OrdMap<K, V, const N: usize> {
    len: usize,
    keys: [MaybeUninit<K>; N],
    values: [MaybeUninit<V>; N],
}

impl<K, V, const N: usize> OrdMap<K, V, N> {
    fn insert_index(&self, key: &K) -> Result<usize, usize>
    where
        K: Ord,
    {
        let last = match self.last_key_value() {
            Some(last) => last.0,
            None => return Ok(0),
        };
        if last < key {
            Ok(self.len)
        } else {
            // We swap these, because Ok for us is an insertion point.
            match self.keys_slice().binary_search(key) {
                Ok(index) => Err(index),
                Err(index) => Ok(index),
            }
        }
    }
    pub fn keys_slice(&self) -> &[K] {
        let ptr = self.keys.as_ptr() as *const K;
        unsafe { core::slice::from_raw_parts(ptr, self.len) }
    }
    pub fn values_slice(&self) -> &[V] {
        let ptr = self.values.as_ptr() as *const V;
        unsafe { core::slice::from_raw_parts(ptr, self.len) }
    }
    pub fn values_mut_slice(&mut self) -> &mut [V] {
        let ptr = self.values.as_mut_ptr() as *mut V;
        unsafe { core::slice::from_raw_parts_mut(ptr, self.len) }
    }
    pub fn slices(&self) -> (&[K], &[V]) {
        (self.keys_slice(), self.values_slice())
    }
    pub fn slices_mut(&mut self) -> (&[K], &mut [V]) {
        let keys = {
            let ptr = self.keys.as_ptr() as *const K;
            unsafe { core::slice::from_raw_parts(ptr, self.len) }
        };
        let values = {
            let ptr = self.values.as_mut_ptr() as *mut V;
            unsafe { core::slice::from_raw_parts_mut(ptr, self.len) }
        };
        (keys, values)
    }

    pub fn new() -> Self {
        Self::default()
    }

    pub fn len(&self) -> usize {
        self.len
    }
    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    pub fn entry(&mut self, key: K) -> Entry<'_, K, V, N>
    where
        K: Ord,
    {
        match self.insert_index(&key) {
            Err(index) => Entry::occupied(self, index, key),
            Ok(index) => Entry::vacant(self, index, key),
        }
    }

    /// Insert the value into the map.  Note that the semantics of this
    /// operation are not exactly the same as for BTreeMap.  This will not
    /// replace existing values.
    ///  The possible results are:
    /// * `Ok(Some((K, V)))`: If the key already existed in the map.  The key
    ///   and value are not inserted, and they are returned back to the caller.
    /// * `Ok(None)`: If the key and value were inserted successfully without
    ///   replacement.
    /// * `Err((K, V))`: If the set was full; the and value are returned to the caller.
    pub fn insert(&mut self, key: K, value: V) -> Result<Option<(K, V)>, (K, V)>
    where
        K: Ord,
    {
        match self.insert_index(&key) {
            Err(_) => Ok(Some((key, value))),
            Ok(index) => self.insert_at(index, key, value).map(|()| None),
        }
    }

    /// A raw insert at a particular index.  The key and value are returned on a
    /// full map.
    fn insert_at(&mut self, index: usize, key: K, value: V) -> Result<(), (K, V)> {
        if self.len == N {
            Err((key, value))
        } else {
            self.keys[index..=self.len].rotate_right(1);
            self.keys[index].write(key);
            self.values[index..=self.len].rotate_right(1);
            self.values[index].write(value);
            self.len += 1;
            Ok(())
        }
    }

    /// Replace the key and value at the index and return the old ones.
    fn replace_at(&mut self, index: usize, key: K, value: V) -> (K, V) {
        let prev_key = mem::replace(&mut self.keys[index], MaybeUninit::new(key));
        let prev_value = mem::replace(&mut self.values[index], MaybeUninit::new(value));
        unsafe { (prev_key.assume_init(), prev_value.assume_init()) }
    }

    /// Insert the value into the map, replacing an existing key if found.
    ///  The possible results are:
    /// * `Ok(Some((K, V)))`: If the key already existed in the map.  The key
    ///   and value are replaced, and the previous key and value are returned
    ///   back to the caller.
    /// * `Ok(None)`: If the key and value were inserted successfully without
    ///   replacement.
    /// * `Err((K, V))`: If the set was full; the and value are returned to the caller.
    pub fn replace(&mut self, key: K, value: V) -> Result<Option<(K, V)>, (K, V)>
    where
        K: Ord,
    {
        match self.insert_index(&key) {
            Err(index) => Ok(Some(self.replace_at(index, key, value))),
            Ok(index) => self.insert_at(index, key, value).map(|()| None),
        }
    }

    pub fn get<Q>(&self, key: &Q) -> Option<&V>
    where
        K: Borrow<Q>,
        Q: Ord + ?Sized,
    {
        self.keys_slice()
            .binary_search_by(move |k| k.borrow().cmp(key))
            .ok()
            .map(|index| unsafe { self.values[index].assume_init_ref() })
    }

    pub fn get_key_value<Q>(&self, key: &Q) -> Option<(&K, &V)>
    where
        K: Borrow<Q>,
        Q: Ord + ?Sized,
    {
        self.keys_slice()
            .binary_search_by(move |k| k.borrow().cmp(key))
            .ok()
            .map(|index| unsafe {
                (
                    self.keys[index].assume_init_ref(),
                    self.values[index].assume_init_ref(),
                )
            })
    }

    pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
    where
        K: Borrow<Q>,
        Q: Ord + ?Sized,
    {
        self.keys_slice()
            .binary_search_by(move |k| k.borrow().cmp(key))
            .ok()
            .map(|index| unsafe { self.values[index].assume_init_mut() })
    }

    pub fn contains_key<Q>(&self, key: &Q) -> bool
    where
        K: Borrow<Q>,
        Q: Ord + ?Sized,
    {
        self.keys_slice()
            .binary_search_by(move |k| k.borrow().cmp(key))
            .is_ok()
    }

    /// Remove the value from the map.
    /// Returns the removed value.
    pub fn remove<Q>(&mut self, key: &Q) -> Option<(K, V)>
    where
        K: Borrow<Q>,
        Q: Ord + ?Sized,
    {
        self.keys_slice()
            .binary_search_by(move |k| k.borrow().cmp(key))
            .ok()
            .map(move |index| self.remove_at(index))
    }

    /// Remove the entry at the index.
    /// Returns the removed key and value.
    fn remove_at(&mut self, index: usize) -> (K, V) {
        self.keys[index..=self.len].rotate_left(1);
        self.values[index..=self.len].rotate_left(1);
        self.len -= 1;
        unsafe {
            (
                self.keys[self.len].assume_init_read(),
                self.values[self.len].assume_init_read(),
            )
        }
    }

    pub fn clear(&mut self) {
        if self.len == 0 {
            return;
        }

        let ptr = self.keys.as_ptr() as *mut K;
        let key_slice = ptr::slice_from_raw_parts_mut(ptr, self.len);
        let ptr = self.values.as_ptr() as *mut V;
        let value_slice = ptr::slice_from_raw_parts_mut(ptr, self.len);
        self.len = 0;
        unsafe {
            ptr::drop_in_place(key_slice);
            ptr::drop_in_place(value_slice);
        }
    }

    pub fn iter(&self) -> Iter<'_, K, V> {
        Iter::new(self.keys_slice(), self.values_slice())
    }

    pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
        let (keys, values) = self.slices_mut();
        IterMut::new(keys, values)
    }

    pub fn keys(&self) -> Keys<'_, K> {
        Keys::new(self.keys_slice())
    }
    pub fn values(&self) -> Values<'_, V> {
        Values::new(self.values_slice())
    }
    pub fn values_mut(&mut self) -> ValuesMut<'_, V> {
        ValuesMut::new(self.values_mut_slice())
    }
    pub fn first_key_value(&self) -> Option<(&K, &V)> {
        if self.len == 0 {
            None
        } else {
            let key_slot = &self.keys[0];
            let value_slot = &self.values[0];
            Some(unsafe { (key_slot.assume_init_ref(), value_slot.assume_init_ref()) })
        }
    }
    pub fn last_key_value(&self) -> Option<(&K, &V)> {
        if self.len == 0 {
            None
        } else {
            let key_slot = &self.keys[self.len - 1];
            let value_slot = &self.values[self.len - 1];
            Some(unsafe { (key_slot.assume_init_ref(), value_slot.assume_init_ref()) })
        }
    }
    pub fn pop_first(&mut self) -> Option<(K, V)> {
        if self.len == 0 {
            None
        } else {
            self.keys[..self.len].rotate_left(1);
            self.values[..self.len].rotate_left(1);
            self.len -= 1;
            let key_slot = &self.keys[self.len];
            let value_slot = &self.values[self.len];
            Some(unsafe { (key_slot.assume_init_read(), value_slot.assume_init_read()) })
        }
    }
    pub fn pop_last(&mut self) -> Option<(K, V)> {
        if self.len == 0 {
            None
        } else {
            self.len -= 1;
            let key_slot = &self.keys[self.len];
            let value_slot = &self.values[self.len];
            Some(unsafe { (key_slot.assume_init_read(), value_slot.assume_init_read()) })
        }
    }
}

impl<K, V, const N: usize> Default for OrdMap<K, V, N> {
    fn default() -> Self {
        Self {
            len: 0,
            keys: unsafe { MaybeUninit::uninit().assume_init() },
            values: unsafe { MaybeUninit::uninit().assume_init() },
        }
    }
}
impl<K, V, const N: usize> Clone for OrdMap<K, V, N>
where
    K: Clone,
    V: Clone,
{
    fn clone(&self) -> Self {
        let mut keys: [MaybeUninit<K>; N] = unsafe { MaybeUninit::uninit().assume_init() };
        let mut values: [MaybeUninit<V>; N] = unsafe { MaybeUninit::uninit().assume_init() };

        for (source, destination) in self.keys_slice().iter().zip(&mut keys[..self.len]) {
            destination.write(source.clone());
        }

        for (source, destination) in self.values_slice().iter().zip(&mut values[..self.len]) {
            destination.write(source.clone());
        }

        Self {
            keys,
            values,
            len: self.len,
        }
    }
}

impl<K, V, const N: usize> Drop for OrdMap<K, V, N> {
    fn drop(&mut self) {
        self.clear();
    }
}

impl<K, V, const N: usize> fmt::Debug for OrdMap<K, V, N>
where
    K: fmt::Debug,
    V: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_map().entries(self.iter()).finish()
    }
}

impl<K, Q, V, const N: usize> Index<&Q> for OrdMap<K, V, N>
where
    K: Borrow<Q>,
    Q: Ord + ?Sized,
{
    type Output = V;

    fn index(&self, key: &Q) -> &Self::Output {
        self.get(key).unwrap()
    }
}

impl<K, V, const N: usize> IntoIterator for OrdMap<K, V, N> {
    type Item = (K, V);

    type IntoIter = IntoIter<K, V, N>;

    fn into_iter(self) -> Self::IntoIter {
        IntoIter::new(self)
    }
}

impl<'a, K, V, const N: usize> IntoIterator for &'a OrdMap<K, V, N> {
    type Item = (&'a K, &'a V);

    type IntoIter = Iter<'a, K, V>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

impl<'a, K, V, const N: usize> IntoIterator for &'a mut OrdMap<K, V, N> {
    type Item = (&'a K, &'a mut V);

    type IntoIter = IterMut<'a, K, V>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter_mut()
    }
}

impl<K, V, const N1: usize, const N2: usize> PartialOrd<OrdMap<K, V, N2>> for OrdMap<K, V, N1>
where
    K: PartialOrd<K>,
    V: PartialOrd<V>,
{
    fn partial_cmp(&self, other: &OrdMap<K, V, N2>) -> Option<core::cmp::Ordering> {
        self.iter().partial_cmp(other.iter())
    }
}
impl<K, V, const N1: usize, const N2: usize> PartialEq<OrdMap<K, V, N2>> for OrdMap<K, V, N1>
where
    K: PartialEq<K>,
    V: PartialEq<V>,
{
    fn eq(&self, other: &OrdMap<K, V, N2>) -> bool {
        self.iter().eq(other.iter())
    }
}

impl<K, V, const N: usize> Eq for OrdMap<K, V, N>
where
    K: Eq,
    V: Eq,
{
}

impl<K, V, const N: usize> Ord for OrdMap<K, V, N>
where
    K: Ord,
    V: Ord,
{
    fn cmp(&self, other: &OrdMap<K, V, N>) -> core::cmp::Ordering {
        self.iter().cmp(other.iter())
    }
}

impl<K, V, const N: usize> Hash for OrdMap<K, V, N>
where
    K: Hash,
    V: Hash,
{
    fn hash<H: Hasher>(&self, state: &mut H) {
        state.write_usize(self.len);
        for pair in self {
            pair.hash(state);
        }
    }
}

/// Panics on overflow
impl<K, V, const N: usize> Extend<(K, V)> for OrdMap<K, V, N>
where
    K: Ord,
{
    fn extend<I>(&mut self, iter: I)
    where
        I: IntoIterator<Item = (K, V)>,
    {
        for (key, value) in iter {
            if let Err(_) = self.insert(key, value) {
                panic!("map overflowed on extend");
            }
        }
    }
}

/// Panics on overflow
impl<'a, K, V, const N: usize> Extend<(&'a K, &'a V)> for OrdMap<K, V, N>
where
    K: Ord + Copy,
    V: Copy,
{
    fn extend<I>(&mut self, iter: I)
    where
        I: IntoIterator<Item = (&'a K, &'a V)>,
    {
        for (key, value) in iter {
            if let Err(_) = self.insert(*key, *value) {
                panic!("map overflowed on extend");
            }
        }
    }
}

/// Panics on overflow
impl<K, V, const N: usize> FromIterator<(K, V)> for OrdMap<K, V, N>
where
    K: Ord,
{
    fn from_iter<I>(iter: I) -> Self
    where
        I: IntoIterator<Item = (K, V)>,
    {
        let mut map = Self::new();
        for (key, value) in iter {
            if let Err(_) = map.insert(key, value) {
                panic!("map overflowed on extend");
            }
        }
        map
    }
}

impl<K, V, const N1: usize, const N2: usize> From<[(K, V); N1]> for OrdMap<K, V, N2>
where
    K: Ord,
{
    fn from(mut value: [(K, V); N1]) -> Self {
        value.sort_unstable_by(|a, b| a.0.cmp(&b.0));
        let mut map = Self::new();
        for (key, value) in value {
            if map.len == 0 || unsafe { map.keys[map.len - 1].assume_init_ref() } < &key {
                if map.len == N2 {
                    panic!("map overflowed on extend");
                }
                map.keys[map.len].write(key);
                map.values[map.len].write(value);
                map.len += 1;
            }
        }
        map
    }
}