{-# LANGUAGE OverloadedStrings #-}
module Voting.Protocol.Trustees.All where

import Control.Monad (Monad(..), foldM, unless)
import Control.Monad.Trans.Except (ExceptT(..), throwE)
import Data.Eq (Eq(..))
import Data.Function (($))
import Data.Maybe (maybe)
import Data.Semigroup (Semigroup(..))
import Text.Show (Show(..))
import qualified Control.Monad.Trans.State.Strict as S
import qualified Data.ByteString as BS
import qualified Data.List as List

import Voting.Protocol.Utils
import Voting.Protocol.Arithmetic
import Voting.Protocol.Credential
import Voting.Protocol.Election

-- * Type 'TrusteePublicKey'
data TrusteePublicKey q = TrusteePublicKey
 { trustee_PublicKey      :: PublicKey q
 , trustee_SecretKeyProof :: Proof q
 } deriving (Eq,Show)

randomSecretKey :: Monad m => RandomGen r => SubGroup q => S.StateT r m (SecretKey q)
randomSecretKey = random

-- ** Type 'ErrorTrusteePublicKey'
data ErrorTrusteePublicKey
 =   ErrorTrusteePublicKey_Wrong
     -- ^ The 'trustee_SecretKeyProof' is wrong.
 deriving (Eq,Show)

-- BELENIOS: MakeSimple.prove
-- | @('proveTrusteePublicKey' trustSecKey)@
-- returns the 'PublicKey' associated to 'trustSecKey'
-- and a 'Proof' of its knowledge.
proveTrusteePublicKey ::
 Monad m => RandomGen r => SubGroup q =>
 SecretKey q -> S.StateT r m (TrusteePublicKey q)
proveTrusteePublicKey trustSecKey = do
        let trustee_PublicKey = publicKey trustSecKey
        trustee_SecretKeyProof <-
                prove trustSecKey [groupGen] $
                        hash (trusteePublicKeyStatement trustee_PublicKey)
        return TrusteePublicKey{..}

-- BELENIOS: MakeSimple.check
-- | @('verifyTrusteePublicKey' trustPubKey)@
-- returns 'True' iif. the given 'trustee_SecretKeyProof'
-- does prove that the 'SecretKey' associated with
-- the given 'trustee_PublicKey' is known by the trustee.
verifyTrusteePublicKey ::
 Monad m => SubGroup q =>
 TrusteePublicKey q ->
 ExceptT ErrorTrusteePublicKey m ()
verifyTrusteePublicKey TrusteePublicKey{..} =
        unless ((proof_challenge trustee_SecretKeyProof ==) $
        hash
         (trusteePublicKeyStatement trustee_PublicKey)
         [commit trustee_SecretKeyProof groupGen trustee_PublicKey]) $
                throwE ErrorTrusteePublicKey_Wrong

-- ** Hashing
trusteePublicKeyStatement :: PublicKey q -> BS.ByteString
trusteePublicKeyStatement trustPubKey = "pok|"<>bytesNat trustPubKey<>"|"

-- BELENIOS: MakeSimple.combine
electionPublicKey :: SubGroup q => [TrusteePublicKey q] -> PublicKey q
electionPublicKey = List.foldr (\TrusteePublicKey{..} -> (trustee_PublicKey *)) one

-- BELENIOS: combine_factors
-- | @('combineDecryptionShares' checker pubKeyByTrustee decShareByTrustee)@
-- returns the 'DecryptionFactor's by voter by question
combineDecryptionShares ::
 SubGroup q =>
 [[Encryption q]] ->
 [PublicKey q] -> DecryptionShareCombinator q
combineDecryptionShares encByQuestByBallot pubKeyByTrustee decShareByTrustee = do
        isoZipWithM_ (throwE ErrorDecryptionShare_Invalid)
         (verifyDecryptionShare encByQuestByBallot)
         pubKeyByTrustee
         decShareByTrustee
        (d0,ds) <- maybe err return $ List.uncons decShareByTrustee
        foldM
         (\decFactorByQuestByBallot DecryptionShare{..} ->
                isoZipWithM err
                 (\acc df -> maybe err return $ isoZipWith (*) acc df)
                 decFactorByQuestByBallot decryptionShare_factors)
         (decryptionShare_factors d0) ds
        where err = throwE ErrorDecryptionShare_Invalid