{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE UndecidableInstances #-} -- for Reifies instances
module Voting.Protocol.Tally where

import Control.DeepSeq (NFData)
import Control.Monad (Monad(..), mapM, unless)
import Control.Monad.Trans.Except (Except, ExceptT, throwE)
import Data.Aeson (ToJSON(..),FromJSON(..))
import Data.Eq (Eq(..))
import Data.Function (($))
import Data.Functor ((<$>))
import Data.Maybe (maybe)
import Data.Semigroup (Semigroup(..))
import Data.Tuple (fst)
import GHC.Generics (Generic)
import Numeric.Natural (Natural)
import Prelude (fromIntegral)
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 qualified Data.Map.Strict as Map

import Voting.Protocol.Utils
import Voting.Protocol.FFC
import Voting.Protocol.Credential
import Voting.Protocol.Election

-- * Type 'Tally'
data Tally c = Tally
 { tally_countMax :: !Natural
   -- ^ The maximal number of supportive 'Opinion's that a choice can get,
   -- which is here the same as the number of 'Ballot's.
   --
   -- Used in 'proveTally' to decrypt the actual
   -- count of votes obtained by a choice,
   -- by precomputing all powers of 'groupGen's up to it.
 , tally_encByChoiceByQuest :: !(EncryptedTally c)
   -- ^ 'Encryption' by 'Question' by 'Ballot'.
 , tally_decShareByTrustee :: ![DecryptionShare c]
   -- ^ 'DecryptionShare' by trustee.
 , tally_countByChoiceByQuest :: ![[Natural]]
   -- ^ The decrypted count of supportive 'Opinion's, by choice by 'Question'.
 } deriving (Eq,Show,Generic,NFData)
deriving instance Reifies c FFC => ToJSON (Tally c)
deriving instance Reifies c FFC => FromJSON (Tally c)

-- ** Type 'EncryptedTally'
-- | 'Encryption' by choice by 'Question'.
type EncryptedTally c = [[Encryption c]]

-- | @('encryptedTally' ballots)@
-- returns the sum of the 'Encryption's of the given @ballots@,
-- along with the number of 'Ballot's.
encryptedTally :: Reifies c FFC => [Ballot c] -> (EncryptedTally c, Natural)
encryptedTally ballots =
        ( List.foldr (\Ballot{..} ->
                List.zipWith (\Answer{..} ->
                        List.zipWith (+)
                         (fst <$> answer_opinions))
                 ballot_answers)
         (List.repeat (List.repeat zero))
         ballots
        , fromIntegral $ List.length ballots
        )

-- ** Type 'DecryptionShareCombinator'
type DecryptionShareCombinator c =
        EncryptedTally c -> [DecryptionShare c] -> Except ErrorTally [[DecryptionFactor c]]

proveTally ::
 Reifies c FFC =>
 (EncryptedTally c, Natural) -> [DecryptionShare c] ->
 DecryptionShareCombinator c ->
 Except ErrorTally (Tally c)
proveTally
 (tally_encByChoiceByQuest, tally_countMax)
 tally_decShareByTrustee
 decShareCombinator = do
        decFactorByChoiceByQuest <-
                decShareCombinator
                 tally_encByChoiceByQuest
                 tally_decShareByTrustee
        dec <- isoZipWithM (throwE ErrorTally_NumberOfQuestions)
         (maybe (throwE ErrorTally_NumberOfChoices) return `o2`
                isoZipWith (\Encryption{..} decFactor -> encryption_vault / decFactor))
         tally_encByChoiceByQuest
         decFactorByChoiceByQuest
        let logMap = Map.fromList $ List.zip groupGenPowers [0..tally_countMax]
        let log x =
                maybe (throwE ErrorTally_CannotDecryptCount) return $
                Map.lookup x logMap
        tally_countByChoiceByQuest <- (log `mapM`)`mapM`dec
        return Tally{..}

verifyTally ::
 Reifies c FFC =>
 Tally c -> DecryptionShareCombinator c ->
 Except ErrorTally ()
verifyTally Tally{..} decShareCombinator = do
        decFactorByChoiceByQuest <- decShareCombinator tally_encByChoiceByQuest tally_decShareByTrustee
        isoZipWith3M_ (throwE ErrorTally_NumberOfQuestions)
         (isoZipWith3M_ (throwE ErrorTally_NumberOfChoices)
                 (\Encryption{..} decFactor count -> do
                        let groupGenPowCount = encryption_vault / decFactor
                        unless (groupGenPowCount == groupGen ^ fromNatural count) $
                                throwE ErrorTally_WrongProof))
         tally_encByChoiceByQuest
         decFactorByChoiceByQuest
         tally_countByChoiceByQuest

-- ** Type 'DecryptionShare'
-- | A decryption share is a 'DecryptionFactor' and a decryption 'Proof', by choice by 'Question'.
-- Computed by a trustee in 'proveDecryptionShare'.
type DecryptionShare c = [[(DecryptionFactor c, Proof c)]]

-- *** Type 'DecryptionFactor'
-- | @'encryption_nonce' '^'trusteeSecKey@
type DecryptionFactor = G

-- @('proveDecryptionShare' encByChoiceByQuest trusteeSecKey)@
proveDecryptionShare ::
 Monad m => Reifies c FFC => RandomGen r =>
 EncryptedTally c -> SecretKey c -> S.StateT r m (DecryptionShare c)
proveDecryptionShare encByChoiceByQuest trusteeSecKey =
        (proveDecryptionFactor trusteeSecKey `mapM`) `mapM` encByChoiceByQuest

proveDecryptionFactor ::
 Monad m => Reifies c FFC => RandomGen r =>
 SecretKey c -> Encryption c -> S.StateT r m (DecryptionFactor c, Proof c)
proveDecryptionFactor trusteeSecKey Encryption{..} = do
        proof <- prove trusteeSecKey [groupGen, encryption_nonce] (hash zkp)
        return (encryption_nonce^trusteeSecKey, proof)
        where zkp = decryptionShareStatement (publicKey trusteeSecKey)

decryptionShareStatement :: Reifies c FFC => PublicKey c -> BS.ByteString
decryptionShareStatement pubKey =
        "decrypt|"<>bytesNat pubKey<>"|"

-- *** Type 'ErrorTally'
data ErrorTally
 =   ErrorTally_NumberOfQuestions
     -- ^ The number of 'Question's is not the one expected.
 |   ErrorTally_NumberOfChoices
     -- ^ The number of choices is not the one expected.
 |   ErrorTally_NumberOfTrustees
     -- ^ The number of trustees is not the one expected.
 |   ErrorTally_WrongProof
     -- ^ The 'Proof' of a 'DecryptionFactor' is wrong.
 |   ErrorTally_CannotDecryptCount
     -- ^ Raised by 'proveTally' when the discrete logarithm of @'groupGen' '^'count@
     -- cannot be computed, likely because 'tally_countMax' is wrong,
     -- or because the 'EncryptedTally' or 'DecryptionShare's have not been verified.
 deriving (Eq,Show,Generic,NFData)

-- | @('verifyDecryptionShare' encTally trusteePubKey trusteeDecShare)@
-- checks that 'trusteeDecShare'
-- (supposedly submitted by a trustee whose 'PublicKey' is 'trusteePubKey')
-- is valid with respect to the 'EncryptedTally' 'encTally'.
verifyDecryptionShare ::
 Monad m => Reifies c FFC =>
 EncryptedTally c -> PublicKey c -> DecryptionShare c ->
 ExceptT ErrorTally m ()
verifyDecryptionShare encByChoiceByQuest trusteePubKey =
        let zkp = decryptionShareStatement trusteePubKey in
        isoZipWithM_ (throwE ErrorTally_NumberOfQuestions)
         (isoZipWithM_ (throwE ErrorTally_NumberOfChoices) $
         \Encryption{..} (decFactor, proof) ->
                unless (proof_challenge proof == hash zkp
                 [ commit proof groupGen trusteePubKey
                 , commit proof encryption_nonce decFactor
                 ]) $ throwE ErrorTally_WrongProof)
         encByChoiceByQuest

verifyDecryptionShareByTrustee ::
 Monad m => Reifies c FFC =>
 EncryptedTally c -> [PublicKey c] -> [DecryptionShare c] ->
 ExceptT ErrorTally m ()
verifyDecryptionShareByTrustee encTally =
        isoZipWithM_ (throwE ErrorTally_NumberOfTrustees)
         (verifyDecryptionShare encTally)