#![cfg_attr(not(feature = "std"), no_std)]
extern crate alloc;

#[allow(clippy::many_single_char_names)]
mod helpers;

#[allow(clippy::many_single_char_names)]
mod offchain;

#[allow(clippy::many_single_char_names)]
mod shuffle;

#[allow(clippy::many_single_char_names)]
mod dkg;

#[allow(clippy::many_single_char_names)]
pub mod types;

mod bench;

#[cfg(test)]
mod mock;

#[cfg(test)]
#[macro_use]
mod tests;

pub mod keys;

use crate::dkg::{
    create::combine_shares,
    tally::combine_shares_and_tally_topic,
    verify::{
        verify_proof_and_store_decrypted_share, verify_proof_and_store_keygen_share,
    },
};
use crate::helpers::{
    assertions::{
        ensure_not_a_voting_authority, ensure_sealer, ensure_vote_does_not_exist,
        ensure_vote_exists, ensure_vote_phase, ensure_voting_authority,
    },
    ballot::store_ballot,
    phase::set_phase,
};
use crate::types::{
    Ballot, Cipher, DecryptedShare, DecryptedShareProof, NrOfShuffles,
    PublicKey as SubstratePK, PublicKeyShare, PublicParameters, ShufflePayload,
    ShuffleState, Title, Topic, TopicId, TopicResult, Vote, VoteId, VotePhase,
};
use frame_support::{
    debug, decl_error, decl_event, decl_module, decl_storage, dispatch::DispatchResult,
    storage::StorageMap, storage::StorageValue, traits::Get, weights::Pays,
};
use frame_system::{
    ensure_signed,
    offchain::{AppCrypto, CreateSignedTransaction},
};
use sp_std::{prelude::*, str, vec::Vec};

/// This is the pallet's configuration trait
pub trait Trait:
    pallet_timestamp::Trait + frame_system::Trait + CreateSignedTransaction<Call<Self>>
{
    /// The identifier type for an offchain worker.
    type AuthorityId: AppCrypto<Self::Public, Self::Signature>;
    /// The overarching dispatch call type.
    type Call: From<Call<Self>>;
    /// The overarching event type.
    type Event: From<Event<Self>> + Into<<Self as frame_system::Trait>::Event>;

    // Wait period between automated fetches. Set to 0 disable this feature.
    //   Then you need to manucally kickoff pricefetch
    type BlockDuration: Get<Self::BlockNumber>;
}

decl_storage! {
    trait Store for Module<T: Trait> as OffchainModule {
        pub VotingAuthorities get(fn voting_authorities) config(): Vec<T::AccountId>;
        pub Sealers get(fn sealers) config(): Vec<T::AccountId>;

        CountsBySealer get(fn count_by_sealer): map hasher(blake2_128_concat) T::AccountId => u32;
        Counts get(fn count): u32;

        /// A vector containing the IDs of voters that have submitted their ballots
        Voters get(fn voters): Vec<T::AccountId>;

        /// Set of all voteIds
        VoteIds get(fn vote_ids): Vec<VoteId>;

        /// Maps a vote (i.e. the voteId) to a due date
        Votes get(fn votes): map hasher(blake2_128_concat) VoteId => Vote<T::AccountId>;

        /// Maps a voteId to a topic (topicId, question)
        Topics get(fn topics): map hasher(blake2_128_concat) VoteId => Vec<Topic>;

        /// Maps an voter and a vote to a ballot. Used to verify if a voter has already voted.
        Ballots get(fn ballots): double_map hasher(blake2_128_concat) VoteId, hasher(blake2_128_concat) T::AccountId => Ballot;

        /// Maps a topicId (question) to a list of Ciphers and how many times each Cipher has been shuffled
        Ciphers get(fn ciphers): double_map hasher(blake2_128_concat) TopicId, hasher(blake2_128_concat) NrOfShuffles => Vec<Cipher>;

        /// Maps a voteId and topicId to a list of shuffle proofs (iteration, ciphers, proof)
        ShuffleProofs: map hasher(blake2_128_concat) (VoteId, TopicId) => Vec<ShufflePayload>;

        /// Maps a voteId and topicid to a shuffle status
        ShuffleStateStore: map hasher(blake2_128_concat) (VoteId, TopicId) => Option<ShuffleState>;

        /// Maps a topic to a map of results. [topic_id -> {message/vote: count}]
        Tally get(fn tally): map hasher(blake2_128_concat) TopicId => Option<TopicResult>;

        /// Maps a sealer and a topic to a vector of decrypted shares.
        DecryptedShares get(fn decrypted_shares): double_map hasher(blake2_128_concat) TopicId, hasher(blake2_128_concat) T::AccountId  => Vec<Vec<u8>>;

        /// Stores the public key of a sealer together with its Schnorr proof.
        PublicKeyShares get(fn key_shares): map hasher(blake2_128_concat) VoteId => Vec<PublicKeyShare>;

        /// Stores the public key of a sealer, indexed by sealer account
        PublicKeyShareBySealer get(fn key_share_by_sealer): map hasher(blake2_128_concat) (VoteId, T::AccountId) => Option<PublicKeyShare>;

        /// Maps a vote to a public key (the vote's/system's public key) used to encrypt ballots.
        PublicKey get(fn public_key): map hasher(blake2_128_concat) VoteId => Option<SubstratePK>;
    }
}

