{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}
module Protocol.Election where

import Control.Monad (Monad(..), mapM, zipWithM)
import Control.Monad.Morph (MFunctor(..))
import Control.Monad.Trans.Class (MonadTrans(..))
import Data.Bool
import Data.Either (either)
import Data.Eq (Eq(..))
import Data.Foldable (Foldable, foldMap, and)
import Data.Function (($), id, const)
import Data.Functor (Functor, (<$>))
import Data.Functor.Identity (Identity(..))
import Data.Maybe (Maybe(..), fromMaybe)
import Data.Ord (Ord(..))
import Data.Semigroup (Semigroup(..))
import Data.String (IsString(..))
import Data.Text (Text)
import Data.Traversable (Traversable(..))
import Data.Tuple (fst, snd)
import GHC.Natural (minusNaturalMaybe)
import Numeric.Natural (Natural)
import Prelude (error, fromIntegral)
import Text.Show (Show(..))
import qualified Control.Monad.Trans.Except as Exn
import qualified Control.Monad.Trans.State.Strict as S
import qualified Data.ByteString as BS
import qualified Data.List as List

import Protocol.Arithmetic
import Protocol.Credential

-- * Type 'Encryption'
-- | ElGamal-like encryption.
-- Its security relies on the /Discrete Logarithm problem/.
--
-- Because ('groupGen' '^'encNonce '^'secKey '==' 'groupGen' '^'secKey '^'encNonce),
-- knowing @secKey@, one can divide 'encryption_vault' by @('encryption_nonce' '^'secKey)@
-- to decipher @('groupGen' '^'clear)@, then @clear@ must be small to be decryptable,
-- because it is encrypted as a power of 'groupGen' to enable the additive homomorphism.
data Encryption q = Encryption
 { encryption_nonce :: G q
   -- ^ Public part of the random 'encNonce': @('groupGen' '^'encNonce)@
 , encryption_vault :: G q
   -- ^ Encrypted clear: @('pubKey' '^'r '*' 'groupGen' '^'clear)@
 } deriving (Eq,Show)

-- | Additive homomorphism.
-- Using the fact that: @'groupGen' '^'x '*' 'groupGen' '^'y '==' 'groupGen' '^'(x'+'y)@.
instance SubGroup q => Additive (Encryption q) where
	zero = Encryption one one
	x+y = Encryption
	 (encryption_nonce x * encryption_nonce y)
	 (encryption_vault x * encryption_vault y)

-- *** Type 'EncryptionNonce'
type EncryptionNonce = E

-- | @('encrypt' pubKey clear)@ returns an ElGamal-like 'Encryption'.
--
-- WARNING: the secret encryption nonce (@encNonce@)
-- is returned alongside the 'Encryption'
-- in order to prove the validity of the encrypted clear in 'prove',
-- but this secret @encNonce@ MUST be forgotten after that,
-- as it may be used to decipher the 'Encryption'
-- without the secret key associated with 'pubKey'.
encrypt ::
 Monad m => RandomGen r => SubGroup q =>
 PublicKey q -> E q ->
 S.StateT r m (EncryptionNonce q, Encryption q)
encrypt pubKey clear = do
	encNonce <- random
	-- NOTE: preserve the 'encNonce' for 'prove'.
	return $ (encNonce,)
		Encryption
		 { encryption_nonce = groupGen^encNonce
		 , encryption_vault = pubKey  ^encNonce * groupGen^clear
		 -- NOTE: 'clear' is put as exponent in order
		 -- to make an additive homomorphism
		 -- instead of a multiplicative homomorphism.
		 -- log (a*b) = log a + log b
		 }

-- * Type 'Proof'
-- | 'Proof' of knowledge of a discrete logarithm:
-- @secret == logBase base (base^secret)@.
--
-- NOTE: Since @(pubKey == 'groupGen' '^'secKey)@, then:
-- @(logBase 'encryption_nonce' ('encryption_vault' '*' 'encryption_nonce') '==' secKey '+' clear)@.
data Proof q = Proof
 { proof_challenge :: Challenge q
   -- ^ 'Challenge' sent by the verifier to the prover
   -- to ensure that the prover really has knowledge
   -- of the secret and is not replaying.
   -- Actually, 'proof_challenge' is not sent in a 'prove',
   -- but derived from the prover's 'Commitment's and statements
   -- with a collision resistant hash.
 , proof_response :: E q
   -- ^ Response sent by the prover to the verifier.
   -- Usually: @nonce '+' sec '*' 'proof_challenge'@.
   --
   -- To be computed efficiently, it requires @sec@:
   -- either the @secKey@ (in 'signature_proof')
   -- or the @encNonce@ (in 'prove').
 } deriving (Eq,Show)

