-- SPDX-FileCopyrightText: 2020 Tocqueville Group
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ

-- | Tests for Indigo Expr

module Test.Expr
  ( test_SmallIndigoExpr
  ) where

import Prelude

import qualified Data.Bits as B
import qualified Data.Map as M
import qualified Data.Set as S
import GHC.Natural (intToNatural, naturalFromInteger, naturalToInt)
import Hedgehog (Gen)
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range
import Test.Tasty (TestTree)

import Cleveland.Util (genTuple2)
import Hedgehog.Gen.Lorentz.UStore (genUStoreFieldExt, genUStoreSubMap)
import Hedgehog.Gen.Michelson (genMText)
import Hedgehog.Gen.Michelson.Typed (genBigMap)
import Hedgehog.Gen.Tezos.Address (genAddress)
import Hedgehog.Gen.Tezos.Core (genChainId, genMutez)
import Hedgehog.Gen.Tezos.Crypto (genKeyHash, genPublicKey, genSignature)
import qualified Indigo as I
import Indigo.Lorentz
import Michelson.Interpret (MichelsonFailed(..), runUnpack)
import Michelson.Interpret.Pack
import Michelson.Runtime.GState (genesisAddress)
import Michelson.Text
import qualified Michelson.Typed as T
import Test.Code.Expr
import Test.Util
import Tezos.Core (dummyChainId, unsafeMkMutez)
import qualified Tezos.Crypto as C

genMyTemplate :: Gen MyTemplate
genMyTemplate = MyTemplate
  <$> genUStoreSubMap (Gen.integral (Range.linearFrom 0 -1000 1000)) (pure ())
  <*> genUStoreFieldExt Gen.bool

genMyUStore :: Gen MyUStore
genMyUStore = mkUStore <$> genMyTemplate

genMySum :: Gen MySum
genMySum = Gen.choice [MySumA <$> Gen.bool, MySumB <$> Gen.integral (Range.linear 0 1000)]

