use crate::bigint::BigUintWrapper;
use crate::crypto::{Rsa, RsaParameter};
use num_bigint_dig::BigUint;
use rand::prelude::SliceRandom;
use rand::{CryptoRng, RngCore};
use serde::{Deserialize, Serialize};
use std::convert::TryInto;
use uuid::Uuid;

#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
/// The [Deck] struct is the starting point to shuffle. It contains n cards. The deck doesn't care
/// what cards it represents.
///
/// # Example
///
/// ```
/// # use distributed_cards::{Deck, Prime, RsaParameter};
/// # use uuid::Uuid;
/// # use rand::thread_rng;
/// let rng = &mut thread_rng();
/// let primes = [Prime::random(128, rng), Prime::random(128, rng)]; // each participant contributes one prime
/// let rsa_parameter = RsaParameter::from_primes(&primes);
/// let uuid = Uuid::from_u128(0x_1905709b_e2ae_469c_9589_4e37dcf3e5bc); // must be unique and unpredictable for each shuffle
/// let deck = Deck::new(104, uuid);
/// let (encrypted_deck, key) = deck.shuffle_encrypt(&rsa_parameter, rng);
/// assert!(deck.is_shuffle_encrypt_valid(&key, &encrypted_deck));
/// ```
pub struct Deck {
    num_cards: u32,
    shuffle_id: Uuid,
}

impl Deck {
    /// Creates a new deck with n cards numbered 0 up to n-1. The given shuffle_id must be unique
    /// over all shuffles, but both parameters have to be the same for each participant
    pub fn new(num_cards: u32, shuffle_id: Uuid) -> Deck {
        Deck {
            num_cards,
            shuffle_id,
        }
    }

    // output format of card: [1] [4 byte card_id] [16 byte shuffle_id]
    // the number of bits required should always be exactly 161
    fn encode_card(card: u32, shuffle_id: Uuid) -> BigUint {
        (BigUint::from(1u32) << 160)
            + (BigUint::from(card) << 128)
            + BigUint::from(shuffle_id.as_u128())
    }

    fn decode_card(card: BigUint) -> Option<(u32, Uuid)> {
        if card.bits() != 161 {
            return None;
        }
        let card = card.to_bytes_be();
        let id = u32::from_be_bytes(card[1..5].try_into().unwrap());
        let uuid = Uuid::from_bytes(card[5..].try_into().unwrap());
        Some((id, uuid))
    }

    fn to_encrypted_deck(&self) -> EncryptedDeck {
        let cards = (0..self.num_cards)
            .into_iter()
            .map(|i| BigUintWrapper(Deck::encode_card(i, self.shuffle_id)))
            .collect();
        EncryptedDeck { cards }
    }
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct EncryptedDeck {
    cards: Vec<BigUintWrapper>,
}

/// 1st shuffling phase
impl Deck {
    pub fn shuffle_encrypt<Rng: CryptoRng + RngCore>(
        &self,
        rsa_parameter: &RsaParameter,
        rng: &mut Rng,
    ) -> (EncryptedDeck, Rsa) {
        let deck: EncryptedDeck = self.to_encrypted_deck();
        deck.shuffle_encrypt(rsa_parameter, rng)
    }

    pub fn is_shuffle_encrypt_valid(&self, key: &Rsa, encrypted_deck: &EncryptedDeck) -> bool {
        let deck: EncryptedDeck = self.to_encrypted_deck();
        deck.is_shuffle_encrypt_valid(key, encrypted_deck)
    }
}

impl EncryptedDeck {
    fn encrypt(&self, key: &Rsa) -> EncryptedDeck {
        let cards = self
            .cards
            .iter()
            .map(|c| BigUintWrapper(key.encrypt(c.0.clone())))
            .collect();
        EncryptedDeck { cards }
    }

    fn decrypt(&self, key: &Rsa) -> EncryptedDeck {
        let cards = self
            .cards
            .iter()
            .map(|c| BigUintWrapper(key.decrypt(c.0.clone())))
            .collect();
        EncryptedDeck { cards }
    }
}

/// 1st shuffling phase
impl EncryptedDeck {
    pub fn shuffle_encrypt<Rng: CryptoRng + RngCore>(
        &self,
        rsa_parameter: &RsaParameter,
        rng: &mut Rng,
    ) -> (EncryptedDeck, Rsa) {
        let key = Rsa::gen_with_parameter(rsa_parameter.clone(), rng);
        let mut deck = self.encrypt(&key);
        deck.cards.shuffle(rng);
        (deck, key)
    }

