module Test.Pos.Crypto.Gen
        (
        -- Protocol Magic Generator
          genProtocolMagic
        , genProtocolMagicId
        , genRequiresNetworkMagic

        -- Sign Tag Generator
        , genSignTag

        -- Key Generators
        , genKeypair
        , genPublicKey
        , genSecretKey
        , genEncryptedSecretKey

        -- Redeem Key Generators
        , genRedeemKeypair
        , genRedeemPublicKey
        , genRedeemSecretKey
        , genUnitRedeemSignature

        -- VSS Key Generators
        , genVssKeyPair
        , genVssPublicKey

        -- Proxy Cert and Key Generators
        , genProxyCert
        , genProxySecretKey
        , genProxySignature
        , genUnitProxyCert
        , genUnitProxySecretKey
        , genUnitProxySignature

        -- Signature Generators
        , genSignature
        , genSignatureEncoded
        , genSigned
        , genRedeemSignature
        , genUnitSignature

        -- Secret Generators
        , genDecShare
        , genEncShare
        , genSharedSecretData
        , genSecret
        , genSecretProof

        -- Hash Generators
        , genAbstractHash
        , genWithHash
        , genUnitAbstractHash

        -- SafeSigner Generators
        , genSafeSigner

        -- PassPhrase Generators
        , genPassPhrase

        -- HD Generators
        , genHDPassphrase
        , genHDAddressPayload
        ) where

import           Universum

import qualified Data.ByteArray as ByteArray
import           Data.List.NonEmpty (fromList)
import           Hedgehog
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range

import           Crypto.Hash (Blake2b_256)
import           Pos.Binary.Class (Bi)
import           Pos.Crypto (PassPhrase)
import           Pos.Crypto.Configuration (ProtocolMagic (..),
                     ProtocolMagicId (..), RequiresNetworkMagic (..))
import           Pos.Crypto.Hashing (AbstractHash (..), HashAlgorithm, WithHash,
                     abstractHash, withHash)
import           Pos.Crypto.HD (HDAddressPayload (..), HDPassphrase (..))
import           Pos.Crypto.Random (deterministic)
import           Pos.Crypto.SecretSharing (DecShare, EncShare, Secret,
                     SecretProof, VssKeyPair, VssPublicKey, decryptShare,
                     deterministicVssKeyGen, genSharedSecret, toVssPublicKey)
import           Pos.Crypto.Signing (EncryptedSecretKey, ProxyCert,
                     ProxySecretKey, ProxySignature, PublicKey,
                     SafeSigner (..), SecretKey, SignTag (..), Signature,
                     Signed, createPsk, deterministicKeyGen, mkSigned,
                     noPassEncrypt, proxySign, safeCreateProxyCert,
                     safeCreatePsk, sign, signEncoded, toPublic)
import           Pos.Crypto.Signing.Redeem (RedeemPublicKey, RedeemSecretKey,
                     RedeemSignature, redeemDeterministicKeyGen, redeemSign)

----------------------------------------------------------------------------
-- Protocol Magic Generator
----------------------------------------------------------------------------

genProtocolMagic :: Gen ProtocolMagic
genProtocolMagic = ProtocolMagic <$> genProtocolMagicId
                                 <*> genRequiresNetworkMagic

genProtocolMagicId :: Gen ProtocolMagicId
genProtocolMagicId = ProtocolMagicId <$> Gen.int32 Range.constantBounded

genRequiresNetworkMagic :: Gen RequiresNetworkMagic
genRequiresNetworkMagic = Gen.element [RequiresNoMagic, RequiresMagic]

----------------------------------------------------------------------------
-- Sign Tag Generator
----------------------------------------------------------------------------

genSignTag :: Gen SignTag
genSignTag = Gen.element
        [ SignForTestingOnly
        , SignTx
        , SignRedeemTx
        , SignVssCert
        , SignUSProposal
        , SignCommitment
        , SignUSVote
        , SignMainBlock
        , SignMainBlockLight
        , SignMainBlockHeavy
        , SignProxySK
        ]

