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

module Michelson.TypeCheck.TypeCheck
  ( TcInstrHandler
  , TcOriginatedContracts
  , TcResult
  , TypeCheckEnv(..)
  , TypeCheckOptions(..)
  , TypeCheck
  , TypeCheckNoExcept
  , runTypeCheck
  , TypeCheckInstr
  , TypeCheckInstrNoExcept
  , runTypeCheckIsolated
  , runTypeCheckInstrIsolated
  , liftNoExcept
  , liftNoExcept'
  , throwingTCError
  , throwingTCError'
  , preserving
  , preserving'
  , guarding
  , guarding_
  , tcEither

  , tcExtFramesL
  , tcModeL
  , TypeCheckMode(..)
  , SomeParamType(..)
  , mkSomeParamType
  , mkSomeParamTypeUnsafe
  ) where


import Control.Monad.Except (throwError)
import Data.Default (Default(..))
import Data.Singletons (Sing)
import Fmt (Buildable, build, pretty)
import qualified Text.Show

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

type TypeCheck =
  (ReaderT TypeCheckOptions
    (ExceptT TCError
      (State TypeCheckEnv)))

-- | A non-throwing alternative for @TypeCheck@. Mainly meant to be used for
-- construction of a partially typed tree (see @TypeCheckedSeq@).
type TypeCheckNoExcept =
  (ReaderT TypeCheckOptions
    (State TypeCheckEnv))

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

-- | @Show@ instance of @SomeParamType@, mainly used in test.
instance Show SomeParamType where
  show = show . someParamToParameterType

-- | @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 s T.ParamNotesUnsafe{..}) =
  U.ParameterType (T.AsUTypeExt s pnNotes) pnRootAnn

-- | Construct @SomeParamType@ from @ParameterType@, mainly used in test.
mkSomeParamTypeUnsafe :: HasCallStack => U.ParameterType -> SomeParamType
mkSomeParamTypeUnsafe p =
  case mkSomeParamType p of
    Right sp ->  sp
    Left err -> error $ "Illegal type in parameter of env contract: " <> pretty err

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 T.sing 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 (U.Value, T.T)
  | TypeCheckContract SomeParamType
  | TypeCheckTest
  | TypeCheckPack

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

data TypeCheckOptions = TypeCheckOptions
  { tcVerbose :: Bool -- ^ Whether to add stack type comments after every
                      -- instruction a la tezos-client.
  }

instance Default TypeCheckOptions where
  def = TypeCheckOptions{ tcVerbose = False }

makeLensesWith postfixLFields ''TypeCheckEnv

runTypeCheck :: TypeCheckMode -> TypeCheckOptions -> TypeCheck a -> Either TCError a
runTypeCheck mode options = evaluatingState (TypeCheckEnv [] mode)
                          . runExceptT
                          . usingReaderT options

-- | 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 a -> Either TCError a
runTypeCheckIsolated = runTypeCheck TypeCheckTest def

type TcResult inp = Either TCError (SomeInstr inp)

type TypeCheckInstr =
  ReaderT InstrCallStack TypeCheck

type TypeCheckInstrNoExcept =
  ReaderT InstrCallStack TypeCheckNoExcept

-- | Similar to 'runTypeCheckIsolated', but for 'TypeCheckInstr.'
runTypeCheckInstrIsolated :: TypeCheckInstr a -> Either TCError a
runTypeCheckInstrIsolated = runTypeCheckIsolated . flip runReaderT def

liftNoExcept :: TypeCheckInstrNoExcept a -> TypeCheckInstr a
liftNoExcept action = (pure action)
                 <**> (usingReaderT <$> ask)
                 <**> (usingReaderT <$> lift ask)
                 <**> (evaluatingState <$> get)

liftNoExcept' :: TypeCheckNoExcept a -> TypeCheck a
liftNoExcept' action = (pure action)
                  <**> (usingReaderT <$> ask)
                  <**> (evaluatingState <$> get)

throwingTCError
  :: TypeCheckInstrNoExcept (TypeCheckedSeq inp) -> TypeCheckInstr (SomeInstr inp)
throwingTCError action = liftNoExcept action
  >>= tcsEither (const throwError) (pure)

throwingTCError'
  :: TypeCheckNoExcept (TypeCheckedSeq inp) -> TypeCheck (SomeInstr inp)
throwingTCError' action = liftNoExcept' 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 = do
  res <- (pure action)
    <**> (usingReaderT <$> ask)
    <**> (usingReaderT <$> lift ask)
    <**> (pure runExceptT)
    <**> (evaluatingState <$> get)
  either onErr onOk res

-- | 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. (Typeable inp, HasCallStack)
      => U.ExpandedInstr
      -> HST inp
      -> TypeCheckInstrNoExcept (TypeCheckedSeq inp)