//! Byte-granularity right shift (division by 256^n).
use super::U256;

impl U256 {
    /// Shifts the value right by `n` bytes (divides by 256^n).
    ///
    /// Bytes shifted past the least significant position are discarded.
    /// Returns zero when `n >= 32`.
    ///
    /// # Examples
    ///
    /// ```
    /// use cnfy_uint::u256::U256;
    ///
    /// let a = U256::from_be_limbs([0, 0, 1, 0]);
    /// assert_eq!(a.shr_bytes(8), U256::from_be_limbs([0, 0, 0, 1]));
    /// ```
    #[inline]
    pub const fn shr_bytes(&self, n: u32) -> U256 {
        if n >= 32 {
            return U256([0, 0, 0, 0]);
        }
        let limb_shift = (n / 8) as usize;
        let bit_shift = (n % 8) * 8;
        let mut result = [0u64; 4];
        let mut i = 0;
        while i + limb_shift < 4 {
            let src = i + limb_shift;
            result[i] = self.0[src] >> bit_shift;
            if bit_shift > 0 && src + 1 < 4 {
                result[i] |= self.0[src + 1] << (64 - bit_shift);
            }
            i += 1;
        }
        U256(result)
    }
}

#[cfg(test)]
mod ai_tests {
    use super::*;
    /// Shifting by zero is the identity.
    #[test]
    fn identity() {
        let a = U256::from_be_limbs([0x1234, 0x5678, 0x9ABC, 0xDEF0]);
        assert_eq!(a.shr_bytes(0), a);
    }

    /// Shift by one byte moves data 8 bits right.
    #[test]
    fn one_byte() {
        assert_eq!(
            U256::from_be_limbs([0, 0, 0, 0xFF00]).shr_bytes(1),
            U256::from_be_limbs([0, 0, 0, 0xFF]),
        );
    }

    /// Shift by one full limb (8 bytes).
    #[test]
    fn one_limb() {
        assert_eq!(
            U256::from_be_limbs([0, 0, 1, 0]).shr_bytes(8),
            U256::from_be_limbs([0, 0, 0, 1]),
        );
    }

    /// Shift across all limbs.
    #[test]
    fn three_limbs() {
        assert_eq!(
            U256::from_be_limbs([1, 0, 0, 0]).shr_bytes(24),
            U256::from_be_limbs([0, 0, 0, 1]),
        );
    }

    /// Cross-limb byte shift carries bits between limbs.
    #[test]
    fn cross_limb_carry() {
        assert_eq!(
            U256::from_be_limbs([0, 0, 1, 0]).shr_bytes(1),
            U256::from_be_limbs([0, 0, 0, 0x0100000000000000]),
        );
    }

    /// Shifting by 32 or more yields zero.
    #[test]
    fn full_shift() {
        let max = U256::from_be_limbs([u64::MAX, u64::MAX, u64::MAX, u64::MAX]);
        assert_eq!(max.shr_bytes(32), U256::ZERO);
    }

    /// shl then shr by the same amount is the identity (when no overflow).
    #[test]
    fn round_trip() {
        let a = U256::from_be_limbs([0, 0, 0x1234, 0x5678]);
        assert_eq!(a.shl_bytes(5).shr_bytes(5), a);
    }
}
