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

import Control.Monad (Monad(..), mapM, sequence)
import Data.Bool
import Data.Eq (Eq(..))
import Data.Foldable (Foldable, foldMap, and)
import Data.Function (($), (.))
import Data.Functor ((<$>))
import Data.Maybe (Maybe(..), fromJust)
import Data.Ord (Ord(..))
import Data.Proxy (Proxy(..))
import Data.Semigroup (Semigroup(..))
import Data.String (IsString(..))
import Data.Text (Text)
import Data.Tuple (fst, snd, uncurry)
import GHC.TypeNats
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 Protocol.Arithmetic
import Utils.MeasuredList as ML
import qualified Utils.Natural as Nat
import qualified Utils.Constraint as Constraint

-- * Type 'Encryption'
-- | ElGamal-like encryption.
data Encryption q = Encryption
 { encryption_nonce :: G q
   -- ^ Public part of the random 'secNonce': @('groupGen' '^'r)@
 , encryption_vault :: G q
   -- ^ Encrypted opinion: @('pubKey' '^'r '*' 'groupGen' '^'opinion)@
 } 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 'PublicKey'
type PublicKey = G
-- ** Type 'SecretKey'
type SecretKey = E
-- *** Type 'SecretNonce'
type SecretNonce = E
-- ** Type 'ZKP'
-- | Zero-knowledge proof
type ZKP = BS.ByteString

-- ** Type 'Opinion'
-- | Index of a 'Disjunction' within a 'ML.MeasuredList' of them.
-- It is encoded as an 'E'xponent in 'encrypt'.
type Opinion = Nat.Index

-- | @('encrypt' pubKey opinion)@ returns an ElGamal-like 'Encryption'.
--
-- WARNING: the secret nonce is returned alongside the 'Encryption'
-- in order to prove the validity of the encrypted content in 'nizkProof',
-- but this secret nonce 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 -> Opinion ds ->
 S.StateT r m (SecretNonce q, Encryption q)
encrypt pubKey (Nat.Index o) = do
	let opinion = inE (natVal o)
	secNonce <- random
	-- NOTE: preserve the 'secNonce' for 'nizkProof'.
	return $ (secNonce,)
		Encryption
		 { encryption_nonce = groupGen^secNonce
		 , encryption_vault = pubKey  ^secNonce * groupGen^opinion
		 -- NOTE: pubKey == groupGen ^ secKey
		 -- NOTE: 'opinion' is put as exponent in order
		 -- to make an additive homomorphism
		 -- instead of a multiplicative homomorphism.
		 }

-- * Type 'Proof'
data Proof q = Proof
 { proof_challenge :: Challenge q
 , proof_response  :: E q
 } deriving (Eq,Show)

-- ** Type 'Challenge'
type Challenge = E
-- ** Type 'Oracle'
type Oracle q = [Commitment q] -> Challenge q

-- | Fiat-Shamir transformation
-- of an interactive zero-knowledge (IZK) proof
-- into a non-interactive zero-knowledge (NIZK) proof.
nizkProof ::
 Monad m => RandomGen r => SubGroup q =>
 SecretNonce q -> [Commitment q] -> Oracle q -> S.StateT r m (Proof q)
nizkProof secNonce commits oracle = do
	nonce <- random
	let commitments = (^ nonce) <$> commits
	let proof_challenge = oracle commitments
	return Proof
	 { proof_challenge
	 , proof_response = nonce + secNonce*proof_challenge
	 }

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

-- ** 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 => ML.MeasuredList 2 (Disjunction q)
booleanDisjunctions = fromJust $ ML.fromList $ List.take 2 groupGenInverses

intervalDisjunctions ::
 forall q mini maxi.
 SubGroup q =>
 Bounds mini maxi ->
 ML.MeasuredList (maxi-mini) (Disjunction q)
intervalDisjunctions Bounds{}
 | Constraint.Proof <- (Nat.<=) @mini @maxi =
	fromJust $
	ML.fromList $
	List.genericTake (Nat.nat @(maxi-mini)) $
	List.genericDrop (Nat.nat @mini) $
	groupGenInverses

-- ** Type 'ValidityProof'
-- | 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 ValidityProof disjs q = ValidityProof (ML.MeasuredList disjs (Proof q))
 deriving (Eq,Show)

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

proveEncryption ::
 forall disjs m r q.
 Nat.Known disjs =>
 Monad m => RandomGen r => SubGroup q =>
 PublicKey q -> ZKP ->
 ML.MeasuredList disjs (Disjunction q) -> Opinion disjs ->
 (SecretNonce q, Encryption q) ->
 S.StateT r m (ValidityProof disjs q)
