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

-- | Type-checking of Morley extension.
module Michelson.TypeCheck.Ext
  ( typeCheckExt
  ) where

import Control.Lens ((%=))
import Control.Monad.Except (MonadError, liftEither, throwError)
import Data.Constraint (Dict(..))
import Data.Map.Lazy (insert, lookup)
import Data.Typeable ((:~:)(..))

import Michelson.ErrorPos
import Michelson.TypeCheck.Error
import Michelson.TypeCheck.Helpers
import Michelson.TypeCheck.TypeCheck
import Michelson.TypeCheck.TypeCheckedSeq (IllTypedInstr(..), TypeCheckedSeq(..))
import Michelson.TypeCheck.Types
import Michelson.Typed (Notes(..), converge, mkUType, notesT, withUType)
import qualified Michelson.Typed as T
import Michelson.Untyped
  (ExpandedOp, StackFn, TyVar(..), Type, Var, VarAnn, sfnInPattern, sfnOutPattern,
  sfnQuantifiedVars, varSet)
import qualified Michelson.Untyped as U
import Util.Peano (SingNat(SS, SZ))

-- | Perform some, possibly throwing, action presumably making use of a supplied
-- external instruction. In case of an error, return @IllTypedSeq@ wrapping the
-- thrown error and the instruction. If the action successfully returns
-- @SomeInstr@, wrap it in a @WellTypedSeq@.
workOnInstr
  :: U.ExpandedExtInstr
  -> TypeCheckInstr (SomeInstr s)
  -> TypeCheckInstrNoExcept (TypeCheckedSeq s)
workOnInstr ext = tcEither
  (\err -> pure $ IllTypedSeq err [NonTypedInstr $ U.PrimEx $ U.EXT ext])
  (pure . WellTypedSeq)

typeCheckExt
  :: forall s. Typeable s
  => TcInstrHandler
  -> U.ExpandedExtInstr
  -> HST s
  -> TypeCheckInstrNoExcept (TypeCheckedSeq s)
