{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE UndecidableInstances #-} -- for Reifies instances
{-# OPTIONS_GHC -fno-warn-orphans #-}
-- | Finite Field Cryptography (FFC)
-- is a method of implementing discrete logarithm cryptography
-- using finite field mathematics.
module Voting.Protocol.FFC
 ( module Voting.Protocol.FFC
 , Natural
 , Random.RandomGen
 , Reifies(..), reify
 , Proxy(..)
 ) 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.Reflection (Reifies(..), reify)
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.
--
-- NOTE: an 'FFC' term-value is brought into the context of many functions
-- through a type-variable @c@ whose 'Reifies' constraint enables to 'reflect'
-- that 'FFC' at the term-level (a surprising technique but a very useful one).
-- Doing like this is simpler than working in a 'Monad' (like a 'Reader'),
-- and enables that 'FFC' term to be used simply in instances' methods
-- not supporting an inner 'Monad', like 'parseJSON', 'randomR', 'fromEnum' or 'arbitrary'.
-- Aside from that, the sharing of 'FFC' amongst several types
-- is encoded at the type-level by including @c@
-- as a phantom type of 'F', 'G' and 'E'.
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 :: !Natural
     -- ^ 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 <- fromMaybe "" <$> (o .:? "name")
		ffc_fieldCharac <- o .: "p"
		ffc_groupGen    <- o .: "g"
		ffc_groupOrder  <- o .: "q"
		-- TODO: check p is probable prime
		-- TODO: check q is probable prime
		unless (nat 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 (nat 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{..}

fieldCharac :: forall c. Reifies c FFC => Natural
fieldCharac = ffc_fieldCharac (reflect (Proxy::Proxy c))

groupGen :: forall c. Reifies c FFC => G c
groupGen = G $ F $ ffc_groupGen (reflect (Proxy::Proxy c))

groupOrder :: forall c. Reifies c FFC => Natural
groupOrder = ffc_groupOrder (reflect (Proxy::Proxy c))

-- ** Examples
-- | Weak parameters for debugging purposes only.
weakFFC :: FFC
weakFFC = FFC
 { ffc_name        = "weakFFC"
 , ffc_fieldCharac = 263
 , ffc_groupGen    = 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    =  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 c = F { unF :: Natural }
 deriving (Eq,Ord,Show)
 deriving newtype NFData
instance ToJSON (F c) where
	toJSON (F x) = JSON.toJSON (show x)
instance Reifies c FFC => FromJSON (F c) where
	parseJSON (JSON.String s)
	 | Just (c0,_) <- Text.uncons s
	 , c0 /= '0'
	 , Text.all Char.isDigit s
	 , Just x <- readMaybe (Text.unpack s)
	 , x < fieldCharac @c
	 = return (F x)
	parseJSON json = JSON.typeMismatch "F" json
instance Reifies c FFC => FromNatural (F c) where
	fromNatural i = F $ abs $ i `mod` fieldCharac @c
		where
		abs x | x < 0 = x + fieldCharac @c
		      | otherwise = x
instance ToNatural (F c) where
	nat = unF
instance Reifies c FFC => Additive (F c) where
	zero = F 0
	F x + F y = F $ (x + y) `mod` fieldCharac @c
instance Reifies c FFC => Negable (F c) where
	neg (F x)
	 | x == 0 = zero
	 | otherwise = F $ fromJust $ nat (fieldCharac @c)`minusNaturalMaybe`x
instance Reifies c FFC => Multiplicative (F c) where
	one = F 1
	F x * F y = F $ (x * y) `mod` fieldCharac @c
instance Reifies c FFC => Random.Random (F c) where
	randomR (F lo, F hi) =
		first (F . fromIntegral) .
		Random.randomR
		 ( 0`max`toInteger lo
		 , toInteger hi`min`(toInteger (fieldCharac @c) - 1) )
	random =
		first (F . fromIntegral) .
		Random.randomR (0, toInteger (fieldCharac @c) - 1)

-- ** 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 c = G { unG :: F c }
 deriving (Eq,Ord,Show)
 deriving newtype NFData
instance ToJSON (G c) where
	toJSON (G x) = JSON.toJSON x
instance Reifies c FFC => FromJSON (G c) where
	parseJSON (JSON.String s)
	 | Just (c0,_) <- Text.uncons s
	 , c0 /= '0'
	 , Text.all Char.isDigit s
	 , Just x <- readMaybe (Text.unpack s)
	 , x < fieldCharac @c
	 , r <- G (F x)
	 , r ^ E (groupOrder @c) == one
	 = return r
	parseJSON json = JSON.typeMismatch "G" json
instance Reifies c FFC => FromNatural (G c) where
	fromNatural = G . fromNatural
instance ToNatural (G c) where
	nat = unF . unG
instance Reifies c FFC => Multiplicative (G c) where
	one = G $ F one
	G x * G y = G (x * y)
instance Reifies c FFC => Invertible (G c) where
	-- | NOTE: add 'groupOrder' so the exponent given to (^) is positive.
	inv = (^ E (fromJust $ groupOrder @c`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 :: forall c. Reifies c FFC => [G c]
groupGenInverses = go one
	where
	invGen = inv $ groupGen @c
	go g = g : go (g * invGen)

groupGenPowers :: forall c. Reifies c FFC => [G c]
groupGenPowers = go one
	where go g = g : go (g * groupGen @c)

-- | @('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 :: Reifies c FFC => BS.ByteString -> [G c] -> E c
hash bs gs = do
	let s = bs <> BS.intercalate (fromString ",") (bytesNat <$> gs)
	let h = Crypto.hashWith Crypto.SHA256 s
	fromNatural $
		BS.foldl' -- NOTE: interpret the SHA256 as a big-endian number.
		 (\acc b -> acc`shiftL`8 + fromIntegral b)
		 (0::Natural)
		 (ByteArray.convert 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 c = E { unE :: Natural }
 deriving (Eq,Ord,Show)
 deriving newtype NFData
instance ToJSON (E c) where
	toJSON (E x) = JSON.toJSON x
instance Reifies c FFC => FromJSON (E c) where
	parseJSON (JSON.String s)
	 | Just (c0,_) <- Text.uncons s
	 , c0 /= '0'
	 , Text.all Char.isDigit s
	 , Just x <- readMaybe (Text.unpack s)
	 , x < groupOrder @c
	 = return (E x)
	parseJSON json = JSON.typeMismatch "E" json

instance Reifies c FFC => FromNatural (E c) where
	fromNatural i =
		E $ abs $ i `mod` groupOrder @c
		where
		abs x | x < 0 = x + groupOrder @c
		      | otherwise = x
instance ToNatural (E c) where
	nat = unE

instance Reifies c FFC => Additive (E c) where
	zero = E zero
	E x + E y = E $ (x + y) `mod` groupOrder @c
instance Reifies c FFC => Negable (E c) where
	neg (E x)
	 | x == 0 = zero
	 | otherwise = E $ fromJust $ nat (groupOrder @c)`minusNaturalMaybe`x
instance Reifies c FFC => Multiplicative (E c) where
	one = E one
	E x * E y = E $ (x * y) `mod` groupOrder @c
instance Reifies c FFC => Random.Random (E c) where
	randomR (E lo, E hi) =
		first (E . fromIntegral) .
		Random.randomR
		 ( 0`max`toInteger lo
		 , toInteger hi`min`(toInteger (groupOrder @c) - 1) )
	random =
		first (E . fromIntegral) .
		Random.randomR (0, toInteger (groupOrder @c) - 1)
instance Reifies c FFC => Enum (E c) 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'.
(^) :: Reifies c FFC => G c -> E c -> G c
(^) b (E e)
 | e == 0 = one
 | otherwise = t * (b*b) ^ E (e`shiftR`1)
	where
	t | testBit e 0 = 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