use crate::crypto::constants::{NEAR_PRG_CTX, SECURITY_PARAMETER};
use crate::errors::ProtocolError;
use auto_ops::impl_op_ex;
use rand_core::CryptoRngCore;
use serde::{Deserialize, Serialize};
use sha3::{
    Shake256,
    digest::{ExtendableOutput, Update, XofReader},
};
use subtle::{Choice, ConditionallySelectable, ConstantTimeEq};

pub const SEC_PARAM_64: usize = SECURITY_PARAMETER.div_ceil(64);
pub const SEC_PARAM_8: usize = SECURITY_PARAMETER.div_ceil(8);

/// Represents a vector of bits.
///
/// This vector will have the size of our security parameter, which is useful
/// for most of our OT extension protocols.
#[derive(Clone, Copy, PartialEq, Serialize, Deserialize, Eq)]
pub struct BitVector([u64; SEC_PARAM_64]);

impl_secret_debug!(BitVector);

impl BitVector {
    pub fn zero() -> Self {
        Self([0u64; SEC_PARAM_64])
    }

    /// Return a random bit vector.
    pub fn random(rng: &mut impl CryptoRngCore) -> Self {
        let mut out = [0u64; SEC_PARAM_64];
        for o in &mut out {
            *o = rng.next_u64();
        }
        Self(out)
    }

    /// Get a specific bit from the vector.
    ///
    /// # Panics
    /// Panics if `j >= SECURITY_PARAMETER`.
    pub fn bit(&self, j: usize) -> u8 {
        let word = self
            .0
            .get(j / 64)
            .expect("bit index must be < SECURITY_PARAMETER");
        ((word >> (j % 64)) & 1) as u8
    }

    pub fn from_bytes(bytes: &[u8; SEC_PARAM_8]) -> Self {
        let u64s = bytes.chunks_exact(8).map(|chunk| {
            u64::from_le_bytes(
                chunk
                    .try_into()
                    .expect("Cannot fail as chunks_exact takes 8 bytes"),
            )
        });
        let mut out = [0u64; SEC_PARAM_64];
        for (o, u) in out.iter_mut().zip(u64s) {
            *o = u;
        }
        Self(out)
    }

    pub fn bytes(&self) -> [u8; SEC_PARAM_8] {
        let mut out = [0u8; SEC_PARAM_8];
        for (chunk, x_i) in out.chunks_exact_mut(8).zip(self.0.iter()) {
            chunk.copy_from_slice(&x_i.to_le_bytes());
        }
        out
    }

    /// Iterate over the bits of this vector.
    pub fn bits(&self) -> impl Iterator<Item = Choice> {
        self.0
            .into_iter()
            .flat_map(|u| (0..64).map(move |j| ((u >> j) & 1).ct_eq(&1)))
    }

    /// Modify this vector by xoring it with another vector.
    pub fn xor_mut(&mut self, other: &Self) {
        for (self_i, other_i) in self.0.iter_mut().zip(other.0.iter()) {
            *self_i ^= other_i;
        }
    }

    /// Xor this vector with another.
    pub fn xor(&self, other: &Self) -> Self {
        let mut out = *self;
        out.xor_mut(other);
        out
    }

    /// Return the bitwise not of this vector.
    pub fn not(&self) -> Self {
        let mut out = *self;
        for out_i in &mut out.0 {
            *out_i = !*out_i;
        }
        out
    }

    pub fn and_mut(&mut self, other: &Self) {
        for (self_i, other_i) in self.0.iter_mut().zip(other.0.iter()) {
            *self_i &= other_i;
        }
    }

    pub fn and(&self, other: &Self) -> Self {
        let mut out = *self;
        out.and_mut(other);
        out
    }

    /// Multiplication in the field.
    ///
    /// This returns an unreduced value, which is fine for our use case.
    pub fn gf_mul(&self, other: &Self) -> DoubleBitVector {
        // Algorithm 2.35 in "Guide to Elliptic Curve Cryptography"
        let mut out = [0u64; 2 * SEC_PARAM_64];

        for k in (0..64).rev() {
            for (j, self_word) in self.0.iter().enumerate() {
                let to_add = Self::conditional_select(
                    &Self::zero(),
                    other,
                    Choice::from(((*self_word >> k) & 1) as u8),
                );

                for (out_elem, add_elem) in out.iter_mut().skip(j).zip(to_add.0.iter()) {
                    *out_elem ^= add_elem;
                }
            }
            if k != 0 {
                let mut prev = 0u64;
                for out_i in &mut out {
                    let next_prev = *out_i >> 63;
                    *out_i = (*out_i << 1) | prev;
                    prev = next_prev;
                }
            }
        }

        DoubleBitVector(out)
    }
}

impl ConditionallySelectable for BitVector {
    fn conditional_select(a: &Self, b: &Self, choice: Choice) -> Self {
        let mut out = [0u64; SEC_PARAM_64];
        for ((o_i, a_i), b_i) in out.iter_mut().zip(a.0.iter()).zip(b.0.iter()) {
            *o_i = u64::conditional_select(a_i, b_i, choice);
        }
        Self(out)
    }
}

