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

import Control.Monad (Monad(..), mapM, forM, join, sequence)
import Data.Bool
import Data.Eq (Eq(..))
import Data.Foldable (Foldable, foldMap, and, toList)
import Data.Function (($), (.))
import Data.Functor (Functor(..), (<$>))
import Data.Int (Int)
import Data.Map.Strict (Map)
import Data.Maybe (Maybe(..), fromJust)
import Data.Monoid (Monoid(..))
import Data.Ord (Ord(..), Ordering(..))
import Data.Proxy (Proxy(..))
import Data.Semigroup (Semigroup(..))
import Data.String (IsString(..))
import Data.Text (Text)
import Data.Tuple (fst, snd, curry)
import GHC.TypeNats
import Numeric.Natural (Natural)
import Prelude (Integral, fromIntegral, undefined, error, min)
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 qualified Data.Text.Encoding as Text

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 = G
type SecretKey = E
type SecretNonce = E
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 'Signature'
type Signature q = (PublicKey q, Proof q)

-- * 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

validBool :: SubGroup q => ML.MeasuredList 2 (Disjunction q)
validBool = fromJust $ ML.fromList $ List.take 2 groupGenInverses

validRange ::
 forall q mini maxi proxy.
 SubGroup q =>
 Bounds mini maxi ->
 ML.MeasuredList (maxi-mini) (Disjunction q)
validRange 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' 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 n q = ValidityProof (ML.MeasuredList n (Proof q))
 deriving (Eq,Show)

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

proveEncryption ::
 forall ds m r q.
 Nat.Known ds =>
 Monad m => RandomGen r => SubGroup q =>
 PublicKey q -> ZKP ->
 ML.MeasuredList ds (Disjunction q) -> Opinion ds ->
 (SecretNonce q, Encryption q) ->
 S.StateT r m (ValidityProof ds q)
proveEncryption pubKey zkp valids
 (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 @ds -- prove that o+1<=ds implies 1<=ds-o and o<=ds
 , Constraint.Proof <- (Nat.<=) @o @ds     -- prove that o<=ds implies ds-o is a Natural and o+(ds-o) ~ ds
 , Constraint.Proof <- (Nat.<=) @1 @(ds-o) -- prove that ((ds-o)-1)+1 ~ ds-o
 = do
	let (prevDisjs, _indexedDisj:#nextDisjs) = ML.splitAt o valids
	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)
instance Eq (Bounds mini maxi)

-- * 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 validBool)
		 opinions encryptions
	sumProof <-
		proveEncryption pubKey zkp
		 (validRange 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 validBool <$> answer_opinions) &&
	validateEncryption pubKey zkp
	 (validRange 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 :: Text
 }