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

import Test.Tasty.HUnit
import qualified Data.List as List
import qualified System.Random as Random

import Protocol.Arithmetic
import Protocol.Credential
import Protocol.Election
import Protocol.Trustees.Simple

import Utils

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
	 ]
 , testGroup "trustee" $
	 [ testsTrustee @WeakParams
	 ]
 ]

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 =
		runExcept $
		(`evalStateT` Random.mkStdGen seed) $ do
			uuid <- randomUUID
			cred <- randomCredential
			let ballotSecKey = credentialSecretKey @q uuid cred
			let elecPubKey = publicKey ballotSecKey -- FIXME: wrong key
			let elec = Election
				 { election_name        = "election"
				 , election_description = "description"
				 , election_publicKey   = elecPubKey
				 , election_questions   = quests
				 , election_uuid        = uuid
				 , election_hash        = Hash "" -- FIXME: when implemented
				 }
			verifyBallot elec
			 <$> encryptBallot elec (Just ballotSecKey) opins
	in
	testCase (show opins) $
		verify @?= exp

testsTrustee :: forall q. Params q => TestTree
testsTrustee =
	testGroup (paramsName @q)
	 [ testTrustee @q 0 (Right ())
	 ]

testTrustee ::
 forall q. SubGroup q =>
 Int -> Either ErrorTrusteePublicKey () -> TestTree
testTrustee seed exp =
	let verify =
		runExcept $
		(`evalStateT` Random.mkStdGen seed) $ do
			trustSecKey <- randomSecretKey @_ @_ @q
			trustPubKey <- proveTrusteePublicKey trustSecKey
			lift $ verifyTrusteePublicKey trustPubKey
	in
	testCase (show seed) $
		verify @?= exp