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

module Morley.Michelson.TypeCheck.TypeCheck
  ( TcInstrHandler
  , TcOriginatedContracts
  , TypeCheckEnv(..)
  , BigMapFinder
  , TypeCheckOptions(..)
  , TypeCheck
  , TypeCheckNoExcept
  , TypeCheckResult
  , runTypeCheck
  , TypeCheckInstr
  , TypeCheckInstrNoExcept
  , runTypeCheckIsolated
  , typeCheckingWith
  , liftNoExcept
  , throwingTCError
  , preserving
  , preserving'
  , guarding
  , guarding_
  , tcEither
  , ask'
  , asks'
  , local'

  , tcModeL
  , TypeCheckMode(..)
  , SomeParamType(..)
  , mkSomeParamType
  , TypeCheckInstrEnv (..)
  , tcieErrorPos
  , tcieNotInView
  ) where

import Control.Lens (makeLenses)
import Control.Monad.Except (Except, MonadError, runExcept, throwError)
import Data.Constraint (Dict)
import Data.Default (Default(..))
import Fmt (Buildable, build, pretty)

import Morley.Michelson.ErrorPos (ErrorSrcPos)
import Morley.Michelson.TypeCheck.Error (TCError(..), TCTypeError(..))
import Morley.Michelson.TypeCheck.TypeCheckedSeq
  (IllTypedInstr(..), TypeCheckedInstr, TypeCheckedOp(..), TypeCheckedSeq(..), someInstrToOp,
  tcsEither)
import Morley.Michelson.TypeCheck.Types
import Morley.Michelson.Typed qualified as T
import Morley.Michelson.Typed.Existential (SomeVBigMap)
import Morley.Michelson.Untyped qualified as U
import Morley.Tezos.Address (ContractHash)
import Morley.Util.Lens
import Morley.Util.MultiReader

-- | Environments avaliable during typecheck.
type TCEnvs = '[TypeCheckEnv, TypeCheckOptions]

-- | A full type-check monad carrying intermediary context (via 'TypeCheckEnv'),
-- general 'TypeCheckOptions' and throwing 'TCError'.
type TypeCheck = MultiReaderT TCEnvs (Except TCError)

-- | A non-throwing alternative for 'TypeCheck'. Mainly meant to be used for
-- construction of a partially typed tree (see 'TypeCheckedSeq').
type TypeCheckNoExcept = MultiReaderT TCEnvs Identity

-- | Monad for performing some typechecking operations with the same options.
--
-- Unlike 'TypeCheck' monad, this does not carry the context of intra-contract
-- or intra-value typechecking.
type TypeCheckResult = ReaderT TypeCheckOptions (Except TCError)

-- | Environments available during instr typecheck
type TCInstrEnvs = TypeCheckInstrEnv ': TCEnvs

-- | Version of 'TypeCheck' additionally carrying instruction-specific
-- 'TypeCheckInstrEnv'
type TypeCheckInstr = MultiReaderT TCInstrEnvs (Except TCError)

-- | Version of 'TypeCheckNoExcept' additionally carrying instruction-specific
-- 'TypeCheckInstrEnv'
type TypeCheckInstrNoExcept = MultiReaderT TCInstrEnvs Identity

data SomeParamType = forall t. (T.ParameterScope t) =>
  SomeParamType (T.ParamNotes t)

-- | @Show@ instance of @SomeParamType@, mainly used in test.
deriving stock instance Show SomeParamType

-- | @Eq@ instance of @SomeParamType@, mainly used in test.
instance Eq SomeParamType where
  s1 == s2 = someParamToParameterType s1 == someParamToParameterType s2

-- | @Buildable@ instance of @SomeParamType@, mainly used in test.
instance Buildable SomeParamType where
  build = build . someParamToParameterType

-- | Helper function means to provide a quick way for creating instance
-- of @SomeParamType@ needed in test.
someParamToParameterType :: SomeParamType -> U.ParameterType
someParamToParameterType (SomeParamType T.UnsafeParamNotes{..}) =
  U.ParameterType (T.AsUType pnNotes) pnRootAnn

