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

import Control.Monad (Monad(..), mapM, forM, join, sequence)
import Data.Bool
import Data.Eq (Eq(..))
import Data.Function (($), (.))
import Data.Functor (Functor(..), (<$>))
import Data.Foldable (Foldable, foldMap, and, toList)
import Data.Int (Int)
import Data.Maybe (Maybe(..))
import Data.Monoid (Monoid(..))
import Data.Ord (Ord(..), Ordering(..))
import Data.Semigroup (Semigroup(..))
import Data.Text (Text)
import Data.Tuple (fst, snd, curry)
import Prelude (Integral, fromIntegral, undefined, error)
import Text.Show (Show(..))
import Data.String (IsString(..))
import qualified Control.Monad.Trans.State.Strict as S
import qualified Data.List as List
import qualified Data.Text.Encoding as Text
import qualified Data.ByteString as BS
import qualified Data.Map.Strict as Map
import Data.Map.Strict (Map)

import Protocol.Arith
import Protocol.List

-- * Type '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 (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 = G
type SecretKey = E
type SecretNonce = E
type PublicKeyString = BS.ByteString

-- *** Type 'Opinion'
-- | Exponent indexing a 'Disjunction' within a list of them.
type Opinion = E

encrypt ::
 Monad m =>
 RandomGen r =>
 SubGroup q =>
 PublicKey q -> Opinion q ->
 S.StateT r m (SecretNonce q, Encryption q)
encrypt pubKey opinion = do
	secNonce <- random
	-- NOTE: preserve the 'secNonce' for 'proof'
	return $ (secNonce,)
		Encryption
		 { encryption_nonce = groupGen^secNonce
		 , encryption_vault = pubKey  ^secNonce * groupGen^opinion
		 -- NOTE: pubKey == groupGen ^ secKey
		 -- NOTE: 'index' is put as exponent in order
		 -- to make an additive homomorphism
		 -- instead of a multiplicative homomorphism.
		 }

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

type Oracle q = [Commitment q] -> Hash q
type Hash = E

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

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

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

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

validRange :: SubGroup q => E q -> E q -> [Disjunction q]
validRange mini maxi =
	List.genericTake (intE maxi - intE mini) $
	List.genericDrop (intE mini) groupGenInverses

-- ** Type 'ValidityProof'
-- | A list of 'Proof' 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 q = ValidityProof [Proof q]
 deriving (Eq,Show)

-- | @('validableEncryption' pubKey zkp ds d (secNonce, enc))@
-- returns a 'ValidityProof' that @'encryption_nonce' == 'groupGen''^''secNonce'@
-- and @'encryption_vault' == pubKey'^'secNonce'/'indexedDisj'@.
validableEncryption ::
 forall m r q.
 Monad m =>
 RandomGen r =>
 SubGroup q =>
 PublicKey q -> PublicKeyString ->
 [Disjunction q] -> Opinion q ->
 (SecretNonce q, Encryption q) ->
 S.StateT r m (ValidityProof q)
validableEncryption pubKey zkp valids index (secNonce, Encryption{..})
 | (prevDisjs,_indexedDisj:nextDisjs) <-
   List.splitAt (fromIntegral (intE index)) valids = do
	prevFakes <- fakeProof `mapM` prevDisjs
	nextFakes <- fakeProof `mapM` nextDisjs
	let challengeSum =
		neg $
		sum (proof_challenge . fst <$> prevFakes) +
		sum (proof_challenge . fst <$> nextFakes)
	genuineProof <- proof secNonce [groupGen, pubKey] $
	 -- | 'Oracle'
	 \nizkCommitments ->
		let statement =
			"prove|"<>zkp<>"|"<>
			fromString (show (intG encryption_nonce))<>","<>
			fromString (show (intG encryption_vault))<>"|" in
		let commitments =
			foldMap snd prevFakes <>
			nizkCommitments <>
			foldMap snd nextFakes in
		hash statement commitments + challengeSum
	return $
		ValidityProof $
		 (fst <$> prevFakes) <>
		 [genuineProof] <>
		 (fst <$> nextFakes)
 | otherwise = error "validableEncryption: bad disjunction index"
	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 -> PublicKeyString ->
 [Disjunction q] ->
 (Encryption q, ValidityProof q) -> Bool
validateEncryption pubKey zkp disjs (Encryption{..}, ValidityProof proofs) =
	List.length disjs == List.length proofs &&
	hash statement commitments == challengeSum
	where
	challengeSum = sum (proof_challenge <$> proofs)
	statement =
		"prove|"<>zkp<>"|"<>
		fromString (show (intG encryption_nonce))<>","<>
		fromString (show (intG encryption_vault))<>"|"
	commitments = join $ List.zipWith commitment 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 q = Question
 {   question_text    :: Text
 ,   question_answers :: [Text]
 ,   question_min     :: E q
 ,   question_max     :: E q
 -- ,   question_blank :: Maybe Bool
 } deriving (Eq, Show)

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

answer ::
 forall m r q.
 Monad m =>
 RandomGen r =>
 SubGroup q =>
 PublicKey q -> PublicKeyString ->
 Question q ->
 [Opinion q] ->
 S.StateT r m (Answer q)
answer pubKey zkp Question{..} opinions = do
	encryptions <- encrypt pubKey `mapM` opinions
	individualProofs :: [ValidityProof q] <-
		sequence $ List.zipWith
		 (validableEncryption pubKey zkp validBool)
		 opinions encryptions
	sumProof <-
		validableEncryption pubKey zkp
		 (validRange question_min question_max)
		 (sum opinions - question_min)
		 ( sum (fst <$> encryptions)
		 , sum (snd <$> encryptions) )
	return Answer
	 { answer_opinions =
		List.zip
		 (snd <$> encryptions) -- NOTE: drop the secretNonce
		 individualProofs
	 , answer_sumProof = sumProof
	 }

validateAnswer ::
 SubGroup q =>
 PublicKey q -> PublicKeyString ->
 Question q ->
 Answer q -> Bool
validateAnswer pubKey zkp Question{..} Answer{..} =
	and (validateEncryption pubKey zkp validBool <$> answer_opinions) &&
	validateEncryption pubKey zkp
	 (validRange question_min question_max)
	 ( sum (fst <$> answer_opinions)
	 , answer_sumProof )