//! Module with the definition of the ServerKey.
//!
//! This module implements the generation of the server public key, together with all the
//! available homomorphic integer operations.
pub mod comparator;
mod crt;
mod crt_parallel;
mod radix;
mod radix_parallel;

use crate::integer::client_key::ClientKey;
use crate::shortint::server_key::MaxDegree;
use serde::{Deserialize, Serialize};

/// Error returned when the carry buffer is full.
pub use crate::shortint::CheckError;
pub use radix::scalar_mul::ScalarMultiplier;
pub use radix::scalar_sub::TwosComplementNegation;
pub use radix_parallel::{MiniUnsignedInteger, Reciprocable};

/// A structure containing the server public key.
///
/// The server key is generated by the client and is meant to be published: the client
/// sends it to the server so it can compute homomorphic integer circuits.
#[derive(Serialize, Deserialize, Clone)]
pub struct ServerKey {
    pub(crate) key: crate::shortint::ServerKey,
}

impl From<ServerKey> for crate::shortint::ServerKey {
    fn from(key: ServerKey) -> crate::shortint::ServerKey {
        key.key
    }
}

/// Compute the [`MaxDegree`] for an integer server key (compressed or uncompressed). This formula
/// provisions a free carry bit. This allows carry propagation between shortint blocks in a
/// [`RadixCiphertext`](`crate::integer::RadixCiphertext`), as that process requires adding a bit of
/// carry from one shortint block to the next, which would overflow and lead to wrong results if we
/// did not provision that carry bit.
fn integer_server_key_max_degree(parameters: crate::shortint::ShortintParameterSet) -> MaxDegree {
    MaxDegree((parameters.message_modulus().0 - 1) * parameters.carry_modulus().0 - 1)
}

impl ServerKey {
    /// Generates a server key.
    ///
    /// # Example
    ///
    /// ```rust
    /// use tfhe::integer::{ClientKey, ServerKey};
    /// use tfhe::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS_PBS;
    ///
    /// // Generate the client key:
    /// let cks = ClientKey::new(PARAM_MESSAGE_2_CARRY_2_KS_PBS);
    ///
    /// // Generate the server key:
    /// let sks = ServerKey::new(&cks);
    /// ```
    pub fn new<C>(cks: C) -> ServerKey
    where
        C: AsRef<ClientKey>,
    {
        // It should remain just enough space to add a carry
        let client_key = cks.as_ref();
        let max_degree = integer_server_key_max_degree(client_key.key.parameters);

        let sks = crate::shortint::server_key::ServerKey::new_with_max_degree(
            &client_key.key,
            max_degree,
        );

        ServerKey { key: sks }
    }

    /// Creates a ServerKey from an already generated shortint::ServerKey.
    ///
    /// # Example
    ///
    /// ```rust
    /// use tfhe::integer::{ClientKey, ServerKey};
    /// use tfhe::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS_PBS;
    ///
    /// let size = 4;
    ///
    /// // Generate the client key:
    /// let cks = ClientKey::new(PARAM_MESSAGE_2_CARRY_2_KS_PBS);
    ///
    /// // Generate the server key:
    /// let sks = ServerKey::new(&cks);
    /// ```
    pub fn from_shortint(
        cks: &ClientKey,
        mut key: crate::shortint::server_key::ServerKey,
    ) -> ServerKey {
        // It should remain just enough space add a carry
        let max_degree = integer_server_key_max_degree(cks.key.parameters);

        key.max_degree = max_degree;
        ServerKey { key }
    }

    pub fn deterministic_pbs_execution(&self) -> bool {
        self.key.deterministic_pbs_execution()
    }

    pub fn set_deterministic_pbs_execution(&mut self, new_deterministic_execution: bool) {
        self.key
            .set_deterministic_pbs_execution(new_deterministic_execution)
    }
}

#[derive(Clone, Serialize, Deserialize)]
pub struct CompressedServerKey {
    pub(crate) key: crate::shortint::CompressedServerKey,
}

impl CompressedServerKey {
    pub fn new(client_key: &ClientKey) -> CompressedServerKey {
        let max_degree = integer_server_key_max_degree(client_key.key.parameters);

        let key =
            crate::shortint::CompressedServerKey::new_with_max_degree(&client_key.key, max_degree);
        Self { key }
    }
}

impl From<CompressedServerKey> for ServerKey {
    fn from(compressed: CompressedServerKey) -> Self {
        let key = compressed.key.into();
        Self { key }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::integer::RadixClientKey;
    use crate::shortint::parameters::PARAM_MESSAGE_2_CARRY_2;

    /// https://github.com/zama-ai/tfhe-rs/issues/460
    /// Problem with CompressedServerKey degree being set to shortint MaxDegree not accounting for
    /// the necessary carry bits for e.g. Radix carry propagation.
    #[test]
    fn test_compressed_server_key_max_degree() {
        let cks = ClientKey::new(crate::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS_PBS);
        // msg_mod = 4, carry_mod = 4, (msg_mod - 1) * carry_mod = 12; minus 1 => 11
        let expected_max_degree = MaxDegree(11);

        let sks = ServerKey::new(&cks);
        assert_eq!(sks.key.max_degree, expected_max_degree);

        let csks = CompressedServerKey::new(&cks);
        assert_eq!(csks.key.max_degree, expected_max_degree);

        let decompressed_sks: ServerKey = csks.into();
        assert_eq!(decompressed_sks.key.max_degree, expected_max_degree);

        // Repro case from the user
        {
            let client_key = RadixClientKey::new(PARAM_MESSAGE_2_CARRY_2, 14);
            let compressed_eval_key = CompressedServerKey::new(client_key.as_ref());
            let evaluation_key = ServerKey::from(compressed_eval_key);
            let modulus = (client_key.parameters().message_modulus().0 as u128)
                .pow(client_key.num_blocks() as u32);

            let mut ct = client_key.encrypt(modulus - 1);
            let mut res_ct = ct.clone();
            for _ in 0..5 {
                res_ct = evaluation_key.smart_add_parallelized(&mut res_ct, &mut ct);
            }
            let res = client_key.decrypt::<u128>(&res_ct);
            assert_eq!(modulus - 6, res);
        }
    }
}