decl_event!(
    /// Events generated by the module.
    pub enum Event<T>
    where
        AccountId = <T as frame_system::Trait>::AccountId,
    {
        /// ballot submission event -> [from/who, ballot]
        BallotSubmitted(AccountId, VoteId, Ballot),

        /// public key stored event -> [from/who, public key]
        PublicKeyStored(AccountId, VoteId, SubstratePK),

        /// A voting authority set the vote's public parameters. [vote, who, params]
        VoteCreatedWithPublicParameters(VoteId, AccountId, PublicParameters),

        /// A voting authority set the question of a topic of a vote [vote, (topic_id, question)]
        VoteTopicQuestionStored(VoteId, Topic),

        /// A voting authority changed the vote phase [vote_id, newPhase]
        VotePhaseChanged(VoteId, VotePhase),

        /// A public key share was submitted. [public key with its proof]
        PublicKeyShareSubmitted(PublicKeyShare),

        /// A system public key has been created. [vote_id, public_key]
        PublicKeyCreated(VoteId, SubstratePK),

        /// A decrypted share was submitted for a vote. [paritial decryptions with its proof]
        DecryptedShareSubmitted(TopicId, AccountId),

        /// A decrypted share was submitted for a vote. [paritial decryptions with its proof]
        TopicTallied(VoteId, TopicId, TopicResult),

        /// A decrypted share was submitted for a vote. [paritial decryptions with its proof]
        ShuffleProofSubmitted(TopicId, AccountId),
    }
);

decl_error! {
    pub enum Error for Module<T: Trait> {
        // Error returned when not sure which off-chain worker function to executed
        UnknownOffchainMux,

        // Error returned when Vec<u8> cannot be parsed into BigUint
        ParseError,

        // Error returned when requester is not a voting authority
        NotAVotingAuthority,

        // Error returned when requester is a voting authority
        IsVotingAuthority,

        // Error returned when requester is not a sealer
        NotASealer,

        // Error returned when making signed transactions in off-chain worker
        NoLocalAcctForSigning,
        OffchainSignedTxError,

        // Error returned when failing to get randomness
        RandomnessGenerationError,

        // Error returned when upper bound is zero
        RandomnessUpperBoundZeroError,

        // Error returned when error occurs in gen_random_range
        RandomRangeError,

        // Error returned when permutation size is zero
        PermutationSizeZeroError,

        // Error returned when public key doesn't exist
        PublicKeyNotExistsError,

        // Error returned when public key share doesn't exist
        PublicKeyShareNotExistsError,

        // Error returned when the public key share proof doesn't verify
        PublicKeyShareProofError,

        // Error returned when there are less than two public key shares
        NotEnoughPublicKeyShares,

        // Error returned when inverse modulo operation fails
        InvModError,

        // Error returned when division modulo operation fails
        DivModError,

        // Error returned when vote_id does not exist yet
        VoteDoesNotExist,

        // Error returned when vote_id already exists
        VoteAlreadyExists,

        // Error returned when vote is in wrong phase
        WrongVotePhase,

        // Error returned when the decrypted share proof doesn't verify
        DecryptedShareProofError,

        // Error returned when not all sealers have submitted their decrypted shares yet
        NotEnoughDecryptedShares,

        // Error returned when a topic has already been tallied and a second attempt to tally the votes is made
        TopicHasAlreadyBeenTallied,

        // Error returned when a shuffle proof verification fails
        ShuffleProofVerifcationFailed,

        // Error returned when a shuffle has already been performed for a set of ciphers
        ShuffleAlreadyPerformed,

        // Error returned when ballots are empty when trying to shuffle them
        ShuffleCiphersSizeZeroError,

        // Error returned when no encryptions exists for given nr_of_shuffles
        NrOfShufflesDoesNotExist,

        /// Error returned in offchain worker when computing shuffle start position fails
        CouldNotComputeShuffleStartPosition,

        /// Error returned when shuffle state information mismatch occurrs in extrinsic
        ShuffleStateIncorrect,

        /// Error returned when shuffle is submitted for (vote_id, topic_id) which is already completed
        ShuffleAlreadyCompleted
    }
}