impl_op_ex!(^ |u: &BitVector, v: &BitVector| -> BitVector { u.xor(v) });
impl_op_ex!(^= |u: &mut BitVector, v: &BitVector| { u.xor_mut(v) });
impl_op_ex!(&|u: &BitVector, v: &BitVector| -> BitVector { u.and(v) });
impl_op_ex!(&= |u: &mut BitVector, v: &BitVector| { u.and_mut(v) });
impl_op_ex!(!|u: &BitVector| -> BitVector { u.not() });

/// A `BitVector` of double the size.
///
/// This is useful because it's quicker to avoid reducing the result of GF multiplication.
#[derive(Clone, Copy, Serialize, Deserialize)]
#[cfg_attr(test, derive(PartialEq, Eq))]
pub struct DoubleBitVector([u64; Self::SIZE]);

impl_secret_debug!(DoubleBitVector);

impl DoubleBitVector {
    const SIZE: usize = 2 * SEC_PARAM_64;

    pub fn zero() -> Self {
        Self([0u64; Self::SIZE])
    }

    pub fn xor_mut(&mut self, other: &Self) {
        for (self_i, other_i) in self.0.iter_mut().zip(other.0.iter()) {
            *self_i ^= *other_i;
        }
    }

    pub fn xor(&self, other: &Self) -> Self {
        let mut out = *self;
        out.xor_mut(other);
        out
    }
}

impl ConditionallySelectable for DoubleBitVector {
    fn conditional_select(a: &Self, b: &Self, choice: Choice) -> Self {
        let mut out = [0u64; 2 * SEC_PARAM_64];
        for ((o_i, a_i), b_i) in out.iter_mut().zip(a.0.iter()).zip(b.0.iter()) {
            *o_i = u64::conditional_select(a_i, b_i, choice);
        }
        Self(out)
    }
}

impl ConstantTimeEq for DoubleBitVector {
    fn ct_eq(&self, other: &Self) -> Choice {
        let mut out = Choice::from(1);
        for (a, b) in self.0.iter().zip(other.0.iter()) {
            out &= a.ct_eq(b);
        }
        out
    }
}

impl_op_ex!(^ |u: &DoubleBitVector, v: &DoubleBitVector| -> DoubleBitVector { u.xor(v) });
impl_op_ex!(^= |u: &mut DoubleBitVector, v: &DoubleBitVector| { u.xor_mut(v) });

/// Represents a matrix of bits.
///
/// Each row of this matrix is a `BitVector`, although we might have more or less
/// rows.
///
/// This is a fundamental object used for our OT extension protocol.
#[derive(Clone, Serialize, Deserialize)]
#[cfg_attr(test, derive(PartialEq, Eq))]
pub struct BitMatrix(Vec<BitVector>);

impl_secret_debug!(
    BitMatrix,
    |self| "BitMatrix(<redacted>, height={})",
    self.0.len()
);

impl BitMatrix {
    /// Create a random matrix of a certain chunk size.
    ///
    /// Each chunk will have a security parameter's worth of rows.
    pub fn random(rng: &mut impl CryptoRngCore, height: usize) -> Result<Self, ProtocolError> {
        if !height.is_multiple_of(SECURITY_PARAMETER) {
            return Err(ProtocolError::InvalidInput(format!(
                "height {height} must be a multiple of SECURITY_PARAMETER ({SECURITY_PARAMETER})"
            )));
        }
        Ok(Self((0..height).map(|_| BitVector::random(rng)).collect()))
    }

    /// Create a new matrix from a list of rows.
    pub fn from_rows<'a>(rows: impl IntoIterator<Item = &'a BitVector>) -> Self {
        Self(rows.into_iter().copied().collect())
    }

    /// Return the number of rows in this matrix.
    pub fn height(&self) -> usize {
        self.0.len()
    }

    /// Iterate over the rows of this matrix.
    pub fn rows(&self) -> impl Iterator<Item = &BitVector> {
        self.0.iter()
    }

    /// Iterate over a given column in chunks.
    pub fn column_chunks(&self, j: usize) -> impl Iterator<Item = BitVector> + '_ {
        self.0.chunks_exact(SECURITY_PARAMETER).map(move |chunk| {
            let mut out = BitVector::zero();
            for (word, word_chunk) in out.0.iter_mut().zip(chunk.chunks(64)) {
                for (bit_idx, c_i) in word_chunk.iter().enumerate() {
                    *word |= u64::from(c_i.bit(j)) << bit_idx;
                }
            }
            out
        })
    }

    /// Modify this matrix by xoring it with another.
    pub fn xor_mut(&mut self, other: &Self) {
        for (self_i, other_i) in self.0.iter_mut().zip(other.0.iter()) {
            *self_i ^= other_i;
        }
    }

    /// The result of xoring this matrix with another.
    pub fn xor(&self, other: &Self) -> Self {
        let mut out = self.clone();
        out.xor_mut(other);
        out
    }

    pub fn and_vec_mut(&mut self, v: &BitVector) {
        for self_i in &mut self.0 {
            *self_i &= v;
        }
    }

    pub fn and_vec(&self, v: &BitVector) -> Self {
        let mut out = self.clone();
        out.and_vec_mut(v);
        out
    }
}