-- ** Type 'Challenge'
type Challenge = E

-- ** Type 'Oracle'
-- An 'Oracle' returns the 'Challenge' of the 'Commitment's
-- by hashing them (eventually with other 'Commitment's).
--
-- Used in 'prove' it enables a Fiat-Shamir transformation
-- of an /interactive zero-knowledge/ (IZK) proof
-- into a /non-interactive zero-knowledge/ (NIZK) proof.
-- That is to say that the verifier does not have
-- to send a 'Challenge' to the prover.
-- Indeed, the prover now handles the 'Challenge'
-- which becomes a (collision resistant) hash
-- of the prover's commitments (and statements to be a stronger proof).
type Oracle list q = list (Commitment q) -> Challenge q

-- | @('prove' sec commitments oracle)@
-- returns a 'Proof' that @sec@ is known.
--
-- The 'Oracle' is given the 'commitments'
-- raised to the power of the secret nonce of the 'Proof',
-- as those are the 'commitments' that the verifier will obtain
-- when composing the 'proof_challenge' and 'proof_response' together
-- (in 'encryptionCommitments').
--
-- NOTE: 'sec' is @secKey@ in 'signature_proof' or @encNonce@ in 'proveEncryption'.
--
-- NOTE: The 'commitments' are @['groupGen']@ in 'signature_proof'
-- or @['groupGen', 'pubKey']@ in 'proveEncryption'.
--
-- WARNING: for 'prove' to be a so-called /strong Fiat-Shamir transformation/ (not a weak):
-- the statement must be included in the hash (not only the commitments).
--
-- NOTE: a 'random' @nonce@ is used to ensure each 'prove'
-- does not reveal any information regarding the secret 'sec'.
prove ::
 Monad m => RandomGen r => SubGroup q => Functor list =>
 E q -> list (Commitment q) -> Oracle list q -> S.StateT r m (Proof q)
prove sec commitments oracle = do
	nonce <- random
	let proof_challenge = oracle $ (^ nonce) <$> commitments
	return Proof
	 { proof_challenge
	 , proof_response = nonce - sec*proof_challenge
	 }

-- ** Type 'Commitment'
type Commitment = G