decl_module! {
    pub struct Module<T: Trait> for enum Call where origin: T::Origin {
        // Errors must be initialized if they are used by the pallet.
        type Error = Error<T>;

        // Events must be initialized if they are used by the pallet.
        fn deposit_event() = default;

        /// Set a vote phase.
        #[weight = (10_000, Pays::No)]
        fn set_vote_phase(origin, vote_id: VoteId, phase: VotePhase) -> DispatchResult {
            // only the voting_authority should be able to store the key
            let who: T::AccountId = ensure_signed(origin)?;
            ensure_voting_authority::<T>(&who)?;

            // check that the vote_id exists
            ensure_vote_exists::<T>(&vote_id)?;

            // set the new phase
            let mut vote: Vote<T::AccountId> = Votes::<T>::get(&vote_id);
            vote.phase = phase.clone();
            Votes::<T>::insert(&vote_id, &vote);
            set_phase::<T>(&who, &vote_id, phase.clone())?;

            // notify that the vote phase has been changed
            debug::info!("updated vote phase: {:?}, {:?}", vote_id, phase);
            Self::deposit_event(RawEvent::VotePhaseChanged(vote_id, phase));
            Ok(())
        }

        /// DEV ONLY
        /// NEEDS TO BE DISABLE IN PRODUCTION
        #[weight = (10000, Pays::No)]
        pub fn store_public_key(origin, vote_id: VoteId, pk: SubstratePK) -> DispatchResult {
            // only the voting_authority should be able to store the key
            let who: T::AccountId = ensure_signed(origin)?;
            ensure_voting_authority::<T>(&who)?;

            // store the public key
            PublicKey::insert(vote_id.clone(), pk.clone());

            // notify that the public key has been stored
            debug::info!("stored public key for vote: {:?}", vote_id);
            Self::deposit_event(RawEvent::PublicKeyStored(who, vote_id, pk));
            Ok(())
        }

        /// Store a public key and its proof.
        /// Can only be called from a sealer.
        #[weight = (10_000, Pays::No)]
        fn store_public_key_share(origin, vote_id: VoteId, pk_share: PublicKeyShare) -> DispatchResult {
            // only sealers can store their public key shares
            let who: T::AccountId = ensure_signed(origin)?;
            ensure_not_a_voting_authority::<T>(&who)?;
            ensure_sealer::<T>(&who)?;

            // verify key generatin proof
            // and store public key share
            verify_proof_and_store_keygen_share::<T>(who.clone(), &vote_id, pk_share.clone())?;

            debug::info!("stored public key share for vote: {:?} (by sealer: {:?})", vote_id, who.clone());
            Self::deposit_event(RawEvent::PublicKeyShareSubmitted(pk_share));
            Ok(())
        }

        /// Combine public key shares into a single public key.
        #[weight = (10_000, Pays::No)]
        fn combine_public_key_shares(origin, vote_id: VoteId) -> DispatchResult {
            // only the voting_authority should be able to combine the public key shares
            let who: T::AccountId = ensure_signed(origin)?;
            ensure_voting_authority::<T>(&who)?;
            ensure_vote_exists::<T>(&vote_id)?;

            // create the system's public key
            let pk: SubstratePK = combine_shares::<T>(who, &vote_id)?;

            debug::info!("combined public key shares for vote: {:?}", vote_id);
            Self::deposit_event(RawEvent::PublicKeyCreated(vote_id, pk));
            Ok(())
        }

        /// Create a vote and store public crypto parameters.
        /// Can only be called from a voting authority.
        #[weight = (10000, Pays::No)]
        fn create_vote(origin, vote_id: VoteId, title: Title, params: PublicParameters, topics: Vec<Topic>, batch_size: u64) -> DispatchResult {
            let who: T::AccountId = ensure_signed(origin)?;
            ensure_voting_authority::<T>(&who)?;

            // create new vote
            let vote = Vote::<T::AccountId> {
                voting_authority: who.clone(),
                title,
                phase: VotePhase::default(),
                params: params.clone()
            };

            // store the vote_id, vote + topic information
            let mut vote_ids: Vec<VoteId> = VoteIds::get();
            ensure_vote_does_not_exist::<T>(&vote_id)?;

            vote_ids.push(vote_id.clone());
            VoteIds::put(vote_ids);
            Votes::<T>::insert(&vote_id, vote);

            // create an empty shuffle state for each topic
            for topic in topics.iter() {
                let (topic_id, _) = topic;
                ShuffleStateStore::insert((&vote_id, &topic_id), ShuffleState {
                    iteration: 0,
                    start_position: 0,
                    batch_size,
                    done: false
                });
            }

            // store all topics (topic_id, question)
            Topics::insert(&vote_id, topics);

            // log success + emit event
            debug::info!("created vote: {:?}", vote_id);
            Self::deposit_event(RawEvent::VoteCreatedWithPublicParameters(vote_id, who, params));
            Ok(())
        }

        /// Add a question to the vote.
        /// Can only be called from a voting authority.
        #[weight = (10000, Pays::No)]
        fn store_question(origin, vote_id: VoteId, topic: Topic, batch_size: u64) -> DispatchResult {
            let who = ensure_signed(origin)?;
            ensure_voting_authority::<T>(&who)?;
            ensure_vote_exists::<T>(&vote_id)?;

            let topic_id = &topic.0;
            let mut topics: Vec<Topic> = Topics::get(&vote_id);
            topics.push(topic.clone());

            // create an empty shuffle state for the topic
            ShuffleStateStore::insert((&vote_id, topic_id), ShuffleState {
                iteration: 0,
                start_position: 0,
                batch_size,
                done: false,
            });

            // store the topic
            Topics::insert(&vote_id, topics);

            debug::info!("added question: {:?} to vote: {:?}", topic, vote_id);
            Self::deposit_event(RawEvent::VoteTopicQuestionStored(vote_id, topic));
            Ok(())
        }

        #[weight = (10000, Pays::No)]
        pub fn cast_ballot(origin, vote_id: VoteId, ballot: Ballot) -> DispatchResult {
          let who = ensure_signed(origin)?;
          ensure_vote_exists::<T>(&vote_id)?;
          ensure_vote_phase::<T>(&vote_id, VotePhase::Voting)?;

          // TODO: ensure that it is a legit voter -> in some other project where identity management is considered

          // store the ballot
          store_ballot::<T>(&who, &vote_id, ballot.clone());

          // notify that the ballot has been submitted and stored
          debug::info!("stored ballot for vote_id: {:?}", vote_id);
          Self::deposit_event(RawEvent::BallotSubmitted(who, vote_id, ballot));
          Ok(())
        }

        /// Test function to check signer.
        #[weight = (10_000, Pays::No)]
        fn submit_shuffled_votes_and_proof(origin, vote_id: VoteId, topic_id: TopicId, payload: ShufflePayload) -> DispatchResult {
            let who: T::AccountId = ensure_signed(origin)?;
            ensure_sealer::<T>(&who)?;
            ensure_vote_exists::<T>(&vote_id)?;

            // TODO: discuss if shuffling should be allowed earlier
            ensure_vote_phase::<T>(&vote_id, VotePhase::Tallying)?;

            Self::verify_proof_store_shuffled_ciphers(&vote_id, &topic_id, payload)?;

            // notify that the decrypted share has been:
            // submitted, the proof verified and stored
            debug::info!("verified shuffle proof for vote_id: {:?}, topic_id: {:?}", vote_id, topic_id);
            Self::deposit_event(RawEvent::ShuffleProofSubmitted(topic_id, who));
            Ok(())
        }

        /// Store a decrypted shares.
        #[weight = (10_000, Pays::No)]
        fn submit_decrypted_shares(origin, vote_id: VoteId, topic_id: TopicId, shares: Vec<DecryptedShare>, proof: DecryptedShareProof, nr_of_shuffles: NrOfShuffles) -> DispatchResult {
            // only sealers should be able to store their decrypted shares
            let who: T::AccountId = ensure_signed(origin)?;
            ensure_vote_exists::<T>(&vote_id)?;
            ensure_vote_phase::<T>(&vote_id, VotePhase::Tallying)?;
            ensure_sealer::<T>(&who)?;

            // verify the decrypted share proof
            // and store the decrypted shares if proof verification is successfull
            verify_proof_and_store_decrypted_share::<T>(who.clone(), &vote_id, &topic_id, shares, proof.clone(), &nr_of_shuffles)?;

            // notify that the decrypted share has been:
            // submitted, the proof verified and stored
            debug::info!("stored decrypted share for vote: {:?} and topic: {:?}, by sealer: {:?}", vote_id, topic_id, who.clone());
            Self::deposit_event(RawEvent::DecryptedShareSubmitted(topic_id, who));
            Ok(())
        }

        /// Combine decrypted shares into a final plain text tally.
        #[weight = (10_000, Pays::No)]
        fn combine_decrypted_shares(origin, vote_id: VoteId, topic_id: TopicId, encoded: bool, nr_of_shuffles: NrOfShuffles) -> DispatchResult {
            // only the voting_authority should be able to create the final tally
            let who: T::AccountId = ensure_signed(origin)?;
            ensure_vote_exists::<T>(&vote_id)?;
            ensure_vote_phase::<T>(&vote_id, VotePhase::Tallying)?;
            ensure_voting_authority::<T>(&who)?;

            // combine the decrypted shares
            // tally the topic
            let result: TopicResult = combine_shares_and_tally_topic::<T>(&vote_id, &topic_id, encoded, &nr_of_shuffles)?;

            // notify that the decrypted shares have been combined
            // and that the result has been tallied!
            debug::info!("result for vote: {:?} and topic: {:?} is: {:?}", vote_id, topic_id, result);
            Self::deposit_event(RawEvent::TopicTallied(vote_id, topic_id, result));
            Ok(())
        }

        /// Empty function that does nothing but needs to be called by an offchain worker
        /// when it's not the offchain worker's turn to shuffle the votes.
        #[weight = (10_000, Pays::No)]
        fn do_nothing_when_its_not_your_turn(origin) -> DispatchResult {
            let who: T::AccountId = ensure_signed(origin)?;
            debug::info!("offchain fn call when not shuffling, who: {:?}", who);
            Ok(())
        }

        fn offchain_worker(block_number: T::BlockNumber) {
            debug::info!("off-chain worker: entering...");

            let offchain_shuffle_result = Self::offchain_shuffling(block_number);
            match offchain_shuffle_result {
                Ok(_) => (),
                Err(err) => debug::error!("error while shuffling in offchain worker: {:?}", err),
             }

            debug::info!("off-chain worker: done...");
        }
    }
}

impl<T: Trait> sp_runtime::offchain::storage_lock::BlockNumberProvider for Module<T> {
    type BlockNumber = T::BlockNumber;
    fn current_block_number() -> Self::BlockNumber {
        <frame_system::Module<T>>::block_number()
    }
}