----------------------------------------------------------------------------
-- Key Generators
----------------------------------------------------------------------------

genKeypair :: Gen (PublicKey, SecretKey)
genKeypair = deterministicKeyGen <$> gen32Bytes

genPublicKey :: Gen PublicKey
genPublicKey =  fst <$> genKeypair

genSecretKey :: Gen SecretKey
genSecretKey = snd <$> genKeypair

genEncryptedSecretKey :: Gen EncryptedSecretKey
genEncryptedSecretKey = noPassEncrypt <$> genSecretKey

----------------------------------------------------------------------------
-- Redeem Key Generators
----------------------------------------------------------------------------

genRedeemKeypair :: Gen (Maybe (RedeemPublicKey, RedeemSecretKey))
genRedeemKeypair = redeemDeterministicKeyGen <$> gen32Bytes

genRedeemPublicKey :: Gen (RedeemPublicKey)
genRedeemPublicKey = do
    rkp <- genRedeemKeypair
    case rkp of
        Nothing      -> error "Error generating a RedeemPublicKey."
        Just (pk, _) -> return pk

genRedeemSecretKey :: Gen (RedeemSecretKey)
genRedeemSecretKey = do
    rkp <- genRedeemKeypair
    case rkp of
        Nothing      -> error "Error generating a RedeemSecretKey."
        Just (_, sk) -> return sk

genUnitRedeemSignature :: Gen (RedeemSignature ())
genUnitRedeemSignature = do pm <- genProtocolMagic
                            genRedeemSignature pm (pure ())

----------------------------------------------------------------------------
-- VSS Key Generators
----------------------------------------------------------------------------

genVssKeyPair :: Gen VssKeyPair
genVssKeyPair =  deterministicVssKeyGen <$> gen32Bytes

genVssPublicKey :: Gen VssPublicKey
genVssPublicKey = toVssPublicKey <$> genVssKeyPair

----------------------------------------------------------------------------
-- Proxy Cert and Key Generators
----------------------------------------------------------------------------

genProxyCert :: Bi w => ProtocolMagic -> Gen w -> Gen (ProxyCert w)
genProxyCert pm genW =
    safeCreateProxyCert pm <$> genSafeSigner <*> genPublicKey <*> genW

genProxySecretKey :: Bi w => ProtocolMagic -> Gen w -> Gen (ProxySecretKey w)
genProxySecretKey pm genW =
    safeCreatePsk pm <$> genSafeSigner <*> genPublicKey <*> genW

genProxySignature
    :: (Bi w, Bi a)
    => ProtocolMagic
    -> Gen a
    -> Gen w
    -> Gen (ProxySignature w a)
genProxySignature pm genA genW = do
    delegateSk  <- genSecretKey
    issuerSk    <- genSecretKey
    w           <- genW
    a           <- genA
    let psk = createPsk pm issuerSk (toPublic delegateSk) w
    return $ proxySign pm SignProxySK delegateSk psk a

genUnitProxyCert :: Gen (ProxyCert ())
genUnitProxyCert = do pm <- genProtocolMagic
                      genProxyCert pm $ pure ()

genUnitProxySecretKey :: Gen (ProxySecretKey ())
genUnitProxySecretKey = do pm <- genProtocolMagic
                           genProxySecretKey pm $ pure ()

genUnitProxySignature :: Gen (ProxySignature () ())
genUnitProxySignature = do
    pm <- genProtocolMagic
    genProxySignature pm (pure ()) (pure ())

----------------------------------------------------------------------------
-- Signature Generators
----------------------------------------------------------------------------

genSignature :: Bi a => ProtocolMagic -> Gen a -> Gen (Signature a)
genSignature pm genA =
    sign pm <$> genSignTag <*> genSecretKey <*> genA

genSignatureEncoded :: Gen ByteString -> Gen (Signature a)
genSignatureEncoded genB =
    signEncoded <$> genProtocolMagic <*> genSignTag <*> genSecretKey <*> genB

