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

-- | Utils for testing Indigo

module Test.Util
  ( testIndigoContract
  , testIndigo
  , testIndigoDoc

  , zeroDivFail
  , notNewKeyFail
  , notNewKeyM
  , negativeResFail
  , negativeResM

  , validateContract
  , validateContractOps
  , validateContractSt
  , validateContractConst

  , validateStSuccess
  , validateStEither
  , validateStack2

  , noOptimizationContract
  ) where

import Fmt (pretty)
import Hedgehog (Gen, MonadTest, PropertyT, annotate, forAll, property, (===))
import Prelude
import Test.HUnit ((@?=))
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.Hedgehog (testProperty)
import Test.Tasty.HUnit (testCase)

import Indigo.Lorentz
import Lorentz.Run (Contract(..))
import Lorentz.Test (ContractPropValidator, contractProp, dummyContractEnv, failedTest)
import Michelson.Interpret (MichelsonFailed(..))
import Michelson.Typed.Haskell.Value (IsoValuesStack)
import Util.IO (readFileUtf8)

type IndigoInstrValidator m pm st out =
  pm -> st -> Either MichelsonFailed (Rec Identity out) -> m ()

--------------------------------------------------------------------------------
-- Tests
--------------------------------------------------------------------------------

-- | Takes a validating function, an Indigo-generated contract and a Lorentz
-- contract and checks that the two have equivalent code (aka made of the same
-- instructions) and that it passes validation for random param and storage.
testIndigoContract
  :: forall pm st.
     ( Show pm, Show st
     , NiceParameterFull pm, NiceStorage st
     )
  => String
  -> Gen pm
  -> Gen st
  -> (pm -> st -> ContractPropValidator (ToT st) (PropertyT IO ()))
  -> ContractCode pm st
  -> FilePath
  -> TestTree
testIndigoContract name genPm genSt propValidator iContract michelsonFile =
  testGroup ("Indigo contract: " <> name)
    [ testCase "matches Michelson reference contract" $ do
        expectedContract <- readFileUtf8 michelsonFile
        printLorentzContract False iContractWithoutOptimization @?= fromStrict expectedContract
    , testProperty "has the correct resulting state and operations" $ property $ do
        pm <- forAll genPm
        st <- forAll genSt
        contProp pm st
    ]
  where
    iContractWithoutOptimization = noOptimizationContract iContract

    contProp :: pm -> st -> PropertyT IO ()
    contProp param storage =
      withDict (niceParameterEvi @pm) $
      contractProp
        (compileLorentzContract iContractWithoutOptimization)
        (propValidator param storage)
        dummyContractEnv param storage

testIndigoDoc
    :: forall pm st.
       String
    -> ContractCode pm st
    -> ContractCode pm st
    -> TestTree
testIndigoDoc name iContract lContract =
  testCase (name <> " matches Lorentz docs content") $
    renderLorentzDoc iContract @?= renderLorentzDoc lContract

