{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
-- | Finite Field Cryptography (FFC)
-- is a method of implementing discrete logarithm cryptography
-- using finite field mathematics.
module Voting.Protocol.Arithmetic
 ( module Voting.Protocol.Arithmetic
 , Natural
 , Random.RandomGen
 ) where

import Control.Arrow (first)
import Control.Applicative (Applicative(..))
import Control.DeepSeq (NFData)
import Control.Monad (Monad(..), unless)
import Control.Monad.Trans.Reader (ReaderT(..), asks)
import Control.Monad.Trans.Class (MonadTrans(..))
import Data.Aeson (ToJSON(..),FromJSON(..),(.:),(.:?),(.=))
import Data.Bits
import Data.Bool
import Data.Eq (Eq(..))
import Data.Foldable (Foldable, foldl')
import Data.Function (($), (.), id)
import Data.Functor ((<$>))
import Data.Int (Int)
import Data.Maybe (Maybe(..), fromMaybe, fromJust)
import Data.Ord (Ord(..))
import Data.Proxy (Proxy(..))
import Data.Semigroup (Semigroup(..))
import Data.String (IsString(..))
import Data.Text (Text)
import GHC.Generics (Generic)
import GHC.Natural (minusNaturalMaybe)
import Numeric.Natural (Natural)
import Prelude (Integer, Integral(..), fromIntegral, Enum(..))
import Text.Read (readMaybe)
import Text.Show (Show(..))
import qualified Control.Monad.Trans.State.Strict as S
import qualified Crypto.Hash as Crypto
import qualified Data.Aeson as JSON
import qualified Data.Aeson.Types as JSON
import qualified Data.ByteArray as ByteArray
import qualified Data.ByteString as BS
import qualified Data.Char as Char
import qualified Data.List as List
import qualified Data.Text as Text
import qualified Prelude as Num
import qualified System.Random as Random

-- * Type 'FFC'
-- | Mutiplicative Sub-Group of a Finite Prime Field
data FFC = FFC
 {   ffc_name :: Text
 ,   ffc_fieldCharac :: Natural
     -- ^ The prime number characteristic of a Finite Prime Field.
     --
     -- ElGamal's hardness to decrypt requires a large prime number
     -- to form the 'Multiplicative' subgroup.
 ,   ffc_groupGen :: G
     -- ^ A generator of the 'Multiplicative' subgroup of the Finite Prime Field.
     --
     -- NOTE: since 'ffc_fieldCharac' is prime,
     -- the 'Multiplicative' subgroup is cyclic,
     -- and there are phi('fieldCharac'-1) many choices for the generator of the group,
     -- where phi is the Euler totient function.
 ,   ffc_groupOrder :: Natural
     -- ^ The order of the subgroup.
     --
     -- WARNING: 'ffc_groupOrder' MUST be a prime number dividing @('ffc_fieldCharac'-1)@
     -- to ensure that ElGamal is secure in terms of the DDH assumption.
 } deriving (Eq,Show,Generic,NFData)
deriving instance ToJSON FFC
instance FromJSON FFC where
	parseJSON = JSON.withObject "FFC" $ \o -> do
		ffc_name <- o .:? "name"
		ffc_fieldCharac <- o .: "p"
		ffc_groupGen    <- o .: "g"
		ffc_groupOrder  <- o .: "q"
		-- TODO: check p is prime
		-- TODO: check q is prime
		unless (ffc_groupGen < ffc_fieldCharac) $
			JSON.typeMismatch "FFC: groupGen is not lower than fieldCharac" (JSON.Object o)
		unless (ffc_groupOrder < ffc_fieldCharac) $
			JSON.typeMismatch "FFC: groupOrder is not lower than fieldCharac" (JSON.Object o)
		unless (ffc_groupGen > 1) $
			JSON.typeMismatch "FFC: groupGen is not greater than 1" (JSON.Object o)
		unless (fromJust (ffc_fieldCharac`minusNaturalMaybe`one) `rem` ffc_groupOrder == 0) $
			JSON.typeMismatch "FFC: groupOrder does not divide fieldCharac-1" (JSON.Object o)
		return FFC
		 { ffc_name = fromMaybe "" ffc_name
		 , ffc_fieldCharac
		 , ffc_groupGen = G $ F $ ffc_groupGen
		 , ffc_groupOrder
		 }

-- ** Examples
-- | Weak parameters for debugging purposes only.
weakFFC :: FFC
weakFFC = FFC
 { ffc_name        = "weakFFC"
 , ffc_fieldCharac = 263
 , ffc_groupGen    = G $ F 2
 , ffc_groupOrder  = 131
 }

-- | Parameters used in Belenios.
-- A 2048-bit 'fieldCharac' of a Finite Prime Field,
-- with a 256-bit 'groupOrder' for a 'Multiplicative' subgroup
-- generated by 'groupGen'.
beleniosFFC :: FFC
beleniosFFC = FFC
 { ffc_name        = "beleniosFFC"
 , ffc_fieldCharac = 20694785691422546401013643657505008064922989295751104097100884787057374219242717401922237254497684338129066633138078958404960054389636289796393038773905722803605973749427671376777618898589872735865049081167099310535867780980030790491654063777173764198678527273474476341835600035698305193144284561701911000786737307333564123971732897913240474578834468260652327974647951137672658693582180046317922073668860052627186363386088796882120769432366149491002923444346373222145884100586421050242120365433561201320481118852408731077014151666200162313177169372189248078507711827842317498073276598828825169183103125680162072880719
 , ffc_groupGen    = G $ F 2402352677501852209227687703532399932712287657378364916510075318787663274146353219320285676155269678799694668298749389095083896573425601900601068477164491735474137283104610458681314511781646755400527402889846139864532661215055797097162016168270312886432456663834863635782106154918419982534315189740658186868651151358576410138882215396016043228843603930989333662772848406593138406010231675095763777982665103606822406635076697764025346253773085133173495194248967754052573659049492477631475991575198775177711481490920456600205478127054728238140972518639858334115700568353695553423781475582491896050296680037745308460627
 , ffc_groupOrder  = 78571733251071885079927659812671450121821421258408794611510081919805623223441
 }

-- * Type 'F'
-- | The type of the elements of a Finite Prime Field.
--
-- A field must satisfy the following properties:
--
-- * @(f, ('+'), 'zero')@ forms an abelian group,
--   called the 'Additive' group of 'f'.
--
-- * @('NonNull' f, ('*'), 'one')@ forms an abelian group,
--   called the 'Multiplicative' group of 'f'.
--
-- * ('*') is associative:
--   @(a'*'b)'*'c == a'*'(b'*'c)@ and
--   @a'*'(b'*'c) == (a'*'b)'*'c@.
--
-- * ('*') and ('+') are both commutative:
--   @a'*'b == b'*'a@ and
--   @a'+'b == b'+'a@
--
-- * ('*') and ('+') are both left and right distributive:
--   @a'*'(b'+'c) == (a'*'b) '+' (a'*'c)@ and
--   @(a'+'b)'*'c == (a'*'c) '+' (b'*'c)@
--
-- The 'Natural' is always within @[0..'fieldCharac'-1]@.
newtype F = F { unF :: Natural }
 deriving (Eq,Ord,Show)
 deriving newtype NFData
instance ToJSON F where
	toJSON (F x) = JSON.toJSON (show x)
instance Monad m => FromJSON (ReaderT FFC m (Maybe F)) where
	parseJSON (JSON.String s)
	 | Just (c0,_) <- Text.uncons s
	 , c0 /= '0'
	 , Text.all Char.isDigit s
	 , Just x <- readMaybe (Text.unpack s)
	 = return $ do
		fieldCharac <- asks ffc_fieldCharac
		return $
			if x < fieldCharac
			then Just (F x)
			else Nothing
	parseJSON json = JSON.typeMismatch "F" json
instance Monad m => FromNatural (ReaderT FFC m F) where
	fromNatural i = do
		fieldCharac <- asks ffc_fieldCharac
		let abs x | x < 0 = x + fieldCharac
		          | otherwise = x
		return $ F $ abs $ i `mod` fieldCharac
instance ToNatural F where
	nat = unF
instance Monad m => Additive (ReaderT FFC m F) where
	zero = return $ F 0
	mx + my = do
		F x <- mx
		F y <- my
		fieldCharac <- asks ffc_fieldCharac
		return $ F $ (x + y) `mod` fieldCharac
instance Monad m => Negable (ReaderT FFC m F) where
	neg mx = do
		F x <- mx
		if x == 0 then zero
		 else do
			fieldCharac <- asks ffc_fieldCharac
			return $ F $ fromJust $ nat fieldCharac`minusNaturalMaybe`x
instance Monad m => Multiplicative (ReaderT FFC m F) where
	one = return $ F 1
	mx * my = do
		F x <- mx
		F y <- my
		fieldCharac <- asks ffc_fieldCharac
		return $ F $ (x * y) `mod` fieldCharac
{-
instance Monad m => Random.Random (ReaderT FFC m F) where
	randomR (mlo,mhi) = do
		return $ do
		F lo <- mlo
		F hi <- mhi
		fieldCharac <- asks ffc_fieldCharac
		return $
			first (F . fromIntegral) .
			Random.randomR
			 ( 0`max`toInteger lo
			 , toInteger hi`min`(toInteger fieldCharac - 1) )
	random gen = do
		fieldCharac <- asks ffc_fieldCharac
		return $
			first (F . fromIntegral) $
			Random.randomR (0, toInteger fieldCharac - 1) gen
-}

-- ** Class 'Additive'
class Additive a where
	zero :: a
	(+) :: a -> a -> a; infixl 6 +
	sum :: Foldable f => f a -> a
	sum = foldl' (+) zero
instance Additive Natural where
	zero = 0
	(+)  = (Num.+)
instance Additive Integer where
	zero = 0
	(+)  = (Num.+)
instance Additive Int where
	zero = 0
	(+)  = (Num.+)

-- *** Class 'Negable'
class Additive a => Negable a where
	neg :: a -> a
	(-) :: a -> a -> a; infixl 6 -
	x-y = x + neg y
instance Negable Integer where
	neg  = Num.negate
instance Negable Int where
	neg  = Num.negate

-- ** Class 'Multiplicative'
class Multiplicative a where
	one :: a
	(*) :: a -> a -> a; infixl 7 *
instance Multiplicative Natural where
	one = 1
	(*) = (Num.*)
instance Multiplicative Integer where
	one = 1
	(*) = (Num.*)
instance Multiplicative Int where
	one = 1
	(*) = (Num.*)

-- ** Class 'Invertible'
class Multiplicative a => Invertible a where
	inv :: a -> a
	(/) :: a -> a -> a; infixl 7 /
	x/y = x * inv y

-- * Type 'G'
-- | The type of the elements of a 'Multiplicative' subgroup of a Finite Prime Field.
newtype G = G { unG :: F }
 deriving (Eq,Ord,Show)
 deriving newtype NFData
instance ToJSON G where
	toJSON (G x) = JSON.toJSON x
instance Monad m => FromJSON (ReaderT FFC m (Maybe G)) where
	parseJSON (JSON.String s)
	 | Just (c0,_) <- Text.uncons s
	 , c0 /= '0'
	 , Text.all Char.isDigit s
	 , Just x <- readMaybe (Text.unpack s)
	 = return $ do
		fieldCharac <- asks ffc_fieldCharac
		groupOrder  <- asks ffc_groupOrder
		isInvertible <- (==) <$> G (F x) ^ E groupOrder <*> one
		return $
			if x < fieldCharac && isInvertible
			then Just (G (F x))
			else Nothing
	parseJSON json = JSON.typeMismatch "G" json
instance Monad m => FromNatural (ReaderT FFC m G) where
	fromNatural = (G <$>) . fromNatural
instance ToNatural G where
	nat = unF . unG
instance Monad m => Multiplicative (ReaderT FFC m G) where
	one = return $ G $ F one
	x * y = G <$> ((unG <$> x) * (unG <$> y))
instance Monad m => Invertible (ReaderT FFC m G) where
	-- | NOTE: add 'groupOrder' so the exponent given to (^) is positive.
	inv x = do
		groupOrder <- asks ffc_groupOrder
		x >>= (^ E (fromJust $ groupOrder`minusNaturalMaybe`1))

-- | 'groupGenInverses' returns the infinite list
-- of 'inv'erse powers of 'groupGen':
-- @['groupGen' '^' 'neg' i | i <- [0..]]@,
-- but by computing each value from the previous one.
--
-- Used by 'intervalDisjunctions'.
groupGenInverses :: Monad m => ReaderT FFC m [G]
groupGenInverses = do
	invGen <- inv $ asks ffc_groupGen
	let go g = (:) <$> g <*> go (g * return invGen)
	go one

groupGenPowers :: Monad m => ReaderT FFC m [G]
groupGenPowers = do
	let go g = (:) <$> g <*> go (g * asks ffc_groupGen)
	go one

-- | @('hash' bs gs)@ returns as a number in 'E'
-- the SHA256 of the given 'BS.ByteString' 'bs'
-- prefixing the decimal representation of given subgroup elements 'gs',
-- with a comma (",") intercalated between them.
--
-- NOTE: to avoid any collision when the 'hash' function is used in different contexts,
-- a message 'gs' is actually prefixed by a 'bs' indicating the context.
--
-- Used by 'proveEncryption' and 'verifyEncryption',
-- where the 'bs' usually contains the 'statement' to be proven,
-- and the 'gs' contains the 'commitments'.
hash :: Monad m => BS.ByteString -> [G] -> ReaderT FFC m E
hash bs gs = do
	let s = bs <> BS.intercalate (fromString ",") (bytesNat <$> gs)
	let h = ByteArray.convert (Crypto.hashWith Crypto.SHA256 s)
	fromNatural (BS.foldl' (\acc b -> acc`shiftL`3 + fromIntegral b) (0::Natural) h)

-- * Type 'E'
-- | An exponent of a (necessarily cyclic) subgroup of a Finite Prime Field.
-- The value is always in @[0..'groupOrder'-1]@.
newtype E = E { unE :: Natural }
 deriving (Eq,Ord,Show)
 deriving newtype NFData
instance ToJSON E where
	toJSON (E x) = JSON.toJSON x
instance Monad m => FromJSON (ReaderT FFC m (Maybe E)) where
	parseJSON (JSON.String s)
	 | Just (c0,_) <- Text.uncons s
	 , c0 /= '0'
	 , Text.all Char.isDigit s
	 , Just x <- readMaybe (Text.unpack s)
	 = return $ do
		groupOrder <- asks ffc_groupOrder
		return $
			if x < groupOrder
			then Just (E x)
			else Nothing
	parseJSON json = JSON.typeMismatch "E" json

instance Monad m => FromNatural (ReaderT FFC m E) where
	fromNatural i = do
		groupOrder <- asks ffc_groupOrder
		let abs x | x < 0 = x + groupOrder
		          | otherwise = x
		return $ E $ abs $ i `mod` groupOrder
instance ToNatural E where
	nat = unE

instance Monad m => Additive (ReaderT FFC m E) where
	zero = return $ E zero
	mx + my = do
		E x <- mx
		E y <- my
		groupOrder <- asks ffc_groupOrder
		return $ E $ (x + y) `mod` groupOrder
instance Monad m => Negable (ReaderT FFC m E) where
	neg mx = do
		E x <- mx
		if x == 0 then zero
		 else do
			groupOrder <- asks ffc_groupOrder
			return $ E $ fromJust $ nat groupOrder`minusNaturalMaybe`x
instance Monad m => Multiplicative (ReaderT FFC m E) where
	one = return $ E one
	mx * my = do
		E x <- mx
		E y <- my
		groupOrder <- asks ffc_groupOrder
		return $ E $ (x * y) `mod` groupOrder
{-
instance Random.Random (ReaderT FFC m E) where
	randomR (mlo, mhi) = do
		E lo <- mlo
		E hi <- mhi
		groupOrder <- asks ffc_groupOrder
		return $
			first (F . fromIntegral) .
			Random.randomR
			 ( 0`max`toInteger lo
			 , toInteger hi`min`(toInteger groupOrder - 1) )
	random gen = do
		groupOrder <- asks ffc_groupOrder
		return $
			first (E . fromIntegral) $
			Random.randomR (0, toInteger groupOrder - 1) gen
instance Monad m => Enum (ReaderT FFC m E) where
	toEnum = fromNatural . fromIntegral
	fromEnum = fromIntegral . nat
	enumFromTo lo hi = List.unfoldr
	 (\i -> if i<=hi then Just (i, i+one) else Nothing) lo
-}

infixr 8 ^
-- | @(b '^' e)@ returns the modular exponentiation of base 'b' by exponent 'e'.
(^) :: forall m. Monad m => G -> E -> ReaderT FFC m G
(^) b (E e)
 | e == 0 = one
 | otherwise = do
	base <- t * (return b * return b)
	base ^ E (e`shiftR`1)
	where
	t | testBit e 0 = return b
	  | otherwise   = one

-- | @('randomR' i)@ returns a random integer in @[0..i-1]@.
randomR ::
 Monad m =>
 Random.RandomGen r =>
 Random.Random i =>
 Negable i =>
 Multiplicative i =>
 i -> S.StateT r m i
randomR i = S.StateT $ return . Random.randomR (zero, i-one)

-- | @('random')@ returns a random integer
-- in the range determined by its type.
random ::
 Monad m =>
 Random.RandomGen r =>
 Random.Random i =>
 Negable i =>
 Multiplicative i =>
 S.StateT r m i
random = S.StateT $ return . Random.random

instance Random.Random Natural where
	randomR (mini,maxi) =
		first (fromIntegral::Integer -> Natural) .
		Random.randomR (fromIntegral mini, fromIntegral maxi)
	random = first (fromIntegral::Integer -> Natural) . Random.random

-- * Conversions

-- ** Class 'FromNatural'
class FromNatural a where
	fromNatural :: Natural -> a

-- ** Class 'ToNatural'
class ToNatural a where
	nat :: a -> Natural
instance ToNatural Natural where
	nat = id

-- | @('bytesNat' x)@ returns the serialization of 'x'.
bytesNat :: ToNatural n => n -> BS.ByteString
bytesNat = fromString . show . nat