genSigned :: Bi a => Gen a -> Gen (Signed a)
genSigned genA =
    mkSigned <$> genProtocolMagic <*> genSignTag <*> genSecretKey <*> genA

genRedeemSignature
    :: Bi a
    => ProtocolMagic
    -> Gen a
    -> Gen (RedeemSignature a)
genRedeemSignature pm genA =
    redeemSign pm <$> gst <*> grsk <*> genA
  where
    gst  = genSignTag
    grsk = genRedeemSecretKey

genUnitSignature :: Gen (Signature ())
genUnitSignature = do pm <- genProtocolMagic
                      genSignature pm (pure ())

----------------------------------------------------------------------------
-- Secret Generators
----------------------------------------------------------------------------

genDecShare :: Gen DecShare
genDecShare = do
    (_, _, xs) <- genSharedSecretData
    case fmap fst (uncons xs) of
        Just (vkp, es) -> return $ deterministic "ds" $ decryptShare vkp es
        Nothing        -> error "Error generating a DecShare."

genEncShare :: Gen EncShare
genEncShare = do
    (_, _, xs) <- genSharedSecretData
    case fmap fst (uncons xs) of
        Just (_, es) -> return es
        Nothing      -> error "Error generating an EncShare."

genSharedSecretData :: Gen (Secret, SecretProof, [(VssKeyPair, EncShare)])
genSharedSecretData = do
    let numKeys = 128 :: Int
    parties     <- Gen.integral (Range.constant 4 (fromIntegral numKeys)) :: Gen Integer
    threshold   <- Gen.integral (Range.constant 2 (parties - 2)) :: Gen Integer
    vssKeyPairs <- replicateM numKeys genVssKeyPair
    let vssPublicKeys = map toVssPublicKey vssKeyPairs
        (s, sp, xs)   = deterministic "ss" $ genSharedSecret threshold (fromList vssPublicKeys)
        ys            = zipWith (\(_, y) x -> (x, y)) xs vssKeyPairs
    return (s, sp, ys)

genSecret :: Gen Secret
genSecret = do
    (s, _, _) <- genSharedSecretData
    return s

genSecretProof :: Gen SecretProof
genSecretProof = do
    (_, sp, _) <- genSharedSecretData
    return sp

----------------------------------------------------------------------------
-- Hash Generators
----------------------------------------------------------------------------

genAbstractHash
    :: (Bi a, HashAlgorithm algo)
    => Gen a
    -> Gen (AbstractHash algo a)
genAbstractHash genA = abstractHash <$> genA

genUnitAbstractHash :: Gen (AbstractHash Blake2b_256 ())
genUnitAbstractHash = genAbstractHash $ pure ()

genWithHash :: Bi a => Gen a -> Gen (WithHash a)
genWithHash genA = withHash <$> genA

----------------------------------------------------------------------------
-- PassPhrase Generators
----------------------------------------------------------------------------

genPassPhrase :: Gen PassPhrase
genPassPhrase = ByteArray.pack <$> genWord8List
  where
    genWord8List :: Gen [Word8]
    genWord8List =
        Gen.list (Range.singleton 32) (Gen.word8 Range.constantBounded)

----------------------------------------------------------------------------
-- SafeSigner Generators
----------------------------------------------------------------------------

genSafeSigner :: Gen SafeSigner
genSafeSigner = Gen.choice gens
  where
    gens = [ SafeSigner <$> genEncryptedSecretKey <*> genPassPhrase
           , FakeSigner <$> genSecretKey
           ]

----------------------------------------------------------------------------
-- HD Generators
----------------------------------------------------------------------------

genHDPassphrase :: Gen HDPassphrase
genHDPassphrase = HDPassphrase <$> gen32Bytes

genHDAddressPayload :: Gen HDAddressPayload
genHDAddressPayload = HDAddressPayload <$> gen32Bytes

----------------------------------------------------------------------------
-- Helper Generators
----------------------------------------------------------------------------

genBytes :: Int -> Gen ByteString
genBytes n = Gen.bytes (Range.singleton n)

gen32Bytes :: Gen ByteString
gen32Bytes = genBytes 32
