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

module Morley.Michelson.TypeCheck.TypeCheck
  ( TcInstr
  , TcInstrBase
  , TcOriginatedContracts
  , TypeCheckEnv(..)
  , BigMapFinder
  , TypeCheckOptions(..)
  , TypeCheck
  , TypeCheckNoExcept
  , TypeCheckResult
  , runTypeCheck
  , TypeCheckInstr
  , TypeCheckInstrNoExcept
  , runTypeCheckIsolated
  , typeCheckingWith
  , liftNoExcept
  , throwingTcError
  , IsInstrOp (..)
  , 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.Data (Data)
import Data.Default (Default(..))
import Data.Singletons (demote)
import Fmt (Buildable, build)

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 op = '[TypeCheckEnv op, TypeCheckOptions]

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

-- | A non-throwing alternative for 'TypeCheck'. Mainly meant to be used for
-- construction of a partially typed tree (see 'TypeCheckedSeq').
type TypeCheckNoExcept op = MultiReaderT (TCEnvs op) 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 op = ReaderT TypeCheckOptions (Except (TcError' op))

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

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

-- | Version of 'TypeCheckNoExcept' additionally carrying instruction-specific
-- 'TypeCheckInstrEnv'
type TypeCheckInstrNoExcept op = MultiReaderT (TCInstrEnvs op) 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 TcTypeError SomeParamType
mkSomeParamType (U.ParameterType (T.AsUType (notescp :: T.Notes t)) ann) = do
  T.Dict <- T.checkScope @(T.ParameterScope t)
    & first (UnsupportedTypeForScope (demote @t))
  paramNotes <- T.mkParamNotes notescp ann
    & first IllegalParamDecl
  return $ SomeParamType paramNotes

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 op
  = TypeCheckValue
    -- ^ We're typechecking a value.
      (U.Value' op, 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 op = TypeCheckEnv
  { tcMode :: ~(TypeCheckMode op)
  }

type BigMapFinder = Natural -> Maybe SomeVBigMap

data TypeCheckOptions = TypeCheckOptions
  { tcVerbose :: Bool
    -- ^ Whether to add stack type comments after every
    -- instruction a la @octez-client@.
  , tcStrict :: Bool
    -- ^ Whether should we behave like in test run or real run
    -- (real run is more strict).
    --
    -- @octez-client run@'s behaviour can slightly differ from the behaviour
    -- of @octez-client originate@ and @octez-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 op -> TypeCheck op a -> TypeCheckResult op 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 op (SomeTcInstr t) -> TypeCheckResult op (SomeTcInstr t)
runTypeCheckIsolated = runTypeCheck TypeCheckTest

typeCheckingWith :: TypeCheckOptions -> TypeCheckResult op a -> Either (TcError' op) 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
  :: forall op inp m m'.
     ( MonadMultiReaderT m Identity
     , m' ~ ChangeMultiReaderBase m (Except (TcError' op))
     , MonadError (TcError' op) m'
     )
  => m (TypeCheckedSeq op inp) -> m' (SomeTcInstr inp)
throwingTcError action = liftNoExcept @_ @(TcError' op) action >>= tcsEither (const throwError) (pure)

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

-- | Operation types that appear in @AbstractInstr@.
--
-- 'Data' superclass constraint is used to search for particular sub-instructions.
class Data op => IsInstrOp op where
  -- | Lift a single instruction to operation.
  --
  -- This is used only in error messages.
  liftInstr :: U.InstrAbstract op -> op

  -- | Get callstack at this instruction, if any given.
  --
  -- This is used only in error messages.
  pickErrorSrcPos :: op -> Maybe ErrorSrcPos

  -- | Do our best to re-interpret an instruction as a value.
  --
  -- For instance, @{ {}, {} }@ may turn out to be represented both as nested
  -- @SeqEx@s and as nested @ValueList@ & @ValueNil@s.
  -- This function, along with 'tryValToOp', let us convert between these
  -- two forms.
  --
  -- TODO [#785]: try to remove these two functions
  tryOpToVal :: op -> Maybe (U.Value' op)

  -- | Do our best to re-interpret a value as an instruction.
  --
  -- This does the opposite to 'tryOpToVal'.
  --
  -- If a value is already a lambda, this can return 'Nothing'.
  tryValToOp :: U.Value' op -> Maybe op

instance IsInstrOp U.ExpandedOp where
  liftInstr = U.PrimEx

  pickErrorSrcPos = \case
    U.WithSrcEx cs _ -> Just cs
    _ -> Nothing

  tryOpToVal = \case
    U.SeqEx []      -> Just U.ValueNil
    U.SeqEx xs      -> U.ValueSeq . fromList <$> traverse tryOpToVal xs
    U.PrimEx {}     -> Nothing
    U.WithSrcEx _ i -> tryOpToVal i

  tryValToOp = \case
    U.ValueNil      -> Just $ U.SeqEx []
    U.ValueSeq xs   -> U.SeqEx . toList <$> traverse tryValToOp xs
    _               -> Nothing

-- | 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 op (TypeCheckedSeq op inp)
  -- ^ Acquiring computation
  -> ([TypeCheckedOp op] -> TypeCheckedInstr op)
  -- ^ The parent instruction constructor
  -> (SomeTcInstr inp -> TypeCheckInstr op (SomeTcInstr inp'))
  -- ^ The throwing action
  -> TypeCheckInstrNoExcept op (TypeCheckedSeq op 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 op (TypeCheckedSeq op inp)
  -- ^ Acquiring computation
  -> ([TypeCheckedOp op] -> TypeCheckedInstr op)
  -- ^ The parent instruction constructor
  -> (SomeTcInstr inp -> TypeCheckInstrNoExcept op (TypeCheckedSeq op inp'))
  -- ^ The action
  -> TypeCheckInstrNoExcept op (TypeCheckedSeq op 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
  :: IsInstrOp op
  => U.InstrAbstract op -- ^ Untyped instruction
  -> TypeCheckInstr op a -- ^ Acquiring computation
  -> (a -> TypeCheckInstrNoExcept op (TypeCheckedSeq op inp)) -- ^ Follow-up action
  -> TypeCheckInstrNoExcept op (TypeCheckedSeq op inp)
guarding instr cond action = do
  cond & tcEither
    (\err -> pure $ IllTypedSeq err [NonTypedInstr $ liftInstr instr])
    (action)

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

-- | Some type checker function.
--
-- It accepts untyped instructions to typecheck as @instr@ value, and @op@ type
-- stands for the basic operation type like @ExpandedOp@.
type TcInstr op instr
   = forall inp. (T.SingI inp)
      => instr
      -> HST inp
      -> TypeCheckInstrNoExcept op (TypeCheckedSeq op inp)

-- | Type checker function for the operation like @ExpandedOp@.
-- Function of this type we carry in most of the other functions to let
-- them recurse.
type TcInstrBase op = TcInstr op op