-- | Tests on single Indigo `Expr`s or simple combinations of them.
-- Param and storage for these are generated randomly and their resulting stack
-- is validated against an Haskell function.
test_SmallIndigoExpr :: [TestTree]
test_SmallIndigoExpr =
  [ testIndigo "Cast" genInteger genInteger (validateStSuccess const) (exprUnary @Integer I.cast)
  , testIndigo "Size" genIntegerList genNatural (validateStSuccess (const . intToNatural . length)) exprSize
  , testIndigo "Add" genInteger genInteger (validateStSuccess (+)) (exprBinary @Integer (I.+))
  , testIndigo "Sub" genInteger genInteger (validateStSuccess (-)) (exprBinary @Integer (I.-))
  , testIndigo "Mul" genInteger genInteger (validateStSuccess (*)) (exprBinary @Integer (I.*))
  , testIndigo "Neg" genInteger genInteger (validateStSuccess (const . negate)) (exprUnary @Integer I.neg)
  , testIndigo "Abs" genInteger genNatural (validateStSuccess (\p _ -> naturalFromInteger $ abs p)) exprAbs
  , testIndigo "DivEq" genInteger genInteger (validateStEither divEqCheck) exprDivEq
  , testIndigo "ModNeq" genInteger genInteger (validateStEither modNeqCheck) exprModNeq
  , testIndigo "Le3" genInteger Gen.bool (validateStSuccess (const . (<= 3))) exprLe3
  , testIndigo "Lt3OrGt10" genInteger Gen.bool (validateStSuccess (\p _ -> p < 3 || p > 10)) exprLt3OrGt10
  , testIndigo "Lt3OrGt10" genInteger Gen.bool (validateStSuccess (\p _ -> p >= 3 && p < 10)) exprGe3AndNotGe10
  , testIndigo "Xor" genNatural genNatural (validateStSuccess xor) (exprBinary @Natural (I.^))
  , testIndigo "Lsl" genNatural genShiftNatural (validateStSuccess (\p st -> B.shiftL p (naturalToInt st))) (exprBinary @Natural (I.<<<))
  , testIndigo "Lsr" genNatural genShiftNatural (validateStSuccess (\p st -> B.shiftR p (naturalToInt st))) (exprBinary @Natural (I.>>>))
  , testIndigo "Ge4OrNeq5AndEq6" genInteger Gen.bool (validateStSuccess (\p _ -> p >= 4 || p /= 5 && p == 6)) exprGe4OrNeq5AndEq6
  , testIndigo "Not" Gen.bool Gen.bool (validateStSuccess (\p _ -> not p)) exprNot
  , testIndigo "IsNat" genInteger (Gen.maybe genNatural) (validateStSuccess isNatCheck) exprIsNat
  , testIndigo "Fst" genIntegerPair genInteger (validateStSuccess (\(a,_) _ -> a)) exprFst
  , testIndigo "Snd" genIntegerPair genInteger (validateStSuccess (\(_,b) _ -> b)) exprSnd
  , testIndigo "Some" genInteger genIntegerMaybe (validateStSuccess (\p _ -> Just p)) (exprSome @Integer)
  , testIndigo "None" genInteger genIntegerMaybe (validateStSuccess (\_ _ -> Nothing)) (exprNone @Integer)
  , testIndigo "UStore" genInteger genMyUStore (validateStack2 ustoreCheck) exprUStore

  -- TODO: no `Arbitrary` instance for `Named`
  -- , ToField
  -- , SetField
  -- , Name
  -- , UnName
  -- , Construct
  -- , ConstructT

  , testIndigo "Set" genIntegerSet genInteger (validateStack2 setCheck) exprSet
  , testIndigo "EmptySet" genUnit genIntegerSet (validateStSuccess (\_ _ -> S.empty)) exprEmptySet
  , testIndigo "BigMapLookup" genBigMapInt genIntegerMaybe (validateStSuccess (\(BigMap p) _st -> M.lookup 2 p)) exprBigMapLookup
  , testIndigo "BigMapDelete" genInteger genBigMapInt (validateStSuccess (\p (BigMap st) -> BigMap $ M.delete p st)) exprBigMapDelete
  , testIndigo "BigMapInsert" genInteger genBigMapInt (validateStSuccess (\p (BigMap st) -> BigMap $ M.insert p p st)) exprBigMapInsert
  , testIndigo "Pack" genSignature genByteString (validateStSuccess (\p _ -> packValue' $ T.VSignature p)) exprPack
  , testIndigo "Unpack" genByteString (Gen.maybe genSignature) (validateStSuccess unpackCheck) exprUnpack
  , testIndigo "Cons" genInteger genIntegerList (validateStSuccess (\(p :: Integer) s -> p : s)) exprCons
  , testIndigo "Concat" genMText genMText (validateStSuccess @_ @MText (\p s -> p <> s)) exprConcat
  , testIndigo "Slice" genNatural (Gen.maybe genMText) (validateStSuccess sliceCheck) exprSlice

  -- TODO: Our current testing framework uses storage type for
  -- validation, meaning that we cannot test contracts that way
  -- because we prohibit contract type from appearing in storage.

  -- , Contract
  -- , ConvertEpAddressToContract
  -- , ContractAddress
  -- , Self
  -- , ContractCallingUnsafe
  -- , RunFutureContract
  -- , ImplicitAccount

  , testIndigo "CheckSignature" Gen.bool Gen.bool (validateStSuccess checkSignatureCheck) exprCheckSignature
  , testIndigo "Crypto" genByteString genByteString (validateStack2 cryptoCheck) exprCrypto
  , testIndigo "HashKey" genPublicKey genKeyHash (validateStSuccess (\p _ -> C.hashKey p)) exprHashKey
  , testIndigo "ChainId" genUnit genChainId (validateStSuccess (\_ _ -> dummyChainId)) (exprNullary I.chainId)
  , testIndigo "Amount" genUnit genMutez (validateStSuccess (\_ _ -> unsafeMkMutez 100)) (exprNullary I.amount)
  , testIndigo "Balance" genUnit genMutez (validateStSuccess (\_ _ -> unsafeMkMutez 100)) (exprNullary I.balance)
  , testIndigo "Sender" genUnit genAddress (validateStSuccess (\_ _ -> genesisAddress)) (exprNullary I.sender)

  -- TODO: ContractEnv needed
  -- , Now

  , testIndigo "NonZero" genInteger genIntegerMaybe (validateStSuccess nonZeroCheck) exprNonZero

  , testIndigo "Wrap" Gen.bool genMySum (validateStSuccess wrapCheck) exprWrap
  ]

  where
    genIntegerList = Gen.list (Range.linear 0 100) genInteger
    genIntegerSet = Gen.set (Range.linear 0 100) genInteger
    genIntegerPair = genTuple2 genInteger genInteger
    genIntegerMaybe = Gen.maybe genInteger
    genNatural = Gen.integral @_ @Natural (Range.linear 0 1000)
    genInteger = Gen.integral @_ @Integer (Range.linearFrom 0 -1000 1000)
    genByteString = Gen.bytes (Range.linear 0 100)
    genUnit = pure ()
    genBigMapInt = genBigMap genInteger genInteger

    -- Cannot shift by more than 256 bits
    genShiftNatural = Gen.integral @_ @Natural (Range.linear 0 256)

----------------------------------------------------------------------------
-- Expected behavior
----------------------------------------------------------------------------

divEqCheck :: Integer -> Integer -> Either MichelsonFailed Integer
divEqCheck param st
  | param == 0 = Left zeroDivFail
  | otherwise = Right $ st `div` param

modNeqCheck :: Integer -> Integer -> Either MichelsonFailed Integer
modNeqCheck param st
  | param == 0 = Left zeroDivFail
  | st `mod` param /= 0 = Right 0
  | otherwise = Right 1

isNatCheck :: Integer -> Maybe Natural -> Maybe Natural
isNatCheck param _st
  | param >= 0 = Just (naturalFromInteger param)
  | otherwise = Nothing

unpackCheck :: ByteString -> Maybe Signature -> Maybe Signature
unpackCheck param _st =
  fmap unwrap . rightToMaybe . runUnpack $ param
  where
    unwrap :: Value 'T.TSignature -> Signature
    unwrap (T.VSignature signature) = signature

setCheck :: Set Integer -> Integer -> Either MichelsonFailed (Set Integer, Integer)
setCheck param _st = Right (newParam, newSt)
  where
    newParam
      | S.member 0 param = S.delete 0 param
      | otherwise = S.insert 1 param
    newSt
      | S.size newParam == 1 = 0
      | otherwise = 1

sliceCheck :: Natural -> Maybe MText -> Maybe MText
sliceCheck param (Just st) = Just . takeMText (naturalToInt param) $ st
sliceCheck _param Nothing = Nothing

checkSignatureCheck :: Bool -> Bool -> Bool
checkSignatureCheck _param _st = check sampleSignature
  where
    check SignatureData{..} = C.checkSignature
      (partialParse C.parsePublicKey sdPublicKey)
      (partialParse C.parseSignature sdSignature)
      sdBytes

ustoreCheck
  :: Integer
  -> MyUStore
  -> Either MichelsonFailed (Integer, MyUStore)
ustoreCheck param st
  | param == 0 || M.member 0 stBigMap = Left notNewKeyFail
  | M.member -1 st1BigMap = Right (param, st)
  | otherwise = Right (param, st2)
  where
    myTemplate = either error id $ ustoreDecomposeFull st
    stBigMap = unUStoreSubMap $ ints myTemplate
    st1BigMap = M.insert param () stBigMap
    st2BigMap = M.insert 0 () st1BigMap
    -- st1 = mkUStore $ myTemplate {ints = UStoreSubMap st1BigMap}
    st2 = mkUStore $ myTemplate {ints = UStoreSubMap st2BigMap}

cryptoCheck
  :: ByteString
  -> ByteString
  -> Either MichelsonFailed (ByteString, ByteString)
cryptoCheck param _st = Right (C.sha512 param, C.blake2b param)

nonZeroCheck :: Integer -> Maybe Integer -> Maybe Integer
nonZeroCheck param _st
  | param == 0 = Nothing
  | otherwise = Just param

wrapCheck :: Bool -> MySum -> MySum
wrapCheck param _st = MySumA param