pub trait Bit {
	/// Test if a bit is set
    fn bit<const BIT: usize>(self) -> bool;

	/// Return a copy of self with a bit set
    #[must_use]
    fn set_bit<const BIT: usize>(self, value: bool) -> Self;
}

macro_rules! impl_bit {
	($($t:ty),*) => {
		$( impl_bit!(@ $t); )*
	};

	(@ $t:ty) => {
		impl Bit for $t {
			#[inline(always)]
			fn bit<const BIT: usize>(self) -> bool {
				self & 1 << BIT != 0
			}

			#[must_use]
			#[inline(always)]
			fn set_bit<const BIT: usize>(self, value: bool) -> Self {
				(self & !(1 << BIT)) | (value as $t) << BIT
			}
		}
	}
}

impl_bit!(u8, u16, u32, u64, u128, i8, i16, i32, i64, i128);

////////////////////////////////////////////////////////////////////////////////////

pub trait BitRange<T> {
    fn bit_range<const START: usize, const END: usize>(self) -> T;

    #[must_use]
    fn set_bit_range<const START: usize, const END: usize>(self, value: T) -> Self;
}

macro_rules! impl_bitrange {
	((),($($bitrange_ty: ty),*)) => {};

    (($t: ty),($($bitrange_ty: ty),*)) => {
        $( impl_bitrange!($t, $bitrange_ty); )*
    };

    (($t_head: ty, $($t_rest: ty),*),($($bitrange_ty: ty),*)) => {
        impl_bitrange!(($t_head), ($($bitrange_ty),*));
        impl_bitrange!(($($t_rest),*), ($($bitrange_ty),*));
    };

    ($storage: ty, $value: ty) => {
        impl BitRange<$value> for $storage {
            #[inline]
            fn bit_range<const START: usize, const END: usize>(self) -> $value {
                const VALUE_BIT_LEN: usize = core::mem::size_of::<$value>() << 3;
                let selected = END - START;
                ((self >> START) as $value) << (VALUE_BIT_LEN - selected)
                    >> (VALUE_BIT_LEN - selected)
            }

            #[inline]
			#[must_use]
            fn set_bit_range<const START: usize, const END: usize>(self, value: $value) -> Self {
                const VALUE_BIT_LEN: usize = core::mem::size_of::<$value>() << 3;
                let selected = END - START;
                let mask = (if selected == VALUE_BIT_LEN {
                    <$value>::MAX
                } else {
                    ((1 as $value) << selected) - 1
                } as $storage)
                    << START;
                (self & !mask) | ((value as $storage) << START & mask)
            }
        }
    };
}

impl_bitrange!(
    (u8, u16, u32, u64, u128, i8, i16, i32, i64, i128),
    (u8, u16, u32, u64, u128, i8, i16, i32, i64, i128)
);

////////////////////////////////////////////////////////////////////////////////////

pub trait Pattern {
	/// Determines if `self` matches a pattern
	/// 
	/// ```ignore
	/// assert!( 0x42.match_pattern("01000010") );
	/// assert!( 0x42.match_pattern("XXXXXXX0") );
	/// assert!( !0x42.match_pattern("1") );
	/// ```
    fn match_pattern(self, pattern: &str) -> bool;
}

macro_rules! impl_pattern {
	($($t:ty),*) => {
		$( impl_pattern!(@ $t); )*
	};

	(@ $t:ty) => {
		impl Pattern for $t {
			fn match_pattern(self, pattern: &str) -> bool {
				let pattern = pattern.as_bytes();
				let bits = Self::BITS as usize;
				let len = bits.min(pattern.len());

				let mut mask = 0;
				let mut expected = 0;

				for i in 0..len {
					match pattern[i] {
						b'0' => mask |= 1 << (bits - i - 1),
						b'1' => {
							mask |= 1 << (bits - i - 1);
							expected |= 1 << (bits - i - 1);
						}
						_ => {}
					}
				}

				mask >>= bits - len;
				expected >>= bits - len;

				self & mask == expected
			}
		}
	}
}

impl_pattern!(u8, u16, u32, u64, u128, i8, i16, i32, i64, i128);

////////////////////////////////////////////////////////////////////////////////////

use core::ptr;
use core::ops::RangeInclusive;

#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Ord)]
pub struct Bits<const BITCOUNT: usize, const OFFSET: isize = 0>
where
    [(); (BITCOUNT + 31) >> 5]:,
{
    storage: [u32; (BITCOUNT + 31) >> 5],
}