-- | Takes a validating function and an Indigo-generated Lorentz `Instr` to check
-- the resulting stack content.
testIndigo
  :: (Show pm, Show st, IsoValue pm, IsoValue st, IsoValuesStack out)
  => String
  -> Gen pm
  -> Gen st
  -> IndigoInstrValidator (PropertyT IO) pm st out
  -> ('[pm, st] :-> out)
  -> TestTree
testIndigo name genPm genSt validator iInstr =
  testProperty (name <> " Indigo Expr has correct resulting stack") $ property $ do
    pm <- forAll genPm
    st <- forAll genSt
    stackProp pm st
  where
    stackProp param storage = validator param storage .
      interpretLorentzInstr dummyContractEnv iInstr $
        Identity param :& Identity storage :& RNil

--------------------------------------------------------------------------------
-- Common failures
--------------------------------------------------------------------------------

zeroDivFail :: MichelsonFailed
zeroDivFail = errorToVal [mt|devision by zero|] MichelsonFailedWith

notNewKeyFail :: MichelsonFailed
notNewKeyFail = errorToVal notNewKeyM MichelsonFailedWith

notNewKeyM :: MText
notNewKeyM = [mt|not new key|]

negativeResFail :: MichelsonFailed
negativeResFail = errorToVal negativeResM MichelsonFailedWith

negativeResM :: MText
negativeResM = [mt|unacceptable negative result|]

--------------------------------------------------------------------------------
-- Contract Validators
--------------------------------------------------------------------------------

-- | Makes a validator for `testIndigoContract` that can expect a failure or a
-- resulting [Operation] from the given function. Ignores new storage value.
validateContract
  :: MonadTest m
  => IsoValue st
  => (pm -> st -> Either MichelsonFailed ([Operation], st))
  -> pm -> st -> ContractPropValidator (ToT st) (m ())
validateContract fn param st (res, _) =
  assertMichelsonResult (fn param st) res $ \(ops, val) (opsRes, resVal) -> do
    annotate "matches resulting Storage and Operations"
    (ops, toVal val) === (opsRes, resVal)


-- | Makes a validator for `testIndigoContract` that can expect a failure or a
-- resulting [Operation] from the given function. Ignores new storage value.
validateContractOps
  :: MonadTest m
  => (pm -> st -> Either MichelsonFailed [Operation])
  -> pm -> st -> ContractPropValidator (ToT st) (m ())
validateContractOps fn param st (res, _) =
  assertMichelsonResult (fn param st) res $ \ops (opsRes, _) -> do
    annotate "matches resulting Operations"
    ops === opsRes

-- | Makes a validator for `testIndigoContract` that can expect a failure or a
-- new storage from the given function. Ignores resulting [Operation]
validateContractSt
  :: MonadTest m
  => IsoValue st
  => (pm -> st -> Either MichelsonFailed st)
  -> pm -> st -> ContractPropValidator (ToT st) (m ())
validateContractSt fn param st (res, _) =
  assertMichelsonResult (fn param st) res $ \val (_, resVal) -> do
    annotate "matches resulting Storage"
    toVal val === resVal

-- | Validator for `testIndigoContract` that expects the storage to remain
-- the same and the resulting [Operation] to be empty. Ignores the parameter.
validateContractConst
  :: MonadTest m
  => IsoValue st
  => pm -> st -> ContractPropValidator (ToT st) (m ())
validateContractConst = validateContract (\_param st -> Right ([], st))

--------------------------------------------------------------------------------
-- Instr Validators
--------------------------------------------------------------------------------

-- | Makes a validator for `testIndigo` that expects the stack not to change type
-- and the "storage" value to have changed as described by the given function.
-- Resulting "param" is ignored.
validateStSuccess
  :: (MonadTest m, Eq st, Show st)
  => (pm -> st -> st)
  -> IndigoInstrValidator m pm st '[pm, st]
validateStSuccess fn = validateStEither (\p s -> Right $ fn p s)

-- | Makes a validator for `testIndigo` that expects the stack not to change type
-- and Either end with a failure or with a new stack. Resulting "param" is ignored.
validateStEither
  :: forall m st pm. (MonadTest m, Eq st, Show st)
  => (pm -> st -> Either MichelsonFailed st)
  -> IndigoInstrValidator m pm st '[pm, st]
validateStEither fn param st res = assertMichelsonResult (fn param st) res checkSt
  where
    checkSt :: st -> (Rec Identity '[pm, st]) -> m ()
    checkSt val resStack = do
      let Identity _ :& Identity newState :& RNil = resStack
      annotate "matches resulting state"
      val === newState

-- | Makes a validator for `testIndigo` that expects the stack to have 3 element
-- (in order and with given values) or a failure to occur.
validateStack2
  :: forall m st pm . (MonadTest m, Eq pm, Eq st, Show pm, Show st)
  => (pm -> st -> Either MichelsonFailed (pm, st))
  -> IndigoInstrValidator m pm st '[pm, st]
validateStack2 fn param st res = assertMichelsonResult (fn param st) res checkSt
  where
    checkSt :: (pm, st) -> (Rec Identity '[pm, st]) -> m ()
    checkSt val resStack = do
      let Identity newParam :& Identity newState :& RNil = resStack
      annotate "matches resulting state"
      val === (newParam, newState)


--------------------------------------------------------------------------------
-- Helpers
--------------------------------------------------------------------------------

assertMichelsonResult
  :: MonadTest m
  => Either MichelsonFailed a
  -> Either MichelsonFailed b
  -> (a -> b -> m ())
  -> m ()
assertMichelsonResult mRes1 mRes2 validatorRight = case (mRes1, mRes2) of
  (Left err, Left e) -> annotate "expected failure" >> err === e
  (Left err, Right _) -> failedTest $ "should have failed with: " <> pretty err
  (Right _, Left e) -> failedTest $ "unexpected failure: " <> pretty e
  (Right val1, Right val2) -> validatorRight val1 val2

noOptimizationContract :: ContractCode param st -> Contract param st
noOptimizationContract code = Contract
  { cCode = code
  , cDisableInitialCast = False
  , cCompilationOptions = noOptimizationOptions
  }

noOptimizationOptions :: CompilationOptions
noOptimizationOptions = defaultCompilationOptions { coOptimizerConf = Nothing }