{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE StrictData #-}
{-# OPTIONS_GHC -Wno-missing-signatures #-}
module Hjugement.CLI.Trustee where

import Control.Applicative (Applicative(..))
import Control.Monad (Monad(..), forM, forM_, join, void)
import Control.Monad.Trans.Class (MonadTrans(..))
import Control.Monad.Trans.State.Strict (runState, runStateT)
import Data.Bits (setBit)
import Data.Bool
import Data.ByteString (ByteString)
import Data.Either (Either(..))
import Data.Eq (Eq(..))
import Data.Foldable (Foldable, foldMap, length, null)
import Data.Function (($), (.), id, flip)
import Data.Functor ((<$>))
import Data.Int (Int)
import Data.Maybe (Maybe(..), maybe, fromMaybe)
import Data.Monoid (Monoid(..))
import Data.Ord (Ord(..))
import Data.Proxy (Proxy(..))
import Data.Semigroup (Semigroup(..))
import Data.String (String)
import Data.Text (Text)
import GHC.Natural (minusNatural, minusNaturalMaybe)
import GHC.Prim (coerce)
import Numeric.Natural (Natural)
import Pipes ((>->))
import Prelude (logBase, ceiling, Num(..), (/), (^), fromIntegral, Double)
import Symantic.CLI as CLI
import Text.Show (Show(..))
import System.IO (IO, FilePath)
import qualified Crypto.Hash as Crypto
import qualified Data.Aeson as JSON
import qualified Data.ByteArray as ByteArray
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as BS8
import qualified Data.ByteString.Lazy as BSL
import qualified Data.List as List
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import qualified Data.Text.Lazy as TL
import qualified Data.Text.Lazy.Builder as TLB
import qualified Data.Text.Lazy.Builder.Int as TLB
import qualified Data.Text.Lazy.Encoding as TL
import qualified Data.Text.Lazy.IO as TL
import qualified Data.Time as Time
import qualified Lens.Family as Lens
import qualified Lens.Family.State.Strict as Lens
import qualified Pipes as Pip
import qualified Pipes.ByteString as PipBS
import qualified Pipes.Group as Pip
import qualified Pipes.Prelude as Pip
import qualified Pipes.Aeson as PipJSON (DecodingError(..))
import qualified Pipes.Aeson.Unchecked as PipJSON
import qualified Pipes.Safe as Pip
import qualified Pipes.Safe.Prelude as Pip
import qualified Pipes.Text as PipText
import qualified Pipes.Text.Encoding as PipText
import qualified Pipes.Text.IO as PipText
import qualified Symantic.Document as Doc
import qualified System.FilePath as FP
import qualified System.IO as IO
import qualified System.Posix as Posix
import qualified System.Random as Rand
import qualified Voting.Protocol as VP

import Hjugement.CLI.Utils

-- * trustee
data Trustee_Options = Trustee_Options
 { trustee_election_crypto :: VP.FFC
 } deriving (Show)

api_trustee =
	"Commands for a trustee."
	`helps`
	command "trustee" $
		rule "TrusteeOptions"
		 (Trustee_Options
			 <$> option_crypto
		 ) <?> (
			api_trustee_generate <!>
			api_trustee_decrypt
		 )
		<!> api_help False
run_trustee globOpts =
	(\opts ->
		run_trustee_generate globOpts opts :!:
		run_trustee_decrypt globOpts opts
	 ) :!:
	run_help api_trustee

-- ** generate
api_trustee_generate =
 "Run by a trustee to generate a share of an election key.\
 \ Such a share consists of a private key and a public key with a certificate.\
 \ Generated files are stored in the current directory with\
 \ a name that starts with "<>fileRef "ID"<>",\
 \ where "<>fileRef "ID"<>" is a short fingerprint of the public key.\
 \ The private key is stored in "<>fileRef "ID.privkey"<>" and must be\
 \ secured by the trustee. The public key is stored in "<>fileRef "ID.pubkey"<>" and must\
 \ be sent to the election administrator."
	`helps`
	command "generate" $
		response @()
run_trustee_generate
 Global_Options{..}
 o@Trustee_Options{..} =
	VP.reify trustee_election_crypto $ \(crypto::Proxy c) -> do
		keys@(secKey, pubKey) <- Pip.liftIO $ Rand.getStdRandom $ runState $ do
			secKey <- VP.randomSecretKey @c
			pubKey <- VP.proveIndispensableTrusteePublicKey secKey
			return (secKey, pubKey)
		let pubIdent =
			T.unpack $ T.toUpper $ T.take 8 $
			VP.hexHash $ VP.bytesNat $
			VP.trustee_PublicKey pubKey
		Pip.runSafeT $ Pip.runEffect $ do
			Pip.each [pubIdent] >-> pipeInfo (\ident ->
				"Generated trustee keypair "<>ident<>
				" in "<>(global_dir FP.</> ident)<>".{privkey,pubkey}"
			 ) >-> Pip.drain
			Pip.each [secKey] >-> writeJSON 0o400 (global_dir FP.</> pubIdent FP.<.>"privkey")
			Pip.each [pubKey] >-> writeJSON 0o444 (global_dir FP.</> pubIdent FP.<.>"pubkey")
		return ()

-- ** decrypt
data TrusteeDecrypt_Options = TrusteeDecrypt_Options
 {   trusteeDecrypt_privkey :: FilePath
 ,   trusteeDecrypt_url     :: FilePath
 } deriving (Show)

api_trustee_decrypt =
	"This command is run by each trustee to perform a partial decryption."
	`helps`
	command "decrypt" $
		rule "TrusteeDecryptOptions"
		 (TrusteeDecrypt_Options
			 <$> option_privkey
			 <*> option_url
		 )
		<?> response @(Maybe (VP.DecryptionShare ()))
	where
	option_privkey =
		"Read private key from file "<>ref"FILE"<>"."
		`helps`
		long "privkey" (var "FILE")
	option_url =
		"Download election files from "<>ref"URL"<>"."<>
		"\nDefaults to "<>fileRef "."<>"."
		`helps`
		longOpt "url" "." (var "URL")

run_trustee_decrypt
 Global_Options{..}
 o@Trustee_Options{..}
 TrusteeDecrypt_Options{..} = do
	VP.reify trustee_election_crypto $ \(crypto::Proxy c) -> do
		JSON.eitherDecodeFileStrict' trusteeDecrypt_privkey >>= \case
		 Left err -> do
			outputError $
				Doc.from trusteeDecrypt_privkey<>": "<>
				Doc.from err<>"\n"
		 Right (secKey::VP.E c) -> do
			let pubKey = VP.publicKey secKey
			let trusteeKeysPath = trusteeDecrypt_url FP.</> "public_keys.jsons"
			-- Check that the public key is amongst the public keys of the election
			(keys, ret) <- Pip.runSafeT $ Pip.runEffect $
				Pip.toListM' $
				readJSON trusteeKeysPath
				 >-> Pip.filter ((pubKey ==) . VP.trustee_PublicKey)
			case ret of
			 Left err -> outputError err
			 Right () | null keys -> outputError $
				"the public key associated with the given secret key "<>
				"is not within the list of public trustee keys of the election.\n"<>
				Doc.ul
				 [ "List of trustees' public keys: "<>Doc.from trusteeKeysPath
				 , "Trustee's public key: "<>Doc.from (show (VP.nat pubKey))
				 ]<>"\n"
			 Right () | List.length keys > 1 -> outputError $
				"the public key associated with the given secret key "<>
				"appears more than one time in the list of public trustee keys of the election.\n"<>
				Doc.ul
				 [ "List of trustees' public keys: "<>Doc.from trusteeKeysPath
				 , "Trustee's public key: "<>Doc.from (show (VP.nat pubKey))
				 ]<>"\n"
			 Right () -> do
				-- Tally the encrypted ballots
				-- FIXME: actually support fetching through an URL
				let ballotsPath = trusteeDecrypt_url FP.</> "ballots.jsons"
				((encTally, numBallots), ret) <- Pip.runSafeT $ Pip.runEffect $ do
					Pip.fold'
					 (flip VP.insertEncryptedTally)
					 VP.emptyEncryptedTally id $
						readJSON ballotsPath
				case ret of
				 Left err -> outputError err
				 Right () -> do
					decShare <-
						Rand.getStdRandom $ runState $
							VP.proveDecryptionShare encTally secKey
					return $ Just (coerce decShare :: VP.DecryptionShare ())