#[derive(Debug, Clone)]
pub struct Lfsr {
    /// Internal state (each element is 1B)
    state: Vec<u8>,
    /// Tap positions (bytes oriented)
    taps: Vec<usize>,
}

impl Lfsr {
    /// Create a new LFSR
    ///
    /// - `length`: LFSR size (number of bytes)
    /// - `taps`: indices of tapped bytes (0 based)
    /// - `init`: initial state
    pub fn new(length: usize, taps: Vec<usize>, init: Vec<u8>) -> Self {
        assert_eq!(init.len(), length, "Initial state length mismatch");

        for &tap in &taps {
            assert!(tap < length, "Tap index out of bounds");
        }

        Self { state: init, taps }
    }

    /// LFSR step
    ///
    /// - Returns the output byte
    /// - Shifts the register and inserts feedback at position 0
    pub fn next(&mut self) -> u8 {
        let output = *self.state.last().unwrap();

        let mut feedback: u8 = 0;
        for &tap in &self.taps {
            feedback ^= self.state[tap];
        }

        // Shift right
        for i in (1..self.state.len()).rev() {
            self.state[i] = self.state[i - 1];
        }

        self.state[0] = feedback;

        output
    }

    /// Get current state
    pub fn state(&self) -> &[u8] {
        &self.state
    }

    /// Get taps
    pub fn taps(&self) -> &[usize] {
        &self.taps
    }

    /// Reset state
    ///
    /// - `new_state`: new initial state
    pub fn reset(&mut self, new_state: Vec<u8>) {
        assert_eq!(new_state.len(), self.state.len(), "State length mismatch");
        self.state = new_state;
    }
}

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

    #[test]
    fn test_initial_state() {
        let lfsr = Lfsr::new(3, vec![0, 1], vec![1, 2, 3]);
        assert_eq!(lfsr.state(), &[1, 2, 3]);
        assert_eq!(lfsr.taps(), vec![0, 1]);
    }

    #[test]
    fn test_steps() {
        let mut lfsr = Lfsr::new(3, vec![0, 1], vec![1, 2, 3]);
        let outputs: Vec<u8> = (0..4).map(|_| lfsr.next()).collect();
        // step1: [1,2,3] -> [3,1,2], out=3
        // step2: [3,1,2] -> [2,3,1], out=2
        // step3: [2,3,1] -> [1,2,3], out=1
        // step4: [1,2,3] -> [3,1,2], out=3
        assert_eq!(outputs, vec![3, 2, 1, 3]);
        assert_eq!(lfsr.state(), &[3, 1, 2]);
    }

    #[test]
    fn test_empty_taps() {
        let mut lfsr = Lfsr::new(3, vec![], vec![1, 2, 3]);
        let out = lfsr.next();
        assert_eq!(out, 3);
        assert_eq!(lfsr.state(), &[0, 1, 2]);
    }

    #[test]
    fn test_reset() {
        let mut lfsr = Lfsr::new(3, vec![0], vec![1, 2, 3]);
        lfsr.next();
        lfsr.reset(vec![9, 8, 7]);
        assert_eq!(lfsr.state(), &[9, 8, 7]);
    }

    #[test]
    #[should_panic(expected = "Initial state length mismatch")]
    fn test_invalid_init_length() {
        Lfsr::new(3, vec![0], vec![1, 2]);
    }

    #[test]
    #[should_panic(expected = "Tap index out of bounds")]
    fn test_invalid_tap() {
        Lfsr::new(3, vec![3], vec![1, 2, 3]);
    }

    #[test]
    #[should_panic(expected = "State length mismatch")]
    fn test_invalid_reset_length() {
        let mut lfsr = Lfsr::new(3, vec![0], vec![1, 2, 3]);
        lfsr.reset(vec![1, 2]);
    }
}