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

import Test.Tasty.HUnit
import qualified System.Random as Random
import qualified Text.Printf as Printf

import Voting.Protocol

import Utils

hunit :: TestTree
hunit = testGroup "Indispensable"
 [ testGroup "verifyIndispensableTrusteePublicKey" $
	 [ testsVerifyIndispensableTrusteePublicKey @WeakParams
	 ]
 , testGroup "verifyTally" $
	 [ testsVerifyTally @WeakParams
	 , testsVerifyTally @BeleniosParams
	 ]
 ]

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

testVerifyIndispensableTrusteePublicKey ::
 forall q. Params q =>
 Int -> Either ErrorTrusteePublicKey () -> TestTree
testVerifyIndispensableTrusteePublicKey seed exp =
	let got =
		runExcept $
		(`evalStateT` Random.mkStdGen seed) $ do
			trusteeSecKey :: SecretKey q <- randomSecretKey
			trusteePubKey <- proveIndispensableTrusteePublicKey trusteeSecKey
			lift $ verifyIndispensableTrusteePublicKey trusteePubKey
	in
	testCase (show (paramsName @q)) $
		got @?= exp

testsVerifyTally :: forall q. Params q => TestTree
testsVerifyTally =
	testGroup (paramsName @q)
	 [ testVerifyTally @q 0 1 1 1
	 , testVerifyTally @q 0 2 1 1
	 , testVerifyTally @q 0 1 2 1
	 , testVerifyTally @q 0 2 2 1
	 , testVerifyTally @q 0 5 10 5
	 ]

testVerifyTally ::
 forall q. Params q =>
 Int -> Natural -> Natural -> Natural -> TestTree
testVerifyTally seed nTrustees nQuests nChoices =
	let clearTallyResult = dummyTallyResult nQuests nChoices in
	let decryptedTallyResult :: Either ErrorTally [[Natural]] =
		runExcept $
		(`evalStateT` Random.mkStdGen seed) $ do
			secKeyByTrustee :: [SecretKey q] <-
				replicateM (fromIntegral nTrustees) $ randomSecretKey
			trusteePubKeys <- forM secKeyByTrustee $ proveIndispensableTrusteePublicKey
			let pubKeyByTrustee = trustee_PublicKey <$> trusteePubKeys
			let elecPubKey = combineIndispensableTrusteePublicKeys trusteePubKeys
			(encTally, countMax) <- encryptTallyResult elecPubKey clearTallyResult
			decShareByTrustee <- forM secKeyByTrustee $ proveDecryptionShare encTally
			lift $ verifyDecryptionShareByTrustee encTally pubKeyByTrustee decShareByTrustee
			tally@Tally{..} <- lift $
				proveTally (encTally, countMax) decShareByTrustee $
					combineIndispensableDecryptionShares pubKeyByTrustee
			lift $ verifyTally tally $
				combineIndispensableDecryptionShares pubKeyByTrustee
			return tally_countByChoiceByQuest
	in
	testCase (Printf.printf "nT=%i,nQ=%i,nC=%i (%i maxCount)"
	 nTrustees nQuests nChoices
	 (dummyTallyCount nQuests nChoices)) $
		decryptedTallyResult @?= Right clearTallyResult

dummyTallyCount :: Natural -> Natural -> Natural
dummyTallyCount quest choice = quest * choice

dummyTallyResult :: Natural -> Natural -> [[Natural]]
dummyTallyResult nQuests nChoices =
	[ [ dummyTallyCount q c | c <- [1..nChoices] ]
	| q <- [1..nQuests]
	]

encryptTallyResult ::
 Monad m => RandomGen r => SubGroup q =>
 PublicKey q -> [[Natural]] -> StateT r m (EncryptedTally q, Natural)
encryptTallyResult pubKey countByChoiceByQuest =
	(`runStateT` 0) $
		forM countByChoiceByQuest $
			mapM $ \count -> do
				modify' $ max count
				(_encNonce, enc) <- lift $ encrypt pubKey (fromNatural count)
				return enc