// Eryn Wells <eryn@erynwells.me>

use chessfriend_core::Square;

macro_rules! bit_scanner {
    ($name:ident) => {
        #[derive(Clone, Debug)]
        pub struct $name {
            bits: u64,
            shift: usize,
        }

        impl $name {
            pub(crate) fn new(bits: u64) -> Self {
                Self { bits, shift: 0 }
            }
        }
    };
}

bit_scanner!(LeadingBitScanner);
bit_scanner!(TrailingBitScanner);

fn index_to_square(index: usize) -> Square {
    debug_assert!(index < Square::NUM);
    unsafe {
        #[allow(clippy::cast_possible_truncation)]
        Square::from_index_unchecked(index as u8)
    }
}

impl Iterator for LeadingBitScanner {
    type Item = Square;

    fn next(&mut self) -> Option<Self::Item> {
        let u64bits = u64::BITS as usize;

        if self.shift == u64bits {
            return None;
        }

        let shifted_bits = self.bits << self.shift;
        let leading_zeros = shifted_bits.leading_zeros() as usize;

        if leading_zeros == u64bits {
            self.shift = leading_zeros;
            return None;
        }

        let position = u64bits - (self.shift + leading_zeros + 1);
        // Shift 1 additional place to account for the 1 that `leading_zeros` found.
        self.shift += leading_zeros + 1;

        Some(index_to_square(position))
    }
}

impl Iterator for TrailingBitScanner {
    type Item = Square;

    fn next(&mut self) -> Option<Self::Item> {
        let u64bits = u64::BITS as usize;

        if self.shift == u64bits {
            return None;
        }

        let shifted_bits = self.bits >> self.shift;
        let trailing_zeros = shifted_bits.trailing_zeros() as usize;

        if trailing_zeros == u64bits {
            self.shift = trailing_zeros;
            return None;
        }

        let position = self.shift + trailing_zeros;
        // Shift 1 additional place to account for the 1 that `leading_zeros` found.
        self.shift += trailing_zeros + 1;

        Some(index_to_square(position))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn leading_zero() {
        let mut scanner = LeadingBitScanner::new(0);
        assert_eq!(scanner.next(), None);
    }

    #[test]
    fn leading_one() {
        let mut scanner = LeadingBitScanner::new(1);
        assert_eq!(scanner.next(), Some(Square::A1));
        assert_eq!(scanner.next(), None);
    }

    #[test]
    fn leading_complex() {
        let mut scanner = LeadingBitScanner::new(0b_1100_0101);
        assert_eq!(scanner.next(), Some(Square::H1));
        assert_eq!(scanner.next(), Some(Square::G1));
        assert_eq!(scanner.next(), Some(Square::C1));
        assert_eq!(scanner.next(), Some(Square::A1));
        assert_eq!(scanner.next(), None);
    }

    #[test]
    fn trailing_zero() {
        let mut scanner = TrailingBitScanner::new(0);
        assert_eq!(scanner.next(), None);
    }

    #[test]
    fn trailing_one() {
        let mut scanner = TrailingBitScanner::new(1);
        assert_eq!(scanner.next(), Some(Square::A1));
        assert_eq!(scanner.next(), None);
    }

    #[test]
    fn trailing_complex() {
        let mut scanner = TrailingBitScanner::new(0b_1100_0101);
        assert_eq!(scanner.next(), Some(Square::A1));
        assert_eq!(scanner.next(), Some(Square::C1));
        assert_eq!(scanner.next(), Some(Square::G1));
        assert_eq!(scanner.next(), Some(Square::H1));
        assert_eq!(scanner.next(), None);
    }
}