#[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
    }

    /// Generic step:
    pub fn next_custom<const N: usize, F>(&mut self, indices: [usize; N], f: F) -> u8
    where
        F: Fn([u8; N]) -> u8,
    {
        let output = *self.state.last().unwrap();

        for &idx in &indices {
            assert!(idx < self.state.len(), "Feedback index out of bounds");
        }

        let selected: [u8; N] = std::array::from_fn(|i| self.state[indices[i]]);

        let feedback = f(selected);

        // 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;
    }

    /// Berlekamp–Massey over GF(2^8).
    ///
    /// Returns the retroaction polynomial coefficients `c` such that:
    ///   s[n] = c[1]*s[n-1] + c[2]*s[n-2] + ... + c[L]*s[n-L]
    ///
    /// The coefficients live in GF(2^8)
    ///
    /// The returned polynomial always has `c[0] = 1`.
    pub fn berlekamp_massey(sequence: &[u8]) -> Vec<u8> {
        assert!(!sequence.is_empty(), "sequence must not be empty");

        // C(x): current retroaction polynomial
        let mut c = vec![1u8];
        // B(x): copy of the last "best" polynomial
        let mut b = vec![1u8];

        let mut l: usize = 0; // current linear complexity
        let mut m: usize = 1; // number of steps since last update
        let mut bb: u8 = 1; // last non-zero discrepancy

        for n in 0..sequence.len() {
            // discrepancy d = s[n] + sum_{i=1..L} c[i] * s[n-i]
            let mut d = sequence[n];
            for i in 1..=l {
                d ^= gf_mul(c[i], sequence[n - i]);
            }

            if d == 0 {
                m += 1;
                continue;
            }

            let coef = gf_mul(d, gf_inv(bb));
            let t = c.clone();

            if c.len() < b.len() + m {
                c.resize(b.len() + m, 0);
            }

            for i in 0..b.len() {
                c[i + m] ^= gf_mul(coef, b[i]);
            }

            if 2 * l <= n {
                l = n + 1 - l;
                b = t;
                bb = d;
                m = 1;
            } else {
                m += 1;
            }
        }

        c.truncate(l + 1);
        c
    }

    pub fn predict_next_from_poly(history: &[u8], poly: &[u8]) -> u8 {
        assert!(poly.len() >= 2, "polynomial must have degree at least 1");
        assert_eq!(
            history.len(),
            poly.len() - 1,
            "history must contain exactly degree-many previous samples"
        );

        let mut next = 0u8;
        let degree = poly.len() - 1;

        for i in 1..=degree {
            next ^= gf_mul(poly[i], history[degree - i]);
        }

        next
    }
}

fn gf_mul(mut a: u8, mut b: u8) -> u8 {
    let mut p = 0u8;
    for _ in 0..8 {
        if (b & 1) != 0 {
            p ^= a;
        }
        let hi = a & 0x80; // hi for high bit
        a <<= 1;
        if hi != 0 {
            a ^= 0x1B; // reduction by 0x11B
        }
        b >>= 1;
    }
    p
}

fn gf_inv(x: u8) -> u8 {
    assert!(x != 0, "cannot invert zero in GF(2^8)");
    for y in 1u16..=255 {
        if gf_mul(x, y as u8) == 1 {
            return y as u8;
        }
    }
    unreachable!("every non-zero element in GF(2^8) has an inverse");
}

#[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_next_custom() {
        // S-box placeholder: S(n) = n + 15
        // + are XOR in this context
        fn s(x: u8) -> u8 {
            x ^ 15
        }

        // feedback = k0 + S(k2 + k7)
        let mut lfsr = Lfsr::new(10, vec![], vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);

        let output = lfsr.next_custom([0, 2, 7], |[k0, k2, k7]| k0 ^ s(k2 ^ k7));

        // Expected:
        // k0 = 1
        // k2 = 3
        // k7 = 8
        // k2 + k7 = 3 + 8 = 11
        // s(11) = 11 + 15 = 4
        // feedback = 1 + 4 = 5

        assert_eq!(output, 10);
        assert_eq!(lfsr.state(), &[5, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
    }

    #[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]);
    }

    #[test]
    fn test_berlekamp_massey() {
        // polynomial: 1 + x + x^3 + X^4 + X^7 + x^10
        let mut lfsr = Lfsr::new(10, vec![0, 1, 3, 4, 7], vec![1, 0, 0, 1, 0, 0, 1, 0, 0, 1]);

        let sequence: Vec<u8> = (0..8).map(|_| lfsr.next()).collect();
        let poly = Lfsr::berlekamp_massey(&sequence);

        // polynomial 1 + X^3
        assert_eq!(poly, vec![1, 0, 0, 1]);

        let history = &sequence[sequence.len() - (poly.len() - 1)..];
        let predicted = Lfsr::predict_next_from_poly(history, &poly);

        assert_eq!(predicted, lfsr.next());
    }
}