-- | @('commit' proof x y)@ returns a 'Commitment'
-- from the given 'Proof' with the knowledge of the verifier.
--
-- NOTE: Contrary to Helios-C specifications,
-- @('*')@ is used instead of @('/')@
-- to avoid the performance cost of a modular exponentiation
-- @('^' ('groupOrder' '-' 'one'))@,
-- this is compensated by using @('-')@ instead of @('+')@ in 'prove'.
commit :: SubGroup q => Proof q -> G q -> G q -> Commitment q
commit Proof{..} x y = x^proof_response * y^proof_challenge
{-# INLINE commit #-}

-- ** Type 'Opinion'
-- | Index of a 'Disjunction' within a list of them.
-- It is encrypted as an 'E'xponent by 'encrypt'.
type Opinion = E

-- ** Type 'Disjunction'
-- | A 'Disjunction' is an 'inv'ersed @('groupGen' '^'opinion)@
-- it's used in 'proveEncryption' to generate a 'Proof'
-- that an 'encryption_vault' contains a given @('groupGen' '^'opinion)@,
type Disjunction = G

booleanDisjunctions :: SubGroup q => [Disjunction q]
booleanDisjunctions = List.take 2 groupGenInverses

intervalDisjunctions :: SubGroup q => Opinion q -> Opinion q -> [Disjunction q]
intervalDisjunctions mini maxi =
	List.genericTake (fromMaybe 0 $ (natE maxi + 1)`minusNaturalMaybe`natE mini) $
	List.genericDrop (natE mini) $
	groupGenInverses

-- ** Type 'DisjProof'
-- | A list of 'Proof's to prove that the 'Opinion' within an 'Encryption'
-- is indexing a 'Disjunction' within a list of them,
-- without knowing which 'Opinion' it is.
newtype DisjProof q = DisjProof [Proof q]
 deriving (Eq,Show)

-- | @('proveEncryption' pubKey zkp disjs opin (encNonce, enc))@
-- returns a 'DisjProof' that 'enc' 'encrypt's
-- one of the 'Disjunction's within 'disjs',
-- without revealing which one it is.
--
-- A /NIZK Disjunctive Chaum Pedersen Logarithm Equality/ is used.
proveEncryption ::
 forall m r q.
 Monad m => RandomGen r => SubGroup q =>
 PublicKey q -> ZKP ->
 [Disjunction q] -> Opinion q ->
 (EncryptionNonce q, Encryption q) ->
 S.StateT r (Exn.ExceptT ErrorProove m) (DisjProof q)
proveEncryption pubKey zkp disjs opinion (encNonce, enc)
 | (prevDisjs, _indexedDisj:nextDisjs) <-
   List.genericSplitAt (natE opinion) disjs = do
	-- Fake proofs for all values except the correct one.
	prevFakes <- fakeProof `mapM` prevDisjs
	nextFakes <- fakeProof `mapM` nextDisjs
	let prevProofs = fst <$> prevFakes
	let nextProofs = fst <$> nextFakes
	let challengeSum =
		sum (proof_challenge <$> prevProofs) +
		sum (proof_challenge <$> nextProofs)
	correctProof <- prove encNonce [groupGen, pubKey] $
	 -- 'Oracle'
	 \correctCommitments ->
		let commitments =
			foldMap snd prevFakes <>
			correctCommitments <>
			foldMap snd nextFakes in
		hash (encryptionStatement zkp enc) commitments - challengeSum
	return $ DisjProof $ prevProofs <> (correctProof : nextProofs)
 | otherwise = lift $ Exn.throwE $
	ErrorProove_InvalidOpinion
	 (fromIntegral $ List.length disjs)
	 (natE opinion)
	where
	fakeProof :: Disjunction q -> S.StateT r (Exn.ExceptT ErrorProove m) (Proof q, [Commitment q])
	fakeProof disj = do
		-- Returns 'Commitment's verifiables by the verifier,
		-- but computed from random 'proof_challenge' and 'proof_response'
		-- instead of correct ones.
		proof_challenge <- random
		proof_response  <- random
		let proof = Proof{..}
		return (proof, encryptionCommitments pubKey enc (disj, proof))

verifyEncryption ::
 Monad m =>
 SubGroup q =>
 PublicKey q -> ZKP ->
 [Disjunction q] ->
 (Encryption q, DisjProof q) ->
 Exn.ExceptT ErrorValidateEncryption m Bool
verifyEncryption pubKey zkp disjs (enc, DisjProof proofs)
 | List.length proofs /= List.length disjs =
	Exn.throwE $ ErrorValidateEncryption_InvalidProofLength
	 (fromIntegral $ List.length proofs)
	 (fromIntegral $ List.length disjs)
 | otherwise = return $ challengeSum == hash (encryptionStatement zkp enc) commitments
	where
	challengeSum = sum (proof_challenge <$> proofs)
	commitments = foldMap (encryptionCommitments pubKey enc) (List.zip disjs proofs)

encryptionStatement :: SubGroup q => ZKP -> Encryption q -> BS.ByteString
encryptionStatement (ZKP zkp) Encryption{..} =
	"prove|"<>zkp<>"|"<>
	fromString (show (natG encryption_nonce))<>","<>
	fromString (show (natG encryption_vault))<>"|"

-- | @('encryptionCommitments' pubKey enc (disj,proof))@
-- returns the 'Commitment's with only the knowledge of the verifier.
--
-- The 'Proof' comes from 'prove' of @fakeProof@ in 'proveEncryption'.
encryptionCommitments ::
 SubGroup q =>
 PublicKey q -> Encryption q ->
 (Disjunction q, Proof q) -> [G q]
encryptionCommitments pubKey Encryption{..} (disj, proof) =
	[ commit proof groupGen encryption_nonce
	  -- == groupGen ^ nonce if 'Proof' comes from 'prove'
	, commit proof pubKey (encryption_vault*disj)
	  -- == pubKey ^ nonce if 'Proof' comes from 'prove'
	  -- and 'encryption_vault' encrypts (- logBase groupGen disj).
	]

-- ** Type 'ZKP'
-- | Zero-knowledge proof
newtype ZKP = ZKP BS.ByteString

-- ** Type 'ErrorProove'
-- | Error raised by 'proveEncryption'.
data ErrorProove
 =   ErrorProove_InvalidOpinion Natural Natural
     -- ^ When the opinion is not within the number of 'Disjunction's.
 deriving (Eq,Show)

-- ** Type 'ErrorValidateEncryption'
-- | Error raised by 'verifyEncryption'.
data ErrorValidateEncryption
 =   ErrorValidateEncryption_InvalidProofLength Natural Natural
     -- ^ When the number of proofs is different than
     -- the number of 'Disjunction's.
 deriving (Eq,Show)

-- * Type 'Question'
data Question q = Question
 { question_text    :: Text
 , question_choices :: [Text]
 , question_mini    :: Opinion q
 , question_maxi    :: Opinion q
 -- , question_blank :: Maybe Bool
 } deriving (Eq, Show)

-- * Type 'Answer'
data Answer q = Answer
 { answer_opinions :: [(Encryption q, DisjProof q)]
   -- ^ Encrypted 'Opinion' for each 'question_choices'
   -- with a 'DisjProof' that they belong to [0,1].
 , answer_sumProof :: DisjProof q
   -- ^ Proofs that the sum of the 'Opinon's encrypted in 'answer_opinions'
   -- is an element of @[mini..maxi]@.
 -- , answer_blankProof ::
 } deriving (Eq,Show)

-- ** Type 'ErrorAnswer'
-- | Error raised by 'encryptAnswer'.
data ErrorAnswer
 =   ErrorAnswer_WrongNumberOfOpinions Natural Natural
     -- ^ When the number of opinions is different than
     -- the number of choices ('question_choices').
 |   ErrorAnswer_WrongSumOfOpinions Natural Natural Natural
     -- ^ When the sum of opinions is not within the bounds
     -- of 'question_mini' and 'question_maxi'.
 deriving (Eq,Show)

-- | @('encryptAnswer' pubKey zkp quest opinions)@
-- returns an 'Answer' validable by 'verifyAnswer',
-- unless an 'ErrorAnswer' is returned.
encryptAnswer ::
 Monad m => RandomGen r => SubGroup q =>
 PublicKey q -> ZKP ->
 Question q -> [Bool] ->
 S.StateT r (Exn.ExceptT ErrorAnswer m) (Answer q)
encryptAnswer pubKey zkp Question{..} opinionsBools
 | not (question_mini <= opinionsSum && opinionsSum <= question_maxi) =
	lift $ Exn.throwE $
		ErrorAnswer_WrongSumOfOpinions
		 (natE opinionsSum)
		 (natE question_mini)
		 (natE question_maxi)
 | List.length opinions /= List.length question_choices =
	lift $ Exn.throwE $
		ErrorAnswer_WrongNumberOfOpinions
		 (fromIntegral $ List.length opinions)
		 (fromIntegral $ List.length question_choices)
 | otherwise = do
	encryptions <- encrypt pubKey `mapM` opinions
	hoist (Exn.withExceptT (\case
		 ErrorProove_InvalidOpinion{} -> error "encryptAnswer: impossible happened"
	 )) $ do
		individualProofs <- zipWithM
		 (proveEncryption pubKey zkp booleanDisjunctions)
		 opinions encryptions
		sumProof <- proveEncryption pubKey zkp
		 (intervalDisjunctions question_mini question_maxi)
		 (opinionsSum - question_mini)
		 ( sum (fst <$> encryptions) -- NOTE: sum the 'encNonce's
		 , sum (snd <$> encryptions) -- NOTE: sum the 'Encryption's
		 )
		return $ Answer
		 { answer_opinions = List.zip
			 (snd <$> encryptions) -- NOTE: drop encNonce
			 individualProofs
		 , answer_sumProof = sumProof
		 }
 where
	opinionsSum = sum opinions
	opinions = (\o -> if o then one else zero) <$> opinionsBools

verifyAnswer ::
 SubGroup q =>
 PublicKey q -> ZKP ->
 Question q -> Answer q -> Bool
verifyAnswer pubKey zkp Question{..} Answer{..}
 | List.length question_choices /= List.length answer_opinions = False
 | otherwise = either (const False) id $ Exn.runExcept $ do
	validOpinions <-
		verifyEncryption pubKey zkp booleanDisjunctions
		 `traverse` answer_opinions
	validSum <- verifyEncryption pubKey zkp
	 (intervalDisjunctions question_mini question_maxi)
	 ( sum (fst <$> answer_opinions)
	 , answer_sumProof )
	return (and validOpinions && validSum)

-- * Type 'Election'
data Election q = Election
 { election_name        :: Text
 , election_description :: Text
 , election_publicKey   :: PublicKey q
 , election_questions   :: [Question q]
 , election_uuid        :: UUID
 , election_hash        :: Hash -- TODO: serialize to JSON to calculate this
 } deriving (Eq,Show)

-- ** Type 'Hash'
newtype Hash = Hash Text
 deriving (Eq,Ord,Show)

-- * Type 'Ballot'
data Ballot q = Ballot
 { ballot_answers       :: [Answer q]
 , ballot_signature     :: Maybe (Signature q)
 , ballot_election_uuid :: UUID
 , ballot_election_hash :: Hash
 }

-- | @('encryptBallot' elec ('Just' secKey) opinionsByQuest)@
-- returns a 'Ballot' signed by 'secKey' (the voter's secret key)
-- where 'opinionsByQuest' is a list of 'Opinion's
-- on each 'question_choices' of each 'election_questions'.
encryptBallot ::
 Monad m => RandomGen r => SubGroup q =>
 Election q -> Maybe (SecretKey q) -> [[Bool]] ->
 S.StateT r (Exn.ExceptT ErrorBallot m) (Ballot q)
encryptBallot Election{..} secKeyMay opinionsByQuest
 | List.length election_questions /= List.length opinionsByQuest =
	lift $ Exn.throwE $
		ErrorBallot_WrongNumberOfAnswers
		 (fromIntegral $ List.length opinionsByQuest)
		 (fromIntegral $ List.length election_questions)
 | otherwise = do
	let (keysMay, zkp) =
		case secKeyMay of
		 Nothing -> (Nothing, ZKP "")
		 Just secKey ->
			( Just (secKey, pubKey)
			, ZKP (fromString (show (natG pubKey))) )
			where pubKey = groupGen ^ secKey
	ballot_answers <-
		hoist (Exn.withExceptT ErrorBallot_Answer) $
			zipWithM (encryptAnswer election_publicKey zkp)
			 election_questions opinionsByQuest
	ballot_signature <- case keysMay of
	 Nothing -> return Nothing
	 Just (secKey, signature_publicKey) -> do
		signature_proof <-
			prove secKey (Identity groupGen) $
			 \(Identity commitment) ->
				hash
				 (signatureCommitments zkp commitment)
				 (signatureStatement ballot_answers)
		return $ Just Signature{..}
	return Ballot
	 { ballot_answers
	 , ballot_election_hash = election_hash
	 , ballot_election_uuid = election_uuid
	 , ballot_signature
	 }

verifyBallot :: SubGroup q => Election q -> Ballot q -> Bool
verifyBallot Election{..} Ballot{..} =
	ballot_election_uuid == election_uuid &&
	ballot_election_hash == election_hash &&
	List.length election_questions == List.length ballot_answers &&
	let (isValidSign, zkpSign) =
		case ballot_signature of
		 Nothing -> (True, ZKP "")
		 Just Signature{..} ->
			let zkp = ZKP (fromString (show (natG signature_publicKey))) in
			(, zkp) $
				proof_challenge signature_proof == hash
				 (signatureCommitments zkp (commit signature_proof groupGen signature_publicKey))
				 (signatureStatement ballot_answers)
	in
	and $ isValidSign :
		List.zipWith (verifyAnswer election_publicKey zkpSign)
		 election_questions ballot_answers

-- ** Type 'Signature'
-- | Schnorr-like signature.
--
-- Used to avoid 'Ballot' stuffing.
data Signature q = Signature
 { signature_publicKey :: PublicKey q
 , signature_proof     :: Proof q
 }

-- | @('signatureStatement' answers)@
-- returns all the 'encryption_nonce's and 'encryption_vault's
-- of the given @answers@.
signatureStatement :: Foldable f => SubGroup q => f (Answer q) -> [G q]
signatureStatement =
	foldMap $ \Answer{..} ->
		(`foldMap` answer_opinions) $ \(Encryption{..}, _proof) ->
			[encryption_nonce, encryption_vault]

-- | @('signatureCommitments' zkp commitment)@
-- returns the hashable content from the knowledge of the verifier.
signatureCommitments :: SubGroup q => ZKP -> Commitment q -> BS.ByteString
signatureCommitments (ZKP zkp) commitment =
	"sig|"<>zkp<>"|"<>fromString (show (natG commitment))<>"|"

-- ** Type 'ErrorBallot'
-- | Error raised by 'encryptBallot'.
data ErrorBallot
 =   ErrorBallot_WrongNumberOfAnswers Natural Natural
     -- ^ When the number of answers
     -- is different than the number of questions.
 |   ErrorBallot_Answer ErrorAnswer
     -- ^ When 'encryptAnswer' raised an 'ErrorAnswer'.
 deriving (Eq,Show)