mkSomeParamType :: U.ParameterType -> Either TCError SomeParamType
mkSomeParamType (U.ParameterType t ann) =
  T.withUType t $ \(notescp :: T.Notes t) -> do
    case T.checkScope @(T.ParameterScope t) of
      Right T.Dict ->
        case T.mkParamNotes notescp ann of
          Right paramNotes -> Right $ SomeParamType paramNotes
          Left err ->
            Left $ TCContractError "invalid parameter declaration: " $ Just $ IllegalParamDecl err
      Left err -> Left $ TCContractError ("Parameter type is invalid: " <> pretty err) Nothing

type TcOriginatedContracts = Map ContractHash SomeParamType

-- | Typechecking mode that tells the type checker whether it is typechecking
-- contract code in actual contract, lambda, or test.
data TypeCheckMode
  = TypeCheckValue
    -- ^ We're typechecking a value.
      (U.Value, T.T)
      (Maybe BigMapFinder)
      -- ^ When this is a `Just`, we simulate the typechecking behavior of the RPC's @/run_code@ endpoint.
      --
      -- If an integer is found where a big_map is expected,
      -- we use 'BigMapFinder' to check if a big_map exists with that ID.
      -- If it does, and if the big_map's value and key have the expected types, we replace the
      -- big_map ID with the corresponding big_map value.
  | TypeCheckContract SomeParamType
    -- ^ We're typechecking a contract.
  | TypeCheckTest
    -- ^ We're typechecking a set of instructions in a "test" environment like a REPL,
    -- where the instruction @SELF@ would not make sense.

-- | The typechecking state
data TypeCheckEnv = TypeCheckEnv
  { tcMode :: ~TypeCheckMode
  }

type BigMapFinder = Natural -> Maybe SomeVBigMap

data TypeCheckOptions = TypeCheckOptions
  { tcVerbose :: Bool
    -- ^ Whether to add stack type comments after every
    -- instruction a la tezos-client.
  , tcStrict :: Bool
    -- ^ Whether should we behave like in test run or real run
    -- (real run is more strict).
    --
    -- @tezos-client run@'s behaviour can slightly differ from the behaviour
    -- of @tezos-client originate@ and @tezos-client transfer@. For instance,
    -- some values can be "forged" in test run, but not in a real one,
    -- see: Note [Tickets forging].
    --
    -- Set this to @True@ when need to match the behaviour in the network,
    -- and to @False@ if you prefer providing the user with some convenient
    -- features.
  }

instance Default TypeCheckOptions where
  def = TypeCheckOptions{ tcVerbose = False, tcStrict = True }

data TypeCheckInstrEnv = TypeCheckInstrEnv
  { _tcieErrorPos :: ErrorSrcPos
  , _tcieNotInView :: Maybe (Dict T.IsNotInView)
  }

makeLensesWith postfixLFields ''TypeCheckEnv
makeLenses ''TypeCheckInstrEnv

runTypeCheck :: TypeCheckMode -> TypeCheck a -> TypeCheckResult a
runTypeCheck = usingReaderT . TypeCheckEnv

-- | Run type checker as if it worked isolated from other world -
-- no access to environment of the current contract is allowed.
--
-- Use this function for test purposes only or for some utilities when
-- environment does not matter. In particular, it is assumed that
-- whatever we typecheck does not depend on the parameter type of the
-- contract which is being typechecked (because there is no contract
-- that we are typechecking).
runTypeCheckIsolated :: TypeCheck (SomeInstr t) -> TypeCheckResult (SomeInstr t)
runTypeCheckIsolated = runTypeCheck TypeCheckTest

typeCheckingWith :: TypeCheckOptions -> TypeCheckResult a -> Either TCError a
typeCheckingWith options = runExcept . usingReaderT options

instance Default TypeCheckInstrEnv where
  def = TypeCheckInstrEnv def Nothing

liftNoExcept
  :: forall m e (a :: Type). MonadMultiReaderT m Identity
  => m a
  -> ChangeMultiReaderBase m (Except e) a
