-- 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 qualified Indigo as I
import Indigo.Lorentz
import Lorentz.UStore.Types (genUStoreFieldExt, genUStoreSubMap)
import Michelson.Interpret (MichelsonFailed(..), runUnpack)
import Michelson.Interpret.Pack
import Michelson.Runtime.GState (genesisAddress)
import Michelson.Test.Gen (genMutez)
import Michelson.Test.Util (genTuple2)
import Michelson.Text
import Michelson.Typed (genBigMap)
import qualified Michelson.Typed as T
import Test.Code.Expr
import Test.Util
import Tezos.Address (genAddress)
import Tezos.Core
import Tezos.Crypto (genKeyHash, genPublicKey, genSignature)
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

-- | 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
  ]

  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