#![no_std]
extern crate alloc;
use alloc::boxed::Box;

/// This is the main data structure that the crate contains. It consists of a tree
/// of bits that is specifically meant for an O(1) `find` function. It has slow
/// indexing, pushes, and removals, however, and so is mainly useful for seldom
/// changed datasets that are searched consistently. 
///
/// Example usage:
/// ```rust
/// use bittree::BitTree;
/// let mut btree = BitTree::new();
/// btree.add(&0usize, Some(0usize));
/// assert_eq!(btree.find(&0usize), Some(0usize));
/// 
/// btree.remove(&0usize);
/// assert_eq!(btree.find(&0usize), None);
/// ```

#[derive(Debug, Clone)]
pub struct BitTree<K: GetBytes, V: GetBytes> {
    is_keys: bool,
    value: Option<V>,
    next_vals: *mut BitTree<V, K>,
    next: [Option<*mut BitTree<K, V>>; 2],
    phdata: core::marker::PhantomData<K>
}

pub mod getbytes;
use getbytes::GetBytes;

#[derive(Debug, Clone)]
struct Bits<'a> {
    inner: &'a [u8],
    bitpos: usize
}

impl<'a> Iterator for Bits<'a> {
    type Item = bool;
    fn next(&mut self) -> Option<Self::Item> {
        if self.bitpos >> 3 >= self.inner.len() {
            None
        } else {
            let shift_by = core::mem::size_of::<usize>() * 8 - 3;
            let current_bit = (self.bitpos << shift_by) >> shift_by;
            let out = Some(self.inner[self.bitpos >> 3] & (1 << current_bit) != 0);
            self.bitpos += 1;

            out
        }
    }
}

impl<K: GetBytes, V: GetBytes> Default for BitTree<K, V> {
    fn default() -> Self {
        let mut out = BitTree {
            is_keys: true,
            value: None,
            next_vals: 1usize as *mut BitTree<V, K>,
            next: [None, None],
            phdata: core::marker::PhantomData
        };

        out.next_vals = Box::leak(Box::new(BitTree {
            value: None,
            is_keys: false,
            next_vals: &mut out as *mut BitTree<K, V>,
            next: [None, None],
            phdata: core::marker::PhantomData
        }));

        out
    }
}

impl<K: GetBytes, V: GetBytes> BitTree<K, V> {
    /// This creates a new BitTree.
    pub fn new() -> BitTree<K, V> {
        Self::default()
    }

    /// To add to the bittree, this function gets all of the bits of the input data and creates new 
    /// entries for each non-existent part of the path of bits until it completes the bitpath. It adds
    /// a new value at the last entry.
    pub fn add(&mut self, key: &K, value: Option<V>) {
        let bytes = key.get_bytes();
        let bits = Bits {
            inner: bytes,
            bitpos: 0
        };

        let mut current = self;
        if current.is_keys && value.is_some() {
            unsafe {
                (&mut *current.next_vals).add(&value.clone().unwrap(), Some(key.clone()));
            }
        }

        for bit in bits {
            unsafe {
                if current.next[bit as usize].as_mut().is_some() {
                    current = &mut **current.next[bit as usize].as_mut().unwrap();
                } else {
                    current.next[bit as usize] = Some(Box::leak(Box::new(BitTree::new())));
                    current = &mut **current.next[bit as usize].as_mut().unwrap();
                }
            }
        }

        current.value = value;
    }

    /// To remove from the tree of bits, this function gets the second to last entry 
    /// for the input data and removes the child entry used for the data.
    pub fn remove(&mut self, key: &K) {
        let bytes = key.get_bytes();

        let bits = Bits {
            inner: bytes,
            bitpos: 0
        };

        let mut current = self;
        if current.is_keys {
            unsafe {
                let res = (&mut *current.next_vals).find(key);
                (&mut *current.next_vals).remove(res.as_ref().unwrap());
            }
        }

        for (i, bit) in bits.enumerate() {
            if current.next[bit as usize].is_some() && i == core::mem::size_of::<K>() * 8 - 2 {
                current.next[bit as usize] = None;
            } else if current.next[bit as usize].is_none() {
                break;
            } else {
                unsafe {
                    current = &mut **current.next[bit as usize].as_mut().unwrap();
                }
            }
        }
    }

    /// To search through the data and find a value, this function checks if all of the required 
    /// bit entries specified in the bitpath exist in the values-to-keys bitpath, which solely 
    /// depends on the size of the data being searched for (a constant, hence constant time) rather 
    /// than the size of the dataset. 
    pub fn find(&mut self, value: &V) -> Option<K> {
        let bytes = value.get_bytes();

        let bits = Bits {
            inner: bytes,
            bitpos: 0
        };

        let mut current = unsafe {
            &mut *self.next_vals
        };
        for bit in bits {
            unsafe {
                if let Some(next) = current.next[bit as usize].as_mut() {
                    current = &mut **next;
                } else {
                    return None;
                }
            }
        }

        current.value.clone()
    }
}

impl<K: GetBytes, V: GetBytes> core::ops::Index<&K> for BitTree<K, V> {
    type Output = Option<V>;
    fn index(&self, index: &K) -> &Self::Output {
        let bytes = index.get_bytes();
        let bits = Bits {
            inner: bytes,
            bitpos: 0
        };

        let mut current = self;
        for bit in bits {
            if current.next[bit as usize].is_some() {
                unsafe {
                    current = &**current.next[bit as usize].as_ref().unwrap();
                }
            } else {
                return &current.value;
            }
        }

        &current.value
    }
}

impl<K: GetBytes, V: GetBytes> core::ops::IndexMut<&K> for BitTree<K, V> {
    fn index_mut(&mut self, index: &K) -> &mut Option<V> {
        let bytes = index.get_bytes();
        let bits = Bits {
            inner: bytes,
            bitpos: 0
        };

        let mut current = self;
        for bit in bits {
            if current.next[bit as usize].is_some() {
                unsafe {
                    current = &mut **current.next[bit as usize].as_mut().unwrap();
                }
            } else {
                return &mut current.value;
            }
        }

        &mut current.value
    }
}

#[test]
fn basic_add_remove() {
    let mut btree = BitTree::new();
    btree.add(&0usize, Some(0usize));
    assert!(btree.find(&0usize).is_some());
    assert_eq!(btree[&0usize], Some(0usize));
    btree.remove(&0usize);
    assert!(btree.find(&0usize).is_none());
    assert!(btree[&0usize].is_none());
}