liftNoExcept action = mapMultiReaderT (pure @(Except e) . runIdentity) action

throwingTCError
  :: ( MonadMultiReaderT m Identity
     , m' ~ ChangeMultiReaderBase m (Except TCError)
     , MonadError TCError m'
     )
  => m (TypeCheckedSeq inp) -> m' (SomeInstr inp)
throwingTCError action = liftNoExcept @_ @TCError action >>= tcsEither (const throwError) (pure)

tcEither
  :: (TCError -> TypeCheckInstrNoExcept a) -- ^ Call this if the action throws
  -> (b -> TypeCheckInstrNoExcept a) -- ^ Call this if it doesn't
  -> TypeCheckInstr b -- ^ The action to perform
  -> TypeCheckInstrNoExcept a -- ^ A non-throwing action
tcEither onErr onOk action = either onErr onOk =<< mapMultiReaderT (pure . runExcept) action

-- | Perform a throwing action on an acquired instruction. Preserve the acquired
-- result by embedding it into a type checking tree with a specified parent
-- instruction.
preserving
  :: TypeCheckInstrNoExcept (TypeCheckedSeq inp)
  -- ^ Acquiring computation
  -> ([TypeCheckedOp] -> TypeCheckedInstr)
  -- ^ The parent instruction constructor
  -> (SomeInstr inp -> TypeCheckInstr (SomeInstr inp'))
  -- ^ The throwing action
  -> TypeCheckInstrNoExcept (TypeCheckedSeq inp')
preserving acquire con action = preserving' acquire con
  (\instr -> action instr & tcEither
    (\err -> pure $ IllTypedSeq err [SemiTypedInstr $ con [someInstrToOp instr]])
    (pure . WellTypedSeq))

-- | Perform a non-throwing action on an acquired instruction. Preserve the
-- acquired result even if the action does not succeed. Embed the result into a
-- type checking tree with a specified parent instruction.
preserving'
  :: TypeCheckInstrNoExcept (TypeCheckedSeq inp)
  -- ^ Acquiring computation
  -> ([TypeCheckedOp] -> TypeCheckedInstr)
  -- ^ The parent instruction constructor
  -> (SomeInstr inp -> TypeCheckInstrNoExcept (TypeCheckedSeq inp'))
  -- ^ The action
  -> TypeCheckInstrNoExcept (TypeCheckedSeq inp')
preserving' acquire con action =
  acquire >>= tcsEither
    (\tcOps err -> pure $ IllTypedSeq err [SemiTypedInstr $ con tcOps])
    (action)

-- | Acquire a resource. If successfully, call a follow-up action on it,
-- otherwise embed the error into a type checking tree along with a specified
-- untyped instruction.
guarding
  :: U.ExpandedInstr -- ^ Untyped instruction
  -> TypeCheckInstr a -- ^ Acquiring computation
  -> (a -> TypeCheckInstrNoExcept (TypeCheckedSeq inp)) -- ^ Follow-up action
  -> TypeCheckInstrNoExcept (TypeCheckedSeq inp)
guarding instr cond action = do
  cond & tcEither
    (\err -> pure $ IllTypedSeq err [NonTypedInstr $ U.PrimEx instr])
    (action)

-- | Same as @guarding@ but doesn't pass an acquired result to a follow-up
-- action.
guarding_
  :: U.ExpandedInstr
  -> TypeCheckInstr a
  -> TypeCheckInstrNoExcept (TypeCheckedSeq inp)
  -> TypeCheckInstrNoExcept (TypeCheckedSeq inp)
guarding_ instr cond action = guarding instr cond (const action)

-- pva701: it's really painful to add arguments to TcInstrHandler
-- due to necessity to refactor @typeCheckInstr@.
-- Also functions which are being called from @typeCheckInstr@ would
-- have to be refactored too.
-- Therefore, I am using ReaderT over TypeCheck.
type TcInstrHandler
   = forall inp. (T.SingI inp)
      => U.ExpandedInstr
      -> HST inp
      -> TypeCheckInstrNoExcept (TypeCheckedSeq inp)