impl<const BITCOUNT: usize, const OFFSET: isize> Bits<BITCOUNT, OFFSET>
where
    [(); (BITCOUNT + 31) >> 5]:,
{
    pub const fn new() -> Self {
        Self {
            storage: [0; (BITCOUNT + 31) >> 5],
        }
    }

	/// Create a `Bits` from a pattern
	/// 
	/// ```ignore
	/// let mut bits: Bits<5> = Bits::new();
	/// 
    /// bits.set(1);
    /// bits.set(2);
    /// bits.set(4);
	/// 
    /// assert_eq!(Bits::new_pattern("01101"), bits);
	/// ```
    pub const fn new_pattern(pat: &str) -> Self {
        let mut bits = Self::new();
        let pat = pat.as_bytes();

        let mut idx = 0;
        while idx < pat.len() {
            if pat[idx] == b'1' {
                bits.set(idx);
            }
            idx += 1;
        }

        bits
    }

    pub const fn clear_all_bits(&mut self) {
        unsafe { ptr::write_bytes(self.storage.as_mut_ptr(), 0, self.storage.len()) }
    }

    pub const fn set_all_bits(&mut self) {
        unsafe { ptr::write_bytes(self.storage.as_mut_ptr(), 0xFF, self.storage.len()) }
    }

    /// Note: n must be less than BITCOUNT
    pub const fn get(&self, n: usize) -> bool {
        let n = n.wrapping_add(OFFSET as _);
        debug_assert!(n < BITCOUNT);

        let mask = 1 << (n & 31);
        (self.storage[n >> 5] & mask) != 0
    }

    /// Note: n must be less than BITCOUNT
    pub const fn set(&mut self, n: usize) {
        let n = n.wrapping_add(OFFSET as _);
        debug_assert!(n < BITCOUNT);

        let mask = 1 << (n & 31);
        self.storage[n >> 5] |= mask;
    }

    /// Note: n must be less than BITCOUNT
    pub const fn clear(&mut self, n: usize) {
        let n = n.wrapping_add(OFFSET as _);
        debug_assert!(n < BITCOUNT);

        let mask = 1 << (n & 31);
        self.storage[n >> 5] &= !mask;
    }

    pub const fn set_range(&mut self, range: RangeInclusive<usize>) {
        let mut n = range.start().wrapping_add(OFFSET as _);
        let mut n2 = range.end().wrapping_add(OFFSET as _);

        n2 -= 1;
        while n <= n2 {
            let a_mod = n & 31;
            let b_mod = (if n2 > (n | 31) { 31 } else { n2 & 31 }) + 1;
            let mask = ((1 << b_mod) - 1) & !((1 << a_mod) - 1);
            self.storage[n >> 5] |= mask;
            n = (n + 32) & !31;
        }
    }

    pub fn iter(&self) -> BitsIter<BITCOUNT, OFFSET> {
        BitsIter {
            bits: self,
            idx: OFFSET as _,
        }
    }
}

impl<const BITCOUNT: usize, const OFFSET: isize> Default for Bits<BITCOUNT, OFFSET>
where
    [(); (BITCOUNT + 31) >> 5]:,
{
    fn default() -> Self {
        Self::new()
    }
}

pub struct BitsIter<'a, const BITCOUNT: usize, const OFFSET: isize>
where
    [(); (BITCOUNT + 31) >> 5]:,
{
    bits: &'a Bits<BITCOUNT, OFFSET>,
    idx: usize,
}

impl<'a, const BITCOUNT: usize, const OFFSET: isize> Iterator for BitsIter<'a, BITCOUNT, OFFSET>
where
    [(); (BITCOUNT + 31) >> 5]:,
{
    type Item = bool;

    fn next(&mut self) -> Option<Self::Item> {
        if self.idx < BITCOUNT {
            let mask = 1 << (self.idx & 31);
            let val = (self.bits.storage[self.idx >> 5] & mask) != 0;
            self.idx += 1;
            Some(val)
        } else {
            None
        }
    }
}

impl<'a, const BITCOUNT: usize, const OFFSET: isize> ExactSizeIterator
    for BitsIter<'a, BITCOUNT, OFFSET>
where
    [(); (BITCOUNT + 31) >> 5]:,
{
    fn len(&self) -> usize {
        BITCOUNT - self.idx
    }
}

macro_rules! impl_froms {
	($($t:ty),+) => {
		$(
			impl From<$t> for Bits<{ <$t>::BITS as usize }, 0> {
				fn from(x: $t) -> Self {
					Self { storage: [x as _] }
				}
			}
		)+
	};

	(BIG: $($t:ty),+) => {
		$(
			impl From<$t> for Bits<{ <$t>::BITS as usize }, 0> {
				fn from(x: $t) -> Self {
					Self { storage: unsafe { core::mem::transmute(x.to_le_bytes()) } }
				}
			}
		)+
	}
}

impl_froms! { u8, i8, u16, i16, u32, i32 }
impl_froms! { BIG: u64, i64, u128, i128, usize, isize }

#[test]
fn from_big() {
    let correct: Bits<64> =
        Bits::new_pattern("1111011101111101101101010111101111110111011111011011010101111011");
    let got = Bits::from(0xDEADBEEFDEADBEEFu64);

    assert_eq!(got, correct);
}

#[test]
fn bits_iter() {
    let mut bits: Bits<5> = Bits::new();

    bits.set(1);
    bits.set(2);
    bits.set(4);

    let mut i = bits.iter();

    assert_eq!(i.next(), Some(false)); // 0
    assert_eq!(i.next(), Some(true)); // 1
    assert_eq!(i.next(), Some(true)); // 2
    assert_eq!(i.next(), Some(false)); // 3
    assert_eq!(i.next(), Some(true)); // 4
    assert_eq!(i.next(), None); // 5
}

#[test]
#[should_panic]
fn bits_set_ob() {
    let mut bits: Bits<1> = Bits::new();

    bits.set(1);
}

#[test]
fn bits_pattern() {
    let mut bits: Bits<5> = Bits::new();

    bits.set(1);
    bits.set(2);
    bits.set(4);

    assert_eq!(Bits::new_pattern("01101"), bits);
}

#[test]
fn bits_iter_exact() {
    let bits: Bits<420> = Bits::new();

    let mut i = bits.iter();
    assert_eq!(i.len(), 420);

    for _ in 0..420 {
        assert!(i.next().is_some());
    }

    assert_eq!(i.len(), 0);
}