typeCheckExt tcInstr ext hst = do
  instrPos <- ask
  case ext of
    U.STACKTYPE s -> workOnInstr ext $
      liftExtError hst $ nopSomeInstr <$ checkStackType noBoundVars s hst
    U.FN t sf op  -> checkFn tcInstr t sf op hst
    U.UPRINT pc   -> workOnInstr ext $
      verifyPrint pc <&> \tpc -> toSomeInstr (T.PRINT tpc)
    U.UTEST_ASSERT U.TestAssert{..} -> do
      let cons = U.EXT . U.UTEST_ASSERT . U.TestAssert tassName tassComment
      preserving (typeCheckImpl tcInstr tassInstrs hst) cons $ \(_ :/ si) -> case si of
        AnyOutInstr _ -> throwError $ TCExtError (SomeHST hst) instrPos $ TestAssertError
                         "TEST_ASSERT has to return Bool, but it always fails"
        instr ::: (((_ :: (T.Notes b, Dict (T.WellTyped b), VarAnn)) ::& _)) -> do
          Refl <- liftEither $ first
            (const $ TCExtError (SomeHST hst) instrPos $
              TestAssertError "TEST_ASSERT has to return Bool, but returned something else")
            (eqType @b @'T.TBool)
          tcom <- verifyPrint tassComment
          pure . toSomeInstr $ T.TEST_ASSERT $ T.TestAssert tassName tcom instr
        _ ->
          throwError $ TCExtError (SomeHST hst) instrPos $
            TestAssertError "TEST_ASSERT has to return Bool, but the stack is empty"

    U.UCOMMENT t ->
      -- TODO if we are going to analyze/parse programs from files,
      -- there should be parsing of string and creation of FunctionStarted/FunctionFinished/etc
      pure $ WellTypedSeq $ toSomeInstr $ T.COMMENT_ITEM $ T.JustComment t
  where
    verifyPrint :: U.PrintComment -> TypeCheckInstr (T.PrintComment s)
    verifyPrint (U.PrintComment pc) = do
      let checkStRef (Left txt)             = pure $ Left txt
          checkStRef (Right (U.StackRef i)) = Right <$> createStackRef i hst
      T.PrintComment <$> traverse checkStRef pc

    toSomeInstr ext' = hst :/ T.Ext ext' ::: hst
    nopSomeInstr = hst :/ T.Nop ::: hst

liftExtError :: Typeable s => HST s -> Either ExtError a -> TypeCheckInstr a
liftExtError hst ei = do
  instrPos <- ask
  liftEither $ first (TCExtError (SomeHST hst) instrPos) ei

-- | Check that the optional "forall" variables are consistent if present
checkVars :: Text -> StackFn -> Either ExtError ()
checkVars t sf = case sfnQuantifiedVars sf of
  Just qs
    | varSet (sfnInPattern sf) /= qs -> Left $ VarError t sf
  _ -> pass

-- | Executes function body, pushing @ExtFrame@ onto the state and checks
-- the pattern in @FN@.
checkFn
  :: Typeable inp
  => TcInstrHandler
  -> Text
  -> StackFn
  -> [ExpandedOp]
  -> HST inp
  -> TypeCheckInstrNoExcept (TypeCheckedSeq inp)
checkFn tcInstr t sf body inp = do
  guarding (con body) (checkStart inp) $ \vars ->
    preserving (typeCheckImplStripped tcInstr body inp) con $ \someI@(_ :/ instrAndOut) ->
      case instrAndOut of
        _ ::: out -> checkEnd vars out $> someI
        AnyOutInstr{} -> pure someI
  where
    checkStart hst = do
      liftExtError hst $ checkVars t sf
      vars <- liftExtError hst $ checkStackType noBoundVars (sfnInPattern sf) hst
      tcExtFramesL %= (vars :)
      return vars

    checkEnd :: Typeable out => BoundVars -> HST out -> TypeCheckInstr ()
    checkEnd vars out = liftExtError out $
      void $ checkStackType vars (sfnOutPattern sf) out

    con :: [op] -> U.InstrAbstract op
    con = U.EXT . U.FN t sf

-- | Check that a @StackTypePattern@ matches the type of the current stack
checkStackType
  :: Typeable xs
  => BoundVars
  -> U.StackTypePattern
  -> HST xs
  -> Either ExtError BoundVars
checkStackType (BoundVars vars boundStkRest) s hst = go vars 0 s hst
  where
    go :: Typeable xs => Map Var Type -> Int -> U.StackTypePattern -> HST xs
       -> Either ExtError BoundVars
    go m _ U.StkRest sr = case boundStkRest of
      Nothing -> pure $ BoundVars m (Just $ SomeHST sr)
      Just si@(SomeHST sr') ->
        bimap (StkRestMismatch s (SomeHST sr) si)
              (const $ BoundVars m (Just si))
              (eqHST sr sr')
    go m _ U.StkEmpty SNil = pure $ BoundVars m Nothing
    go _ _ U.StkEmpty _    = Left $ LengthMismatch s
    go _ _ _ SNil        = Left $ LengthMismatch s
    go m n (U.StkCons tyVar ts) ((xann :: Notes xt, _, _) ::& xs) =
      let
        handleType :: U.Type -> Either ExtError BoundVars
        handleType t =
          withUType t $ \(tann :: Notes t) -> do
            Refl <- first
              (\_ -> TypeMismatch s n $ TypeEqError (notesT xann) (notesT tann))
              (eqType @xt @t)
            void $ first (TypeMismatch s n . AnnError) (converge tann xann)
            go m (n + 1) ts xs
      in case tyVar of
        TyCon t -> handleType t
        VarID v -> case lookup v m of
          Nothing -> let t = mkUType xann in go (insert v t m) (n + 1) ts xs
          Just t -> handleType t

-- | Create stack reference accessing element with a given index.
--
-- Fails when index is too large for the given stack.
createStackRef
  :: (MonadError TCError m, MonadReader InstrCallStack m, Typeable s)
  => Natural -> HST s -> m (T.StackRef s)
createStackRef idx hst =
  case doCreate (hst, idx) of
    Just sr -> pure sr
    Nothing -> do
      instrPos <- ask
      throwError $
        TCExtError (SomeHST hst) instrPos $
          InvalidStackReference (U.StackRef idx) (StackSize $ lengthHST hst)
  where
    doCreate :: forall s. (HST s, Natural) -> Maybe (T.StackRef s)
    doCreate = \case
      (SNil, _) -> Nothing
      ((_ ::& _), 0) -> Just (T.StackRef SZ)
      ((_ ::& st), i) -> do
        T.StackRef ns <- doCreate (st, i - 1)
        return $ T.StackRef (SS ns)