impl FromIterator<BitVector> for BitMatrix {
    fn from_iter<T: IntoIterator<Item = BitVector>>(iter: T) -> Self {
        Self(iter.into_iter().collect())
    }
}

impl_op_ex!(^ |u: &BitMatrix, v: &BitMatrix| -> BitMatrix { u.xor(v) });
impl_op_ex!(^= |u: &mut BitMatrix, v: &BitMatrix| { u.xor_mut(v) });
impl_op_ex!(&|u: &BitMatrix, v: &BitVector| -> BitMatrix { u.and_vec(v) });
#[derive(Clone)]
#[cfg_attr(test, derive(PartialEq, Eq))]
pub struct SquareBitMatrix {
    pub matrix: BitMatrix,
}

impl_secret_debug!(SquareBitMatrix);

impl TryFrom<BitMatrix> for SquareBitMatrix {
    type Error = ();

    fn try_from(matrix: BitMatrix) -> Result<Self, Self::Error> {
        if matrix.height() != SECURITY_PARAMETER {
            return Err(());
        }
        Ok(Self { matrix })
    }
}

impl SquareBitMatrix {
    /// Expand transpose expands each row to contain `chunks * SECURITY_PARAMETER` bits, and then transposes
    /// the resulting matrix.
    pub fn expand_transpose(&self, sid: &[u8], rows: usize) -> Result<BitMatrix, ProtocolError> {
        if !rows.is_multiple_of(SECURITY_PARAMETER) {
            return Err(ProtocolError::InvalidInput(format!(
                "rows {rows} must be a multiple of SECURITY_PARAMETER ({SECURITY_PARAMETER})"
            )));
        }

        let mut hasher = Shake256::default();
        hasher.update(NEAR_PRG_CTX);
        hasher.update(b"sid");
        hasher.update(sid);

        let mut out = BitMatrix(vec![BitVector::zero(); rows]);

        // How many bytes to get rows bits?
        let row8 = rows.div_ceil(8);
        hasher.update(b"row");

        for (j, row) in self.matrix.0.iter().enumerate() {
            // We need to clone to make each row use the same prefix.
            let mut hasher_row = hasher.clone();
            for u in row.0 {
                hasher_row.update(&u.to_le_bytes());
            }
            let mut reader = hasher_row.finalize_xof();
            // Expand the row
            let mut expanded = vec![0u8; row8];
            reader.read(&mut expanded);

            // Now, write into the correct column
            let j_word = j / 64;
            let j_bit = j % 64;
            for (i, out_row) in out.0.iter_mut().enumerate() {
                let byte = expanded.get(i / 8).copied().unwrap_or(0);
                if let Some(word) = out_row.0.get_mut(j_word) {
                    *word |= u64::from((byte >> (i % 8)) & 1) << j_bit;
                }
            }
        }
        Ok(out)
    }
}

/// A choice vector holds an arbitrary number of choice bits.
///
/// This vector must always be non-empty.
#[derive(Clone)]
pub struct ChoiceVector(Vec<BitVector>);

impl_secret_debug!(ChoiceVector);

impl ChoiceVector {
    /// Generate a random vector with a certain number of bits.
    pub fn random(rng: &mut impl CryptoRngCore, size: usize) -> Result<Self, ProtocolError> {
        if size == 0 || !size.is_multiple_of(SECURITY_PARAMETER) {
            return Err(ProtocolError::InvalidInput(format!(
                "size {size} must be a positive multiple of SECURITY_PARAMETER ({SECURITY_PARAMETER})"
            )));
        }

        let data = (0..(size / SECURITY_PARAMETER))
            .map(|_| BitVector::random(rng))
            .collect();

        Ok(Self(data))
    }

    /// Iterate over the bits in this vector.
    pub fn bits(&self) -> impl Iterator<Item = Choice> + '_ {
        self.0.iter().flat_map(BitVector::bits)
    }

    /// Iterate over bitvector chunks from this vector.
    ///
    /// If the size of this vector is not evenly divided into chunks,
    /// then the last bitvector will be padded with 0s up until the MSB.
    pub fn chunks(&self) -> impl Iterator<Item = &BitVector> {
        self.0.iter()
    }
}

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

    #[test]
    fn test_gf_multiplication() {
        let a = BitVector([0b10, 0b10]);
        let b = BitVector([0b100, 0b100]);
        let c = DoubleBitVector([0b1000, 0, 0b1000, 0]);
        assert_eq!(a.gf_mul(&b), c);
    }
}