-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

{-# OPTIONS_HADDOCK not-home #-}

-- | Integration with integrational testing engine (pun intended).
module Test.Cleveland.Internal.Pure
  ( PureM(..)
  , runEmulatedT

  -- * Capability implementations
  , PureState
  , TestError(..)
  , emulatedImpl
  , clevelandOpsImpl
  , clevelandMiscImpl

  -- * Initial environment for Emulated tests
  , initEnv

  -- * Support functions
  , moneybagAlias

  -- * Optics
  , psSecretKeys
  , psDefaultAliasesCounter
  , psRefillableAddresses
  , psNow
  , psLevel
  , psMinBlockTime
  , psGState
  ) where

import Control.Lens (assign, at, makeLenses, modifying, to, (%=), (.=), (?=), (?~))
import Control.Lens.Unsound (lensProduct)
import Control.Monad.Catch.Pure (CatchT, runCatchT)
import Control.Monad.Writer (MonadWriter, WriterT, listen, runWriterT, tell)
import Data.Constraint (Dict(..), withDict, (\\))
import Data.Default (def)
import Data.Map qualified as Map
import Data.Ratio ((%))
import Data.Set qualified as Set
import Fmt (Buildable(..), Builder, build, pretty, unlinesF, (+|), (|+))
import Time (Second, toNum, toUnit)

import Lorentz (Mutez, NiceComparable, pattern DefEpName)
import Lorentz qualified as L
import Lorentz.Entrypoints (HasEntrypointArg, TrustEpName(..), useHasEntrypointArg)
import Morley.AsRPC (HasRPCRepr(AsRPC), notesAsRPC, rpcStorageScopeEvi, valueAsRPC)
import Morley.Client (OperationInfo(..))
import Morley.Michelson.Interpret
  (InterpretError(..), InterpretResult(..), MichelsonFailed(..), MichelsonFailureWithStack(..))
import Morley.Michelson.Runtime hiding (ExecutorOp(..), transfer)
import Morley.Michelson.Runtime qualified as Runtime
import Morley.Michelson.Runtime.Dummy
  (dummyLevel, dummyMaxSteps, dummyMinBlockTime, dummyNow, dummyOrigination)
import Morley.Michelson.Runtime.GState
  (AddressStateFam, GState(..), ImplicitState(..), addressesL, genesisAddress, genesisSecretKey,
  gsChainIdL, gsContractAddressAliasesL, gsContractAddressesL, gsImplicitAddressAliasesL,
  gsVotingPowersL, initGState, lookupBalance)
import Morley.Michelson.TypeCheck (TcError, typeCheckContractAndStorage, typeCheckingWith)
import Morley.Michelson.Typed
  (BigMapId(..), IsoValue, SingI, SomeAnnotatedValue(..), SomeVBigMap(..), ToT, Value, Value'(..),
  castM, fromVal, toVal)
import Morley.Michelson.Typed qualified as T
import Morley.Michelson.Typed.Operation
  (EmitOperation(..), OriginationOperation(..), SetDelegateOperation(..), TransferOperation(..))
import Morley.Michelson.Untyped qualified as U
import Morley.Tezos.Address
import Morley.Tezos.Address.Alias
import Morley.Tezos.Core (Timestamp, timestampPlusSeconds, unsafeSubMutez, zeroMutez)
import Morley.Tezos.Crypto (KeyHash, SecretKey(..), detSecretKey, sign, toPublic)
import Morley.Util.Constrained
import Morley.Util.MismatchError
import Morley.Util.Named

import Morley.Util.Bimap qualified as Bimap
import Test.Cleveland.Internal.Abstract
import Test.Cleveland.Internal.Exceptions
import Test.Cleveland.Lorentz
import Test.Cleveland.Util (ceilingUnit)

data PureState = PureState
  { _psSecretKeys :: Map ImplicitAddress SecretKey
  , _psDefaultAliasesCounter :: DefaultAliasCounter
  , _psRefillableAddresses :: Set ImplicitAddress
  , _psNow :: Timestamp
  , _psLevel :: Natural
  , _psMinBlockTime :: Natural
  , _psGState :: GState
  }
  deriving stock (Eq, Show)

instance MonadState PureState PureM where
  get = ask >>= readIORef
  put = (ask >>=) . flip writeIORef

newtype PureM a = PureM
  { unPureM :: ReaderT (IORef PureState) (CatchT (WriterT LogsInfo IO)) a
  }
  deriving newtype (Functor, Applicative, Monad, MonadIO, MonadThrow, MonadCatch,
                    MonadReader (IORef PureState), MonadWriter LogsInfo, MonadFail)

data TestError
  = UnexpectedTypeCheckError TcError
  | UnexpectedStorageType (MismatchError T.T)
  | UnexpectedBigMapType Natural (MismatchError T.T)
  | CustomTestError Text
  deriving stock Show

makeLenses ''PureState

instance Buildable TestError where
  build = \case
    UnexpectedTypeCheckError tcErr ->
      "Unexpected type check error. Reason: " +| tcErr |+ ""
    UnexpectedStorageType merr ->
      "Unexpected storage type.\n" +| merr |+ ""
    UnexpectedBigMapType bigMapId mismatchError ->
      unlinesF
        [ "A big_map with the ID " +| bigMapId |+ " was found, but it does not have the expected type."
        , build mismatchError
        ]
    CustomTestError msg -> pretty msg

instance Exception TestError where
  displayException = pretty
  fromException = fromPossiblyAnnotatedException

moneybagAlias :: ImplicitAlias
moneybagAlias = ImplicitAlias "moneybag"

runEmulatedT :: ImplicitAlias -> EmulatedT PureM a -> IO a
runEmulatedT moneybagAlias' scenario = do
  let clevelandCaps = ClevelandCaps
        { ccSender = Sender genesisAddress
        , ccMoneybag = Moneybag genesisAddress
        , ccMiscCap = clevelandMiscImpl
        , ccOpsCap = clevelandOpsImpl
        }
      caps = EmulatedCaps emulatedImpl clevelandCaps
  let pureM = runReaderT scenario caps
  env <- newIORef (initEnv moneybagAlias')
  (res, _logs) <- runWriterT $ runCatchT $ runReaderT (unPureM pureM) env
  either throwM pure $ res

emulatedImpl :: EmulatedImpl PureM
emulatedImpl =
  EmulatedImpl
    { eiBranchout = \(scenarios :: [(Text, PureM ())]) ->
        forM_ scenarios $ \(name, scenario) -> do
          aliasesState <- get
          newRef <- newIORef aliasesState
          local (\_ -> newRef) scenario & annotateExceptions (ScenarioBranchName [name])
      , eiGetStorage = addCallStack . exceptionHandler . getStorageImpl
    , eiGetMorleyLogs = getMorleyLogsImpl
    , eiSetVotingPowers = assign (psGState . gsVotingPowersL)
    }

clevelandOpsImpl :: Sender -> ClevelandOpsImpl PureM
clevelandOpsImpl (Sender sender) =
  mapClevelandOpsImplExceptions (addCallStack . exceptionHandler)
    ClevelandOpsImpl
    { coiRunOperationBatch = mapM \case
        OpOriginate UntypedOriginateData{..} -> do
          OpOriginate <$> originate uodContract uodName uodStorage uodBalance uodDelegate
        OpTransfer TransferData{..} -> do
          let fromAddr = #from :! sender
          let toAddr = #to :! toL1Address tdTo
          refillable <- isAddressRefillable sender
          when refillable $ do
            balance <- getBalance sender
            when (balance < tdAmount) $ do
              let moneybag = #from :! genesisAddress
                  toSender = #to :! sender
              void $ transfer @() moneybag toSender (unsafeSubMutez tdAmount balance)
                (TrustEpName DefEpName) ()
          -- Here @toAddr@ is 'Address', so we can not check anything
          -- about it and assume that entrypoint is correct. We pass
          -- unit as contract parameter because it won't be checked
          -- anyway.
          emitOps <- transfer @() fromAddr toAddr tdAmount
            (TrustEpName tdEntrypoint) tdParameter

          return $ OpTransfer $ emitOpToContractEvent <$> emitOps
        OpReveal{} -> do
          -- We do not care about reveals in our Morley runtime
          return $ OpReveal ()
        OpDelegation mbAddress -> OpDelegation <$> setDelegate sender mbAddress
    }

emitOpToContractEvent :: EmitOperation -> ContractEvent
emitOpToContractEvent EmitOperation{eoEmit=T.Emit{..},..} = ContractEvent
  { cePayload = case T.notesT emNotes of
      -- A dirty hack to match reference
      T.TUnit -> Nothing
      _ -> Just $ SomeAnnotatedValue emNotes emValue
  , ceTag = emTag
  , ceSource = eoSource
  }

clevelandMiscImpl :: ClevelandMiscImpl PureM
clevelandMiscImpl =
  mapClevelandMiscImplExceptions (addCallStack . exceptionHandler)
    ClevelandMiscImpl
    { cmiRunIO = \action -> liftIO (try action) >>= \case
        Right res -> pure res
        Left (err :: SomeException) -> throwM err

    , cmiResolveAddress = \case
        a@ImplicitAlias{} -> resolveImplicit a
        a@ContractAlias{} -> resolveContract a

    , cmiSignBytes = \bs addr -> do
        sk <- getSecretKey addr
        liftIO $ sign sk bs

    , cmiGenKey = \sodAlias -> do
      alias <- resolveSpecificOrDefaultAlias sodAlias
      smartGenKey Nothing alias

    , cmiGenFreshKey =
        \sodAlias -> do
          alias <- resolveSpecificOrDefaultAlias sodAlias
          existingAddr <- use (psGState . gsImplicitAddressAliasesL . at alias)
          smartGenKey existingAddr alias

    , cmiOriginateLargeUntyped = originateUntyped

    -- Comments are not supported by integrational testing engine (yet).
    , cmiComment = const pass
    , cmiGetPublicKey = \addr -> do
        toPublic <$> getSecretKey addr
    , cmiGetDelegate = \(Constrained addr) -> addressState addr <&> case addr of
        ContractAddress{} -> csDelegate
        ImplicitAddress{} -> isDelegate
    , cmiGetChainId = use $ psGState . gsChainIdL
    , cmiAdvanceTime = \time -> do
        modifying psNow . flip timestampPlusSeconds $
          toNum @Second @Integer $ ceilingUnit $ toUnit @Second time

    , cmiAdvanceToLevel = \fn ->
        -- do not go back in levels
        modifying psLevel (\cl -> max (fn cl) cl)

    , cmiGetNow = use psNow
    , cmiGetLevel = use psLevel
    , cmiGetApproximateBlockInterval = sec . (% 1) <$> use psMinBlockTime
    , cmiAttempt = try
    , cmiThrow = throwM
    , cmiMarkAddressRefillable = setAddressRefillable
    , cmiGetBalance = foldConstrained getBalance
    , cmiUnderlyingImpl = pure $ Left emulatedImpl
    , cmiFailure = failure
    , ..
    }
  where
    setAddressRefillable addr = psRefillableAddresses %= Set.insert addr

    originateUntyped :: Sender -> UntypedOriginateData 'IsLarge -> PureM ContractAddress
    originateUntyped _ UntypedOriginateData {..} = do
      originate uodContract uodName uodStorage uodBalance uodDelegate

    cmiGetBigMapValueMaybe
      :: forall k v.
         (NiceComparable k, IsoValue v)
      => BigMapId k v
      -> k
      -> PureM (Maybe v)
    cmiGetBigMapValueMaybe (BigMapId bmId) k = runMaybeT do
      -- The RPC does not distinguish between "the bigmap does not exist"
      -- and "the bigmap exists, but the key doesn't", so we mimic the RPC's behaviour here.
      -- We simply return `Nothing` in both cases.
      VBigMap _ bigMap <- findBigMapByIdMaybe @(ToT k) @(ToT v) bmId
      hoistMaybe $ fromVal @v <$> Map.lookup (toVal k) bigMap

    cmiGetAllBigMapValuesMaybe
      :: forall k v.
         (NiceComparable k, IsoValue v)
      => BigMapId k v
      -> PureM (Maybe [v])
    cmiGetAllBigMapValuesMaybe (BigMapId bmId) = runMaybeT do
      VBigMap _ bigMap <- findBigMapByIdMaybe @(ToT k) @(ToT v) bmId
      pure $ fromVal @v <$> Map.elems bigMap

    -- In a real chain, when we retrieve a contract's storage via the Tezos RPC,
    -- the storage expression will have all the big_maps replaced with their respective big_map IDs.
    --
    -- Here, we mimic the RPC's behaviour.
    --
    -- We expect all big_maps in the storage to already have an ID.
    -- IDs are assigned to big_maps by the interpreter/runtime when:
    --   * A contract with big_maps in its storage is originated
    --   * A transfer is made and the parameter contains big_maps
    --   * A contract's code is run and it calls `EMPTY_BIG_MAP`, `DUP` or `DUP n`.
    cmiGetSomeStorage :: ContractAddress -> PureM SomeAnnotatedValue
    cmiGetSomeStorage addr = do
      ContractState _ contract (storage :: Value t) _ <- addressState addr
      pure $
        SomeAnnotatedValue
          (notesAsRPC $ T.cStoreNotes contract)
          (valueAsRPC storage)
          \\ rpcStorageScopeEvi @t

    -- Generate a fresh address for a given alias.
    --
    -- If this alias is not yet associated with any address,
    -- we use the alias as the seed for generating a brand new address.
    --
    -- If this alias is already associated with an address,
    -- then we have to use a _different_ seed to guarantee we'll get a brand new address.
    -- Therefore, we concatenate the alias with the existing address and use the
    -- result as a seed for generating a brand new address.
    smartGenKey :: Maybe ImplicitAddress -> ImplicitAlias -> PureM ImplicitAddress
    smartGenKey existingAddr alias@(ImplicitAlias aliasTxt) = do
      let
        seed = maybe aliasTxt (mappend aliasTxt . pretty) existingAddr
        sk = detSecretKey (encodeUtf8 seed)
        addr = detGenKeyAddress (encodeUtf8 seed)

      -- Save alias/address association.
      psGState . gsImplicitAddressAliasesL . at alias ?= addr
      -- Save the address's secret key.
      psSecretKeys . at addr .= Just sk
      pure addr

    resolveSpecificOrDefaultAlias (SpecificAlias alias) =
      return alias
    resolveSpecificOrDefaultAlias DefaultAlias = do
      DefaultAliasCounter counter <- use psDefaultAliasesCounter
      psDefaultAliasesCounter %= \(DefaultAliasCounter i) -> DefaultAliasCounter $ i + 1
      return $ mkDefaultAlias counter

    cmiRunCode
      :: forall cp st vd. (HasRPCRepr st, T.IsoValue (AsRPC st))
      => Sender -> RunCode cp st vd -> PureM (AsRPC st)
    cmiRunCode (Sender sender) (RunCode rcContract rcParameter rcStorage rcAmount rcLevel rcNow rcBalance rcSource) = do
      contract@T.Contract{..} <- pure $ L.toMichelsonContract rcContract
      T.MkEntrypointCallRes _ (epc :: T.EntrypointCallT (T.ToT cp) arg) <- pure $
        T.mkDefEntrypointCall cParamNotes

      (now, level) <- use $ psNow `lensProduct` psLevel
      knownContracts <- use $ psGState . gsContractAddressesL
      bigMapFinder <- registerInterpretation Runtime.mkBigMapFinder
      let tcBm :: forall (t :: T.T). T.SingI t => U.Value -> PureM (T.Value t)
          tcBm = either (throwM . UnexpectedTypeCheckError) pure .
            resolveRunCodeBigMaps bigMapFinder
      rcParameterT <- tcBm rcParameter
      rcStorageT <- tcBm rcStorage
      (_, finalStorage) <- either (throwEE . EEInterpreterFailed (toAddress sender)) pure $
        Runtime.runCode
          (Runtime.runCodeParameters contract rcStorageT epc rcParameterT)
            { Runtime.rcAmount = rcAmount
            , Runtime.rcLevel = fromMaybe level rcLevel
            , Runtime.rcNow = fromMaybe now rcNow
            , Runtime.rcBalance = rcBalance
            , Runtime.rcSource = Constrained $ fromMaybe genesisAddress rcSource
            , Runtime.rcKnownContracts = knownContracts
            , Runtime.rcSender = Constrained sender
            }
      pure $ fromVal $ valueAsRPC finalStorage

-- | Traverse storage values of all contracts and looks for a big_map with the given ID.
findBigMapByIdMaybe
  ::forall k v. (SingI v, SingI k)
  => Natural -> MaybeT PureM (Value ('T.TBigMap k v))
findBigMapByIdMaybe bigMapId = do
  SomeVBigMap (v@VBigMap{} :: Value t) <- MaybeT $
    registerInterpretation Runtime.mkBigMapFinder <*> pure bigMapId
  castM @t @('T.TBigMap k v) v $ throwM . UnexpectedBigMapType bigMapId

isAddressRefillable :: ImplicitAddress -> PureM Bool
isAddressRefillable addr = Set.member addr <$> use psRefillableAddresses

getBalance :: L1AddressKind kind => KindedAddress kind -> PureM Mutez
getBalance addr = do
  gs <- use psGState
  pure $ fromMaybe zeroMutez $ lookupBalance addr gs

exceptionHandler :: PureM a -> PureM a
exceptionHandler action = try action >>= \case
  Left err -> exceptionToTransferFailure err >>= throwM
  Right res -> return res
  where
    exceptionToTransferFailure :: ExecutorError' AddressAndAlias -> PureM TransferFailure
    exceptionToTransferFailure err = case err of
      EEZeroTransaction addr -> return $
        TransferFailure addr EmptyTransaction
      EEIllTypedParameter addr _ -> return $
        TransferFailure addr BadParameter
      EEUnexpectedParameterType addr _ -> return $
        TransferFailure addr BadParameter
      EEInterpreterFailed addr (InterpretError (MichelsonFailureWithStack{..}, _)) ->
        case mfwsFailed of
          MichelsonFailedWith val -> return $
            TransferFailure addr $
              FailedWith (EOTVTypedValue val) (Just mfwsErrorSrcPos)
          MichelsonArithError (T.ShiftArithError{}) -> return $
            TransferFailure addr ShiftOverflow
          MichelsonArithError (T.MutezArithError errType _ _) -> return $
            TransferFailure addr $ MutezArithError errType
          MichelsonGasExhaustion -> return $
            TransferFailure addr GasExhaustion
          _ -> throwM err
      _ -> throwM err


getMorleyLogsImpl :: PureM a -> PureM (LogsInfo, a)
getMorleyLogsImpl action = swap <$> listen action

getStorageImpl
  :: forall st addr. (ToStorageType st addr)
  => addr -> PureM st
getStorageImpl addr = do
  withDict (pickNiceStorage @st addr) $ do
    ContractState _ _ (storage :: Value actualT) _ <- addressState (toContractAddress addr)
    val <- castM @actualT @(ToT st) storage (throwM . UnexpectedStorageType)
    pure $ T.fromVal val

addressState :: KindedAddress kind -> PureM (AddressStateFam kind)
addressState addr = maybe (unknownAddress addr) pure =<<
  use (psGState . addressesL addr . at addr)

resolveImplicit :: ImplicitAlias -> PureM ImplicitAddress
resolveImplicit alias = do
  use (psGState . gsImplicitAddressAliasesL . at alias) >>= maybe (unknownAlias alias) pure

resolveContract :: ContractAlias -> PureM ContractAddress
resolveContract alias = do
  use (psGState . gsContractAddressAliasesL . at alias) >>= maybe (unknownAlias alias) pure

unknownAddress :: KindedAddress kind -> PureM whatever
unknownAddress =
  throwM . CustomTestError .
  mappend "Unknown address provided: " . pretty

unknownAlias :: Alias kind -> PureM whatever
unknownAlias =
  throwM . CustomTestError .
  mappend "Unknown address alias: " . pretty

failure :: forall a. Builder -> PureM a
failure = throwM . CustomTestError . pretty

getSecretKey :: ImplicitAddress -> PureM SecretKey
getSecretKey addr = do
  use (psSecretKeys . at addr) >>= \case
    Nothing -> unknownAddress addr
    Just sk -> pure sk

----------------------------------------------------------------------------
-- Support functions
----------------------------------------------------------------------------

initEnv :: ImplicitAlias -> PureState
initEnv alias = PureState
  { _psSecretKeys = one (genesisAddress, genesisSecretKey)
  , _psDefaultAliasesCounter = DefaultAliasCounter 0
  , _psRefillableAddresses = Set.empty
  , _psNow = dummyNow
  , _psLevel = dummyLevel
  , _psGState = initGState & gsImplicitAddressAliasesL . at alias ?~ genesisAddress
  , _psMinBlockTime = dummyMinBlockTime
  }

----------------------------------------------------------------------------
-- Emulator internals
----------------------------------------------------------------------------

transfer
  :: forall cp vd epRef epArg addr.
     (HasEntrypointArg cp epRef epArg, T.IsoValue epArg, L.ToTAddress cp vd addr)
  => "from" :! ImplicitAddress
  -> "to" :! addr
  -> Mutez
  -> epRef
  -> epArg
  -> PureM [EmitOperation]
transfer (arg #from -> from) (arg #to -> toAddr) money epRef param =
  let L.TAddress to' = L.toTAddress @cp @vd toAddr in
  case useHasEntrypointArg @cp @epRef @epArg epRef of
    (Dict, epName) -> registerInterpretation do
      executeGlobalOperations def $ one $
        Runtime.TransferOp $ TransferOperation
          { toDestination = to'
          , toCounter = 0
          , toTxData = TxData
            { tdSenderAddress = Constrained from
            , tdParameter = TxTypedParam (T.toVal param)
            , tdEntrypoint = epName
            , tdAmount = money
            }
          }

-- | Originate a contract with given initial storage and balance. Its
-- address is returned.
originate :: U.Contract -> ContractAlias -> U.Value -> Mutez -> Maybe KeyHash -> PureM ContractAddress
originate uContract alias uStorage balance delegate =
  case typeCheckingWith def $ typeCheckContractAndStorage uContract uStorage of
    Left tcErr -> throwM $ UnexpectedTypeCheckError tcErr
    Right (T.SomeContractAndStorage contract storage) -> do
      counter <- gsCounter <$> use psGState
      registerInterpretation . executeGlobalOrigination $
        (dummyOrigination storage contract counter)
          { ooBalance = balance
          , ooAlias = Just alias
          , ooDelegate = delegate
          }

throwEE :: ExecutorError -> PureM a
throwEE err =
  -- Replace all `Address`es with `AddressAndAlias` in the error and throw it.
  traverse addrToAddressAndAlias err >>= throwM
  where
    addrToAddressAndAlias :: Address -> PureM AddressAndAlias
    addrToAddressAndAlias (MkAddress kindedAddr) =
      case kindedAddr of
        ContractAddress{} ->
          use $ psGState . gsContractAddressAliasesL . Bimap.flipped . at kindedAddr . to (AddressAndAlias kindedAddr)
        ImplicitAddress{} ->
          use $ psGState . gsImplicitAddressAliasesL . Bimap.flipped . at kindedAddr . to (AddressAndAlias kindedAddr)
        _ ->
          pure $ AddressAndAlias kindedAddr Nothing

-- | Runs a set of operations and updates the engine's state.
registerInterpretation :: ExecutorM a -> PureM a
registerInterpretation action = do
  interpretedResult <- interpret action
  tell $ extractLogs interpretedResult

  case interpretedResult of
    Right (executorRes, res) -> do
      psGState .= _erGState executorRes
      pure res
    Left executorError -> throwEE executorError

  where
    extractLogs :: Either ExecutorError (ExecutorRes, a) -> [ScenarioLogs]
    extractLogs = \case
      Left (EEInterpreterFailed addr (InterpretError e)) -> [ScenarioLogs addr $ snd e]
      Right (res, _) -> res ^. erInterpretResults <&> \(addr, InterpretResult{..}) ->
        ScenarioLogs addr iurMorleyLogs
      _ -> []

-- | Interpret an action and return the result _without_ updating the engine's state.
interpret :: ExecutorM a -> PureM (Either ExecutorError (ExecutorRes, a))
interpret action = do
  now <- use psNow
  level <- use psLevel
  gState <- use psGState
  minBlockTime <- use psMinBlockTime
  pure $ runExecutorM now level minBlockTime dummyMaxSteps gState action

setDelegate :: ImplicitAddress -> Maybe KeyHash -> PureM ()
setDelegate addr mbKh = void $ registerInterpretation $
  executeGlobalOperations def $ one $
    Runtime.SetDelegateOp SetDelegateOperation
      { sdoContract = Constrained addr
      , sdoDelegate = mbKh
      , sdoCounter = 0
      }