    /// checks if the shuffle_encrypt mechanism was executed correctly and it can be reversed again
    pub fn is_shuffle_encrypt_valid(&self, key: &Rsa, result: &EncryptedDeck) -> bool {
        let mut decrypted_result = result.decrypt(key);
        decrypted_result.cards.sort();
        let mut cards = self.cards.clone();
        cards.sort();
        cards == decrypted_result.cards
    }
}

// 2nd shuffling phase
impl EncryptedDeck {
    fn card_specific(&self, shuffle_key: &Rsa, keys: &[Rsa]) -> EncryptedDeck {
        let cards = self
            .cards
            .iter()
            .zip(keys)
            .map(|(card, key)| BigUintWrapper(shuffle_key.decrypt(key.encrypt(card.0.clone()))))
            .collect();
        EncryptedDeck { cards }
    }

    /// removes shuffling key from cards and encrypts each card with a unique random key
    /// returns the deck and all card specific keys
    pub fn encrypt_card_specific<Rng: CryptoRng + RngCore>(
        &self,
        shuffle_key: &Rsa,
        rsa_parameter: &RsaParameter,
        rng: &mut Rng,
    ) -> (EncryptedDeck, Vec<Rsa>) {
        let keys: Vec<Rsa> = (0..self.cards.len())
            .into_iter()
            .map(|_| Rsa::gen_with_parameter(rsa_parameter.clone(), rng))
            .collect();
        let deck = self.card_specific(shuffle_key, &keys);
        (deck, keys)
    }
    /// checks whether the encrypt_card_specific function was executed correctly
    /// needs to know the keys involved, so can only be called for peers after the game has ended
    /// and everyone published all their keys
    pub fn is_encrypt_card_specific_valid(
        &self,
        shuffle_key: &Rsa,
        result: &EncryptedDeck,
        keys: &[Rsa],
    ) -> bool {
        &self.card_specific(shuffle_key, keys) == result
    }
}

impl EncryptedDeck {
    /// Returns None on error
    ///
    /// # Panics
    ///
    /// If card_id >= num_cards
    pub fn decrypt_card(
        &self,
        card_id: u32,
        shuffle_id: Uuid,
        own_key: &Rsa,
        other_keys: &[Rsa],
    ) -> Option<u32> {
        let decrypted = other_keys.iter().fold(
            own_key.decrypt(self.cards[card_id as usize].0.clone()),
            |card, key| key.decrypt(card),
        );
        let (card, uuid) = Deck::decode_card(decrypted)?;
        if uuid == shuffle_id {
            Some(card)
        } else {
            None
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::crypto::Prime;
    use rand::thread_rng;

    #[test]
    fn encode_card() {
        let expected = [
            1u8, 0, 0, 0, 14, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
        ];
        assert_eq!(
            Deck::encode_card(14, Uuid::from_slice(&expected[5..]).unwrap()).to_bytes_be(),
            expected
        )
    }

    #[test]
    fn decode_card() {
        let expected = [
            1u8, 0, 0, 0, 14, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
        ];
        let card = BigUint::from_bytes_be(&expected);
        assert_eq!(
            Deck::decode_card(card),
            Some((14, Uuid::from_bytes(expected[5..].try_into().unwrap())))
        );
    }

    #[test]
    fn decode_card_err() {
        let expected = [
            1u8, 0, 0, 0, 14, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
        ];
        assert_eq!(
            Deck::decode_card(BigUint::from_bytes_be(&expected[1..])),
            None
        );
        assert_eq!(
            Deck::decode_card(BigUint::from_bytes_be(&expected) << 1),
            None
        );
    }

    #[test]
    fn shuffle_encrypt() {
        let rng = &mut thread_rng();
        let primes = [Prime::random(128, rng), Prime::random(128, rng)];
        let rsa_parameter = RsaParameter::from_primes(&primes);
        let uuid = Uuid::from_u128(0x_1905709b_e2ae_469c_9589_4e37dcf3e5bc);
        let deck = Deck::new(4, uuid);
        let (encrypted_deck, key) = deck.shuffle_encrypt(&rsa_parameter, rng);
        assert!(deck.is_shuffle_encrypt_valid(&key, &encrypted_deck));
    }

    #[test]
    fn encrypt_card_specific() {
        let rng = &mut thread_rng();
        let primes = [Prime::random(128, rng), Prime::random(128, rng)];
        let rsa_parameter = RsaParameter::from_primes(&primes);
        let uuid = Uuid::from_u128(0x_1905709b_e2ae_469c_9589_4e37dcf3e5bc);
        let deck = Deck::new(4, uuid);
        let (encrypted_deck, key) = deck.shuffle_encrypt(&rsa_parameter, rng);
        assert!(deck.is_shuffle_encrypt_valid(&key, &encrypted_deck));

        let (phase_2, phase_2_keys) =
            encrypted_deck.encrypt_card_specific(&key, &rsa_parameter, rng);
        assert!(encrypted_deck.is_encrypt_card_specific_valid(&key, &phase_2, &phase_2_keys));
    }

    #[test]
    fn decode() {
        let rng = &mut thread_rng();
        let primes = [Prime::random(128, rng), Prime::random(128, rng)];
        let rsa_parameter = RsaParameter::from_primes(&primes);
        let uuid = Uuid::from_u128(0x_1905709b_e2ae_469c_9589_4e37dcf3e5bc);
        let deck = Deck::new(16, uuid);
        let (encrypted_deck, key) = deck.shuffle_encrypt(&rsa_parameter, rng);
        assert!(deck.is_shuffle_encrypt_valid(&key, &encrypted_deck));

        let (phase_2, phase_2_keys) =
            encrypted_deck.encrypt_card_specific(&key, &rsa_parameter, rng);
        assert!(encrypted_deck.is_encrypt_card_specific_valid(&key, &phase_2, &phase_2_keys));

        let mut decrypted: Vec<u32> = phase_2_keys
            .iter()
            .enumerate()
            .map(|(i, key)| {
                phase_2
                    .decrypt_card(i as u32, uuid.clone(), key, &[])
                    .unwrap()
            })
            .collect();
        decrypted.sort();
        assert_eq!(decrypted, (0..16).collect::<Vec<_>>());
    }

    #[test]
    fn encrypt_card_2() {
        let rng = &mut thread_rng();
        let primes = [Prime::random(128, rng), Prime::random(128, rng)];
        let rsa_parameter = RsaParameter::from_primes(&primes);
        let shuffle_id = Uuid::from_u128(0x_1905709b_e2ae_469c_9589_4e37dcf3e5bc);
        let deck = Deck::new(1, shuffle_id);

        // shuffle phase 1
        let (a_phase_1, a_shuffle_key) = deck.shuffle_encrypt(&rsa_parameter, rng);
        assert!(deck.is_shuffle_encrypt_valid(&a_shuffle_key, &a_phase_1));
        let (b_phase_1, b_shuffle_key) = a_phase_1.shuffle_encrypt(&rsa_parameter, rng);
        assert!(a_phase_1.is_shuffle_encrypt_valid(&b_shuffle_key, &b_phase_1));

        // shuffle phase 2
        let (a_phase_2, a_phase_2_keys) =
            b_phase_1.encrypt_card_specific(&a_shuffle_key, &rsa_parameter, rng);
        assert!(b_phase_1.is_encrypt_card_specific_valid(
            &a_shuffle_key,
            &a_phase_2,
            &a_phase_2_keys
        ));
        let (b_phase_2, b_phase_2_keys) =
            a_phase_2.encrypt_card_specific(&b_shuffle_key, &rsa_parameter, rng);
        assert!(a_phase_2.is_encrypt_card_specific_valid(
            &b_shuffle_key,
            &b_phase_2,
            &b_phase_2_keys
        ));

        // decrypt ingame
        assert_eq!(
            b_phase_2.decrypt_card(
                0,
                shuffle_id,
                &a_phase_2_keys[0],
                &[b_phase_2_keys[0].clone()]
            ),
            Some(0)
        );
        assert_eq!(
            b_phase_2.decrypt_card(
                0,
                shuffle_id,
                &b_phase_2_keys[0],
                &[a_phase_2_keys[0].clone()]
            ),
            Some(0)
        );
    }

    #[test]
    fn serialize_encrypted() {
        let encrypted_deck = EncryptedDeck {
            cards: vec![BigUintWrapper(BigUint::from(
                0x_1905709b_e2ae_469c_9589_4e37dcf3e5bc_u128,
            ))],
        };
        assert_eq!(
            serde_json::to_string(&encrypted_deck).unwrap(),
            r#"{"cards":["GQVwm+KuRpyViU433PPlvA=="]}"#.to_owned()
        );
    }
}