proveEncryption pubKey zkp disjs
 (Nat.Index (o::Proxy o))
 (secNonce, enc@Encryption{..})
 -- NOTE: the 'Constraint.Proof's below are needed to prove
 -- that the returned 'ValidityProof' has the same length
 -- than the given list of 'Disjunction's.
 | Constraint.Proof <- (Nat.+<=) @o @1 @disjs -- prove that o+1<=disjs implies 1<=disjs-o and o<=disjs
 , Constraint.Proof <- (Nat.<=) @o @disjs     -- prove that o<=disjs implies disjs-o is a Natural and o+(disjs-o) ~ disjs
 , Constraint.Proof <- (Nat.<=) @1 @(disjs-o) -- prove that ((disjs-o)-1)+1 ~ disjs-o
 = do
	let (prevDisjs, ML.uncons -> (_indexedDisj,nextDisjs)) = ML.splitAt o disjs
	prevFakes <- fakeProof `mapM` prevDisjs
	nextFakes <- fakeProof `mapM` nextDisjs
	let challengeSum =
		neg $
		sum (proof_challenge . fst <$> prevFakes) +
		sum (proof_challenge . fst <$> nextFakes)
	genuineProof <- nizkProof secNonce [groupGen, pubKey] $
	 -- 'Oracle'
	 \nizkCommitments ->
		let commitments =
			foldMap snd prevFakes <>
			nizkCommitments <>
			foldMap snd nextFakes in
		-- NOTE: this is a so-called strong Fiat-Shamir transformation (not a weak):
		-- because the statement is included in the hash (not only the commitments).
		hash (encryptionStatement zkp enc) commitments + challengeSum
	return $
		ValidityProof $
			ML.concat
			 (fst <$> prevFakes)
			 (ML.cons genuineProof (fst <$> nextFakes))
	where
	fakeProof :: Disjunction q -> S.StateT r m (Proof q, [Commitment q])
	fakeProof disj = do
		proof_challenge <- random
		proof_response  <- random
		let commitments =
			 [ groupGen^proof_response / encryption_nonce       ^proof_challenge
			 , pubKey  ^proof_response / (encryption_vault*disj)^proof_challenge
			 ]
		return (Proof{..}, commitments)

validateEncryption ::
 SubGroup q =>
 PublicKey q -> ZKP ->
 ML.MeasuredList n (Disjunction q) ->
 (Encryption q, ValidityProof n q) -> Bool
validateEncryption pubKey zkp disjs (enc@Encryption{..}, ValidityProof proofs) =
	hash (encryptionStatement zkp enc) commitments == challengeSum
	where
	challengeSum = sum (proof_challenge <$> proofs)
	commitments = foldMap commitment (ML.zip disjs proofs)
		where commitment (disj, Proof{..}) =
			-- g = groupGen
			-- h = pubKey
			-- y1 = encryption_nonce
			-- y2 = (encryption_vault * disj)
			-- com1 = g^res / y1 ^ ch
			-- com2 = h^res / y2 ^ ch
			[ groupGen^proof_response / encryption_nonce       ^proof_challenge
			, pubKey  ^proof_response / (encryption_vault*disj)^proof_challenge
			]

-- * Type 'Question'
data Question choices (mini::Nat) (maxi::Nat) q =
 Question
 { question_text    :: Text
 , question_answers :: ML.MeasuredList choices Text
 , question_bounds  :: Bounds mini maxi
 -- , question_blank :: Maybe Bool
 } deriving (Eq, Show)

-- ** Type 'Bounds'
data Bounds mini maxi =
 ((mini<=maxi), Nat.Known mini, Nat.Known maxi) =>
 Bounds (Proxy mini) (Proxy maxi)
instance Show (Bounds mini maxi) where
	showsPrec p Bounds{} = showsPrec p (Nat.nat @mini, Nat.nat @maxi)
instance Eq (Bounds mini maxi) where
	_==_ = True

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

-- | @('answer' pubKey zkp quest opinions)@
-- returns a validable 'Answer',
-- unless the given 'opinions' do not respect 'question_bounds'.
answer ::
 forall m r q mini maxi choices.
 Monad m => RandomGen r => SubGroup q =>
 PublicKey q -> ZKP ->
 Question choices mini maxi q ->
 ML.MeasuredList choices (Opinion 2) ->
 S.StateT r m (Maybe (Answer choices mini maxi q))
answer pubKey zkp Question{..} opinions
 | Bounds{} <- question_bounds
 , SomeNat (_opinionsSum::Proxy opinionsSum) <-
   someNatVal $ sum $ (\(Nat.Index o) -> natVal o) <$> opinions
   -- prove that opinionsSum-mini is a Natural
 , Just Constraint.Proof <- (Nat.<=?) @mini @opinionsSum
   -- prove that (opinionsSum-mini)+1 is a Natural
 , Constraint.Proof <- (Nat.+) @(opinionsSum-mini) @1
   -- prove that maxi-mini is a Natural
 , Constraint.Proof <- (Nat.<=) @mini @maxi
   -- prove that (opinionsSum-mini)+1 <= maxi-mini
 , Just Constraint.Proof <- (Nat.<=?) @((opinionsSum-mini)+1) @(maxi-mini)
 = do
	encryptions <- encrypt pubKey `mapM` opinions
	individualProofs <-
		sequence $ ML.zipWith
		 (proveEncryption pubKey zkp booleanDisjunctions)
		 opinions encryptions
	sumProof <-
		proveEncryption pubKey zkp
		 (intervalDisjunctions question_bounds)
		 (Nat.Index $ Proxy @(opinionsSum-mini))
		 ( sum (fst <$> encryptions)
		 , sum (snd <$> encryptions) )
	return $ Just Answer
	 { answer_opinions = ML.zip
		 (snd <$> encryptions) -- NOTE: drop secNonce
		 individualProofs
	 , answer_sumProof = sumProof
	 }
 | otherwise = return Nothing

