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

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

import Control.Monad.Except (MonadError, liftEither, throwError)
import Data.Constraint (Dict(..))
import Data.Map.Lazy (insert, lookup)
import Data.Singletons (SingI, fromSing)
import Data.Typeable ((:~:)(..))

import Morley.Michelson.TypeCheck.Error
import Morley.Michelson.TypeCheck.Helpers
import Morley.Michelson.TypeCheck.TypeCheck
import Morley.Michelson.TypeCheck.TypeCheckedSeq (IllTypedInstr(..), TypeCheckedSeq(..))
import Morley.Michelson.TypeCheck.Types
import Morley.Michelson.Typed (Notes(..), notesT, withUType)
import Morley.Michelson.Typed qualified as T
import Morley.Michelson.Untyped (Ty, TyVar(..), Var)
import Morley.Michelson.Untyped qualified as U
import Morley.Util.MismatchError
import Morley.Util.PeanoNatural (PeanoNatural(..))

-- | 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
-- @SomeTcInstr@, wrap it in a @WellTypedSeq@.
workOnInstr
  :: IsInstrOp op
  => U.ExtInstrAbstract [] op
  -> TypeCheckInstr op (SomeTcInstr s)
  -> TypeCheckInstrNoExcept op (TypeCheckedSeq op s)
workOnInstr ext = tcEither
  (\err -> pure $ IllTypedSeq err [NonTypedInstr $ liftInstr $ U.EXT ext])
  (pure . WellTypedSeq)

typeCheckExt
  :: forall s op.
     (SingI s, IsInstrOp op)
  => TcInstrBase op
  -> U.ExtInstrAbstract [] op
  -> HST s
  -> TypeCheckInstrNoExcept op (TypeCheckedSeq op s)
typeCheckExt tcOp ext hst = do
  instrPos <- view tcieErrorPos
  case ext of
    U.STACKTYPE s -> workOnInstr ext $
      liftExtError hst $ stackTypeSomeInstr s <$ checkStackType noBoundVars s 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 tcOp 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.SingT b, Dict (T.WellTyped b))) ::& _)) -> 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 ->
      pure $ WellTypedSeq $ toSomeInstr $ T.COMMENT_ITEM $ T.JustComment t
  where
    verifyPrint :: U.PrintComment -> TypeCheckInstr op (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
    stackTypeSomeInstr s = hst :/ T.Ext (T.STACKTYPE s) ::: hst

liftExtError :: SingI s => HST s -> Either ExtError a -> TypeCheckInstr op a
liftExtError hst ei = do
  instrPos <- view tcieErrorPos
  liftEither $ first (TcExtError (SomeHST hst) instrPos) ei

-- | Check that a @StackTypePattern@ matches the type of the current stack
checkStackType
  :: SingI xs
  => BoundVars
  -> U.StackTypePattern
  -> HST xs
  -> Either ExtError BoundVars
checkStackType (BoundVars vars boundStkRest) s hst = go vars 0 s hst
  where
    go :: SingI xs => Map Var Ty -> 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) ((xsing :: T.SingT xt, _) ::& xs) =
      let
        handleType :: U.Ty -> Either ExtError BoundVars
        handleType t =
          withUType t $ \(tann :: Notes t) -> do
            Refl <- first
              (\_ -> TypeMismatch s n $ TypeEqError MkMismatchError
                { meExpected = fromSing xsing, meActual = notesT tann })
              (eqType @xt @t)
            go m (n + 1) ts xs
      in case tyVar of
        TyCon t -> handleType t
        VarID v -> case lookup v m of
          Nothing -> let t = T.toUType $ fromSing xsing 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' op) m, MonadReader TypeCheckInstrEnv m, SingI s)
  => Natural -> HST s -> m (T.StackRef s)
createStackRef idx hst =
  case doCreate (hst, idx) of
    Just sr -> pure sr
    Nothing -> do
      instrPos <- view tcieErrorPos
      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 Zero)
      ((_ ::& st), i) -> do
        T.StackRef ns <- doCreate (st, i - 1)
        return $ T.StackRef (Succ ns)