{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
module HUnit.Election where

-- import Control.Applicative (Applicative(..))
import qualified Control.Monad.Trans.Except as Exn
import qualified Control.Monad.Trans.State.Strict as S
import qualified Data.List as List
import qualified System.Random as Random

import Protocol.Arithmetic
import Protocol.Credential
import Protocol.Election
import HUnit.Utils

-- * Type 'Params'
class SubGroup q => Params q where
	paramsName :: String
instance Params WeakParams where
	paramsName = "WeakParams"
instance Params BeleniosParams where
	paramsName = "BeleniosParams"

hunit :: TestTree
hunit = testGroup "Election"
 [ testGroup "groupGenInverses"
	 [ testCase "WeakParams" $
		List.take 10 (groupGenInverses @WeakParams) @?=
			[groupGen^neg (fromNatural n) | n <- [0..9]]
	 , testCase "BeleniosParams" $
		List.take 10 (groupGenInverses @BeleniosParams) @?=
			[groupGen^neg (fromNatural n) | n <- [0..9]]
	 ]
 , testGroup "encryptBallot" $
	 [ testsEncryptBallot @WeakParams
	 , testsEncryptBallot @BeleniosParams
	 ]
 ]

testsEncryptBallot :: forall q. Params q => TestTree
testsEncryptBallot =
	testGroup (paramsName @q)
	 [ testEncryptBallot @q 0
		 [Question "q1" ["a1","a2","a3"] zero one]
		 [[True, False, False]]
		 (Right True)
	 , testEncryptBallot @q 0
		 [Question "q1" ["a1","a2","a3"] zero one]
		 [[False, False, False]]
		 (Right True)
	 , testEncryptBallot @q 0
		 [Question "q1" ["a1","a2","a3"] zero one]
		 [[False, False, False]]
		 (Right True)
	 , testEncryptBallot @q 0
		 [Question "q1" [] zero one]
		 []
		 (Left (ErrorBallot_WrongNumberOfAnswers 0 1))
	 , testEncryptBallot @q 0
		 [Question "q1" ["a1","a2"] one one]
		 [[True]]
		 (Left (ErrorBallot_Answer (ErrorAnswer_WrongNumberOfOpinions 1 2)))
	 , testEncryptBallot @q 0
		 [Question "q1" ["a1","a2","a3"] zero one]
		 [[True, True, False]]
		 (Left (ErrorBallot_Answer (ErrorAnswer_WrongSumOfOpinions 2 0 1)))
	 , testEncryptBallot @q 0
		 [Question "q1" ["a1","a2","a3"] one one]
		 [[False, False, False]]
		 (Left (ErrorBallot_Answer (ErrorAnswer_WrongSumOfOpinions 0 1 1)))
	 , testEncryptBallot @q 0
		 [Question "q1" ["a1","a2"] one one]
		 [[False, False, True]]
		 (Left (ErrorBallot_Answer (ErrorAnswer_WrongNumberOfOpinions 3 2)))
	 , testEncryptBallot @q 0
		 [ Question "q1" ["a11","a12","a13"] zero (one+one)
		 , Question "q2" ["a21","a22","a23"] one one
		 ]
		 [ [True, False, True]
		 , [False, True, False] ]
		 (Right True)
	 ]

testEncryptBallot ::
 forall q. SubGroup q =>
 Int -> [Question q] -> [[Bool]] ->
 Either ErrorBallot Bool ->
 TestTree
testEncryptBallot seed quests opins exp =
	let verify =
		Exn.runExcept $
		(`S.evalStateT` Random.mkStdGen seed) $ do
			uuid <- randomUUID
			cred <- randomCredential
			let secKey = secretKey @q uuid cred
			let pubKey = publicKey secKey
			let elec = Election
				 { election_name        = "election"
				 , election_description = "description"
				 , election_publicKey   = pubKey
				 , election_questions   = quests
				 , election_uuid        = uuid
				 , election_hash        = Hash ""
				 }
			verifyBallot elec
			 <$> encryptBallot elec (Just secKey) opins
	in
	testCase (show opins) $
		verify @?= exp