validateAnswer ::
 SubGroup q =>
 PublicKey q -> ZKP ->
 Question choices mini maxi q ->
 Answer choices mini maxi q -> Bool
validateAnswer pubKey zkp Question{..} Answer{..} =
	and (validateEncryption pubKey zkp booleanDisjunctions <$> answer_opinions) &&
	validateEncryption pubKey zkp
	 (intervalDisjunctions question_bounds)
	 ( sum (fst <$> answer_opinions)
	 , answer_sumProof )

-- * Type 'Election'
data Election quests choices mini maxi q = Election
 { election_name        :: Text
 , election_description :: Text
 , election_publicKey   :: PublicKey q
 , election_questions   :: ML.MeasuredList quests (Question choices mini maxi q)
 , election_uuid        :: UUID
 , election_hash        :: Hash
 } deriving (Eq,Show)

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

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

-- * Type 'Ballot'
data Ballot quests choices mini maxi q = Ballot
 { ballot_answers       :: ML.MeasuredList quests (Answer choices mini maxi q)
 , ballot_signature     :: Maybe (Signature q)
 , ballot_election_uuid :: UUID
 , ballot_election_hash :: Hash
 }

ballot ::
 Monad m =>
 RandomGen r =>
 SubGroup q =>
 Election quests choices mini maxi q ->
 Maybe (SecretKey q) ->
 ML.MeasuredList quests (ML.MeasuredList choices (Opinion 2)) ->
 S.StateT r m (Maybe (Ballot quests choices mini maxi q))
ballot Election{..} secKeyMay opinionsByQuest = do
	let (keysMay, zkp) =
		case secKeyMay of
		 Nothing -> (Nothing, "")
		 Just secKey ->
			( Just (secKey, pubKey)
			, fromString (show (natG pubKey)) )
			where pubKey = groupGen ^ secKey
	answersByQuestMay <-
		(sequence <$>) $
		uncurry (answer election_publicKey zkp) `mapM`
		ML.zip election_questions opinionsByQuest
	case answersByQuestMay of
	 Nothing -> return Nothing
	 Just answersByQuest -> do
		ballot_signature <- case keysMay of
		 Nothing -> return Nothing
		 Just (secKey, pubKey) -> do
			w <- random
			let commitment = groupGen ^ w
			let proof_challenge = hash
				 (signatureCommitments zkp commitment)
				 (signatureStatement answersByQuest)
			return $ Just Signature
			 { signature_publicKey = pubKey
			 , signature_proof = Proof
				 { proof_challenge
				 , proof_response = w - secKey*proof_challenge
				 }
			 }
		return $ Just Ballot
		 { ballot_answers = answersByQuest
		 , ballot_election_hash = election_hash
		 , ballot_election_uuid = election_uuid
		 , ballot_signature
		 }

validateBallot ::
 SubGroup q =>
 Election quests choices mini maxi q ->
 Ballot quests choices mini maxi q ->
 Bool
validateBallot Election{..} Ballot{..} =
	ballot_election_uuid == election_uuid &&
	ballot_election_hash == election_hash &&
	let (validSign, zkp) =
		case ballot_signature of
		 Nothing -> (True, "")
		 Just (Signature pubKey Proof{..}) ->
			let zkp = fromString (show (natG pubKey)) in
			let validSign =
				let commitment = groupGen ^ proof_response * pubKey ^ proof_challenge in
				let prefix     = signatureCommitments zkp commitment in
				let contents   = signatureStatement ballot_answers in
				hash prefix contents == proof_challenge
			in (validSign, zkp)
	in
	validSign &&
	and (ML.zipWith (validateAnswer election_publicKey zkp)
	 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 ::
 Foldable f => SubGroup q =>
 f (Answer choices mini maxi q) -> [G q]
signatureStatement =
	foldMap $ \Answer{..} ->
		(`foldMap` answer_opinions) $ \(Encryption{..}, _vp) ->
			[encryption_nonce, encryption_vault]

signatureCommitments :: SubGroup q => ZKP -> Commitment q -> BS.ByteString
signatureCommitments zkp commitment =
	"sig|"<>zkp<>"|"<>fromString (show (natG commitment))<>"|"