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

import Test.Tasty.HUnit
import qualified Data.Text as Text
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 weakFFC
	 ]
 , testGroup "verifyTally" $
	 [ testsVerifyTally weakFFC
	 , testsVerifyTally beleniosFFC
	 ]
 ]

testsVerifyIndispensableTrusteePublicKey :: FFC -> TestTree
testsVerifyIndispensableTrusteePublicKey ffc =
	testGroup (Text.unpack $ ffc_name ffc)
	 [ testVerifyIndispensableTrusteePublicKey ffc 0 (Right ())
	 ]

testVerifyIndispensableTrusteePublicKey ::
 FFC -> Int -> Either ErrorTrusteePublicKey () -> TestTree
testVerifyIndispensableTrusteePublicKey ffc seed exp =
	let got =
		reify ffc $ \(Proxy::Proxy c) ->
		runExcept $
		(`evalStateT` Random.mkStdGen seed) $ do
			trusteeSecKey :: SecretKey c <- randomSecretKey
			trusteePubKey <- proveIndispensableTrusteePublicKey trusteeSecKey
			lift $ verifyIndispensableTrusteePublicKey trusteePubKey
	in
	testCase (Text.unpack $ ffc_name ffc) $
		got @?= exp

testsVerifyTally :: FFC -> TestTree
testsVerifyTally ffc =
	testGroup (Text.unpack $ ffc_name ffc)
	 [ testVerifyTally ffc 0 1 1 1
	 , testVerifyTally ffc 0 2 1 1
	 , testVerifyTally ffc 0 1 2 1
	 , testVerifyTally ffc 0 2 2 1
	 , testVerifyTally ffc 0 5 10 5
	 ]

testVerifyTally :: FFC -> Int -> Natural -> Natural -> Natural -> TestTree
testVerifyTally ffc seed nTrustees nQuests nChoices =
	let clearTallyResult = dummyTallyResult nQuests nChoices in
	let decryptedTallyResult :: Either ErrorTally [[Natural]] =
		reify ffc $ \(Proxy::Proxy c) ->
		runExcept $
		(`evalStateT` Random.mkStdGen seed) $ do
			secKeyByTrustee :: [SecretKey c] <-
				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 "#T=%i,#Q=%i,#C=%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 ::
 Reifies c FFC =>
 Monad m => RandomGen r =>
 PublicKey c -> [[Natural]] -> StateT r m (EncryptedTally c, Natural)
encryptTallyResult pubKey countByChoiceByQuest =
	(`runStateT` 0) $
		forM countByChoiceByQuest $
			mapM $ \count -> do
				modify' $ max count
				(_encNonce, enc) <- lift $ encrypt pubKey (fromNatural count)
				return enc