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

module Morley.Michelson.TypeCheck.Value
    ( typeCheckValImpl
    , tcFailedOnValue
    ) where

import Control.Monad.Except (liftEither, throwError)
import Data.Constraint (Dict(..))
import Data.List.NonEmpty qualified as NE
import Data.Map qualified as M
import Data.Set qualified as S
import Data.Singletons (Sing, SingI(..), demote, fromSing, withSingI)
import Data.Typeable ((:~:)(..))
import Fmt (pretty)
import Prelude hiding (EQ, GT, LT)

import Morley.Michelson.Text
import Morley.Michelson.TypeCheck.Error (TcError'(..), TcTypeError(..))
import Morley.Michelson.TypeCheck.Helpers
import Morley.Michelson.TypeCheck.TypeCheck
import Morley.Michelson.TypeCheck.Types
import Morley.Michelson.Typed
  (EpAddress(..), Notes(..), SingT(..), SomeVBigMap(..), Value'(..), castM)
import Morley.Michelson.Typed qualified as T
import Morley.Michelson.Typed.Contract (giveNotInView)
import Morley.Michelson.Untyped qualified as U
import Morley.Tezos.Address
import Morley.Tezos.Core
import Morley.Tezos.Crypto
import Morley.Tezos.Crypto.BLS12381 qualified as BLS
import Morley.Tezos.Crypto.Timelock (chestFromBytes, chestKeyFromBytes)
import Morley.Util.Type (onFirst)

tcFailedOnValue :: U.Value' [] op -> T.T -> Text -> Maybe TcTypeError -> TypeCheckInstr op a
tcFailedOnValue v t msg err = do
  loc <- view tcieErrorPos
  throwError $ TcFailedOnValue v t msg loc err

-- | Function @typeCheckValImpl@ converts a single Michelson value
-- given in representation from @Morley.Michelson.Type@ module to representation
-- in strictly typed GADT.
--
-- @typeCheckValImpl@ is polymorphic in the expected type of value.
--
-- Type checking algorithm pattern-matches on parse value representation,
-- expected type @t@ and constructs @Val t@ value.
--
-- If there was no match on a given pair of value and expected type,
-- that is interpreted as input of wrong type and type check finishes with
-- error.
--
-- @typeCheckValImpl@ also has a 'Maybe TcOriginatedContracts' argument that
-- should contain the originated contracts when typechecking a parameter and
-- 'Nothing' otherwise.
typeCheckValImpl
  :: forall ty op.
     (SingI ty, IsInstrOp op)
  => Maybe TcOriginatedContracts
  -> TcInstrBase op
  -> U.Value' [] op
  -> TypeCheckInstr op (T.Value ty)
typeCheckValImpl mOriginatedContracts tcDo uv = doTypeCheckVal (uv, sing @ty)
  where
    doTypeCheckVal :: forall tz. (U.Value' [] op, Sing tz) -> TypeCheckInstr op (T.Value tz)
    doTypeCheckVal (uvalue, sng) = do
      typeCheckMode <- asks' (tcMode @op)
      case (uvalue, sng) of
        (U.ValueInt i, STInt) -> pure $ T.VInt i
        (v@(U.ValueInt i), t@STNat)
          | i >= 0 -> pure $ VNat (fromInteger i)
          | otherwise -> tcFailedOnValue v (fromSing t) "" (Just NegativeNat)
        (v@(U.ValueInt i), t@STMutez) -> case mkMutez i of
          Right m -> pure $ VMutez m
          Left err -> tcFailedOnValue v (fromSing t) err (Just InvalidTimestamp)

        -- If `tcBigMaps` is a `Just`, then we simulate the behavior of the RPC
        -- @/run_code@ endpoint.
        -- If a value contains a natural number where a big_map is expected,
        -- we assume that number represents a big_map ID and replace it with
        -- the corresponding big_map value (if it exists and has the expected key/value types).
        (U.ValueInt bigMapId, STBigMap {}) | TypeCheckValue _ (Just bigMapFinder) <- typeCheckMode -> do
          let bigMapMaybe = bigMapFinder =<< fromIntegralMaybe @Integer @Natural bigMapId
          case bigMapMaybe of
            Just (SomeVBigMap (bigMap@VBigMap{} :: T.Value actualT)) ->
              withSingI sng $ castM @actualT @tz bigMap $
                tcFailedOnValue uvalue (fromSing sng) "" . Just . UnexpectedBigMapType bigMapId
            Nothing ->
              tcFailedOnValue uvalue (fromSing sng) "" (Just (InvalidBigMapId bigMapId))
        (U.ValueString s, STString) -> pure $ VString s
        (v@(U.ValueString s), t@STAddress) -> case T.parseEpAddress (unMText s) of
          Right addr -> pure $ VAddress addr
          Left err -> tcFailedOnValue v (fromSing t) "" (Just $ InvalidAddress err)
        (v@(U.ValueBytes b), t@STAddress) -> case T.parseEpAddressRaw (U.unInternalByteString b) of
          Right addr -> pure $ VAddress addr
          Left err -> tcFailedOnValue v (fromSing t) "" (Just $ InvalidAddress err)
        (v@(U.ValueString s), t@STKeyHash) -> case parseHash (unMText s)  of
          Right kHash -> pure $ VKeyHash kHash
          Left err -> tcFailedOnValue v (fromSing t) "" (Just $ InvalidKeyHash err)
        (v@(U.ValueBytes b), t@STKeyHash) ->
          case parseKeyHashRaw (U.unInternalByteString b) of
            Right kHash -> pure $ VKeyHash kHash
            Left err -> tcFailedOnValue v (fromSing t) "" (Just $ InvalidKeyHash err)
        (U.ValueInt i, STBls12381Fr) ->
          pure $ VBls12381Fr (fromIntegralOverflowing @Integer @Bls12381Fr i)
        (v@(U.ValueBytes b), t@STBls12381Fr) ->
          case BLS.fromMichelsonBytes (U.unInternalByteString b) of
            Right val -> pure $ VBls12381Fr val
            Left err -> tcFailedOnValue v (fromSing t) "" (Just $ InvalidBls12381Object err)
        (v@(U.ValueBytes b), t@STBls12381G1) ->
          case BLS.fromMichelsonBytes (U.unInternalByteString b) of
            Right val -> pure $ VBls12381G1 val
            Left err -> tcFailedOnValue v (fromSing t) "" (Just $ InvalidBls12381Object err)
        (v@(U.ValueBytes b), t@STBls12381G2) ->
          case BLS.fromMichelsonBytes (U.unInternalByteString b) of
            Right val -> pure $ VBls12381G2 val
            Left err -> tcFailedOnValue v (fromSing t) "" (Just $ InvalidBls12381Object err)
        (v@(U.ValueString s), t@STTimestamp) -> case parseTimestamp $ unMText s of
          Just tstamp -> pure $ VTimestamp tstamp
          Nothing -> tcFailedOnValue v (fromSing t) "" (Just InvalidTimestamp)
        (U.ValueInt i, STTimestamp) -> pure $ VTimestamp (timestampFromSeconds i)
        (U.ValueBytes (U.InternalByteString s), STBytes) -> pure $ VBytes s
        (U.ValueTrue, STBool) -> pure $ VBool True
        (U.ValueFalse, STBool) -> pure $ VBool False

        (U.ValueString (parsePublicKey . unMText -> Right s), STKey) ->
          pure $ T.VKey s
        (U.ValueBytes (parsePublicKeyRaw . U.unInternalByteString -> Right s), STKey) ->
          pure $ VKey s

        (U.ValueString (parseSignature . unMText -> Right s), STSignature) ->
          pure $ VSignature s
        (U.ValueBytes (parseSignatureRaw . U.unInternalByteString  -> Right s), STSignature) ->
          pure $ VSignature s

        (U.ValueString (parseChainId . unMText -> Right ci), STChainId) ->
          pure $ VChainId ci
        (U.ValueBytes (mkChainId . U.unInternalByteString -> Right ci), STChainId) ->
          pure $ VChainId ci

        (cv@(U.ValueString (T.parseEpAddress . unMText -> Right addr))
          , STContract pc) -> typecheckContractValue cv addr pc
        (cv@(U.ValueBytes (T.parseEpAddressRaw . U.unInternalByteString -> Right addr))
          , STContract pc) -> typecheckContractValue cv addr pc

        (cv, s@(STTicket vt)) -> asks' tcStrict >>= \case
          -- Note [Tickets forging]
          -- Normally, ticket values cannot be constructed manually (i.e. forged)
          -- by design, and the only valid way to make a ticket is calling
          -- @TICKET@ instruction.
          --
          -- However @octez-client run@ adds an exception, it allows passing a
          -- manually constructed ticket value as parameter or initial storage.
          --
          -- Since protocol lima zero amount tickets are no longer allowed
          -- event in "lax" mode, so we forbid them here is well. (see !1270)
          True ->
            tcFailedOnValue cv (fromSing sng)
              "ticket values cannot be forged in real operations"
              Nothing
          False -> case cv of
            U.ValuePair addrV (U.ValuePair datV amV) -> withComparable vt cv s $ do
              addrValue <- doTypeCheckVal (addrV, STAddress)
              dat <- doTypeCheckVal (datV, vt)
              amountValue <- doTypeCheckVal (amV, STNat)
              case (addrValue, amountValue) of
                (_, VNat 0) ->
                  tcFailedOnValue cv (fromSing sng) "ticket amount of 0 is not allowed"
                  Nothing
                (VAddress (EpAddress' addr ep), VNat am) -> do
                  unless (U.isDefEpName ep) $
                    tcFailedOnValue cv (fromSing sng)
                      "it is pointless to provide an address with entrypoint as a \
                      \ticketer, we do not support that"
                      Nothing
                  -- ↑ Tezos allows passing addresses with entrypoint, but it is
                  -- not possible to sanely work with them after that anyway,
                  -- since @TICKET@ instruction (the only valid way of constructing
                  -- tickets in real run) uses @SELF_ADDRESS@ result as ticketer.
                  pure $ VTicket addr dat am
            _ ->
              tcFailedOnValue cv (fromSing sng)
                "ticket type expects a value of type `(pair address <data> nat)`"
                Nothing

        (U.ValueUnit, STUnit) -> pure $ VUnit
        (U.ValuePair ml mr, STPair lt rt) -> do
          l <- doTypeCheckVal (ml, lt)
          r <- doTypeCheckVal (mr, rt)
          pure $ VPair (l, r)
        (U.ValueLeft ml, STOr lt rt) -> do
          l <- doTypeCheckVal (ml, lt)
          withSingI lt $ withSingI rt $ pure $ VOr (Left l)
        (U.ValueRight mr, STOr lt rt) -> do
          r <- doTypeCheckVal (mr, rt)
          withSingI lt $ withSingI rt $ pure $ VOr (Right r)
        (U.ValueSome mv, STOption op) -> do
          v <- doTypeCheckVal (mv, op)
          withSingI op $ pure $ VOption (Just v)
        (U.ValueNone, STOption op) -> do
          withSingI op $ pure $ VOption Nothing

        (U.ValueNil, STList l) -> withSingI l $
          pure $ T.VList []

        -- If ValueSeq contains at least 2 elements, it can be type checked as a
        -- right combed pair.
        (U.ValueSeq vals@(_ :| (_ : _)), STPair _ _) ->
          doTypeCheckVal (seqToRightCombedPair vals, sng)

        (U.ValueSeq (toList -> mels), STList l) -> do
          els <- typeCheckValsImpl (mels, l)
          withSingI l $ pure $ VList els

        (U.ValueNil, STSet s) -> do
          instrPos <- view tcieErrorPos
          case T.comparabilityPresence s of
            Just Dict -> withSingI s $ pure (T.VSet S.empty)
            Nothing -> throwError $ TcFailedOnValue uvalue (fromSing s) "Non comparable types are not allowed in Sets"
              instrPos (Just $ UnsupportedTypeForScope (fromSing s) T.BtNotComparable)

        (sq@(U.ValueSeq (toList -> mels)), s@(STSet vt)) -> withComparable vt sq s $ do
          instrPos <- view tcieErrorPos

          els <- typeCheckValsImpl (mels, vt)
          elsS <- liftEither $ S.fromDistinctAscList <$> ensureDistinctAsc id els
                    `onFirst` \msg -> TcFailedOnValue sq (fromSing vt) msg instrPos Nothing
          withSingI vt $ pure $ VSet elsS

        (v@U.ValueNil, s@(STMap st vt)) ->
          withSingI st $ withSingI vt $ withComparable st v s $ pure $ T.VMap M.empty

        (sq@(U.ValueMap (toList -> mels)), s@(STMap kt vt)) ->
          withComparable kt sq s $ do
            keyOrderedElts <- typeCheckMapVal mels sq kt vt
            withSingI kt $ withSingI vt $
              pure $ VMap (M.fromDistinctAscList keyOrderedElts)

        (v@U.ValueNil, s@(STBigMap st1 st2)) ->
          withSingI st1 $ withSingI st2 $ withComparable st1 v s $ withBigMapAbsence st2 v s $
            pure $ VBigMap Nothing M.empty

        (sq@(U.ValueMap (toList -> mels)), s@(STBigMap kt vt)) ->
          withComparable kt sq s $ withBigMapAbsence vt sq s $ do
            keyOrderedElts <- typeCheckMapVal mels sq kt vt
            withSingI kt $ withSingI vt $
              pure $ VBigMap Nothing (M.fromDistinctAscList keyOrderedElts)

        -- `{ {} }` can be typechecked either as `VLam` or `VList`.
        (U.ValueLambda s, STList l) ->
          case traverse tryOpToVal s of
            Just xs -> doTypeCheckVal (maybe U.ValueNil U.ValueSeq $ nonEmpty xs, sng)
            Nothing -> tcFailedOnValue uvalue (fromSing l) "" Nothing

        (U.ValueSeq s, STLambda _ _) -> case traverse tryValToOp s of
          Just xs -> doTypeCheckVal (U.ValueLambda $ toList xs, sng)
          Nothing -> tcFailedOnValue uvalue (fromSing sng) "" Nothing

        (U.ValueLamRec mp, lam@(STLambda (var :: Sing it) (b :: Sing ot)) :: Sing lam) ->
          withSingI var $ withSingI b $ do
            _ :/ instr <-
              withWTP @lam uvalue $ throwingTcError $
                typeCheckImpl
                  -- lambdas can contain operations forbidden inside views, hence
                  -- we invent a "not in view" constraint here.
                  (giveNotInView $ local (set tcieNotInView $ Just Dict) ... tcDo)
                  (toList mp)
                  ((sing @it, Dict) ::& (lam, Dict) ::& SNil)
            case instr of
              code ::: (lo :: HST lo) -> withWTP @ot uvalue $ do
                case eqHST1 @ot lo of
                  Right Refl -> do
                    pure $ T.mkVLamRec (T.RfNormal code)
                  Left m ->
                    tcFailedOnValue uvalue (demote @ty)
                            "wrong output type of lambda's value:" (Just m)
              AnyOutInstr code ->
                pure $ T.mkVLamRec (T.RfAlwaysFails code)

        (v, STLambda (var :: Sing it) (b :: Sing ot)) -> withSingI var $ withSingI b $ do
          mp <- case v of
            U.ValueNil       -> pure []
            U.ValueLambda mp -> pure $ toList mp
            _ -> tcFailedOnValue v (demote @ty) "unexpected value" Nothing
          _ :/ instr <-
            withWTP @it uvalue $ throwingTcError $
              typeCheckImpl
                -- lambdas can contain operations forbidden inside views, hence
                -- we invent a "not in view" constraint here.
                (giveNotInView $ local (set tcieNotInView $ Just Dict) ... tcDo)
                mp
                ((sing @it, Dict) ::& SNil)
          case instr of
            lam ::: (lo :: HST lo) -> withWTP @ot uvalue $ do
              case eqHST1 @ot lo of
                Right Refl -> do
                  pure $ T.mkVLam (T.RfNormal lam)
                Left m ->
                  tcFailedOnValue v (demote @ty)
                          "wrong output type of lambda's value:" (Just m)
            AnyOutInstr lam ->
              pure $ T.mkVLam (T.RfAlwaysFails lam)

        (v@(U.ValueBytes (U.InternalByteString bs)), STChest) -> do
          isStrict <- asks' tcStrict
          when isStrict $ tcFailedOnValue v T.TChest "chest type temporarily deprecated" Nothing
          case chestFromBytes bs of
            Right res -> pure $ VChest res
            Left err -> tcFailedOnValue v T.TChest err Nothing

        (v@(U.ValueBytes (U.InternalByteString bs)), STChestKey) -> do
          isStrict <- asks' tcStrict
          when isStrict $
            tcFailedOnValue v T.TChestKey "chest_key type temporarily deprecated" Nothing
          case chestKeyFromBytes bs of
            Right res -> pure $ VChestKey res
            Left err -> tcFailedOnValue v T.TChestKey err Nothing

        (v, t@(STSaplingState _)) -> tcFailedOnValue v (fromSing t)
          "sapling_state is not supported" Nothing

        (v, t@(STSaplingTransaction _)) -> tcFailedOnValue v (fromSing t)
          "sapling_transaction is not supported" Nothing

        (v, t) -> tcFailedOnValue v (fromSing t) "unknown value" Nothing

    seqToRightCombedPair :: (NonEmpty $ U.Value' [] op) -> U.Value' [] op
    seqToRightCombedPair (x :| [y]) = U.ValuePair x y
    seqToRightCombedPair (x :| xs) = U.ValuePair x $ seqToRightCombedPair (NE.fromList xs)

    withWTP
      :: forall t a. SingI t
      => U.Value' [] op
      -> (T.WellTyped t => TypeCheckInstr op a)
      -> TypeCheckInstr op a
    withWTP uvalue fn = case T.getWTP @t of
      Right Dict -> fn
      Left err -> tcFailedOnValue uvalue (demote @ty) (pretty err) Nothing

    typeCheckValsImpl
      :: ([U.Value' [] op], Sing t)
      -> TypeCheckInstr op [T.Value t]
    typeCheckValsImpl (mvs, sng) =
      fmap reverse $ foldM (\res mv -> (: res) <$> (doTypeCheckVal (mv, sng))) [] mvs

    -- Typechecks given list of @Elt@s and ensures that its keys are in ascending order.
    --
    -- It return list of pairs (key, value) with keys in ascending order
    -- so it is safe to call @fromDistinctAscList@ on returned list
    typeCheckMapVal
      :: [U.Elt [] op]
      -> U.Value' [] op
      -> Sing kt
      -> Sing vt
      -> TypeCheckInstr op [(T.Value kt, T.Value vt)]
    typeCheckMapVal mels sq kt vt = withComparable kt sq kt $ do
      instrPos <- view tcieErrorPos
      ks <- typeCheckValsImpl (map (\(U.Elt k _) -> k) mels, kt)
      vals <- typeCheckValsImpl (map (\(U.Elt _ v) -> v) mels, vt)
      ksS <- liftEither $ ensureDistinctAsc id ks
            `onFirst` \msg -> TcFailedOnValue sq (fromSing kt) msg instrPos Nothing
      pure $ zip ksS vals

    typecheckContractValue
      :: forall cp tz. (tz ~ 'T.TContract cp)
      => U.Value' [] op -> EpAddress -> Sing cp -> TypeCheckInstr op (T.Value tz)
    typecheckContractValue cv (EpAddress addr epName) pc = do

      instrPos <- view tcieErrorPos
      let unsupportedType =
            TcFailedOnValue cv (demote @ty) "wrong contract parameter" instrPos . Just
          checkParameterScope :: TypeCheckInstr op (Dict (T.ParameterScope cp))
          checkParameterScope = liftEither
            $ first (unsupportedType . UnsupportedTypeForScope (fromSing pc))
            $ withSingI pc $ T.checkScope @(T.ParameterScope cp)
      case addr of
        ImplicitAddress _ -> case pc of
          STUnit -> pure $ VContract (MkAddress addr) T.sepcPrimitive
          STTicket _ -> do
            Dict <- checkParameterScope
            pure $ VContract (MkAddress addr) T.sepcPrimitive
          _ -> throwError $ unsupportedType $ UnexpectedType (one "unit" :| [one "ticket 'a"])
        ContractAddress ca -> case mOriginatedContracts of
          Nothing -> throwError $ TcFailedOnValue cv (fromSing $ STContract pc)
            "'contract' type can not be used in this context" instrPos $
              Just $ UnsupportedTypeForScope (fromSing $ STContract pc) T.BtHasContract
          Just originatedContracts -> case M.lookup ca originatedContracts of
            Just (SomeParamType paramNotes) ->
              case T.mkEntrypointCall epName paramNotes of
                Nothing ->
                  throwError $
                  TcFailedOnValue cv (demote @ty) "unknown entrypoint" instrPos . Just $
                  EntrypointNotFound epName
                Just (T.MkEntrypointCallRes (_ :: Notes t') epc) -> do
                  Refl <- liftEither $ first unsupportedType $ withSingI pc $ eqType @cp @t'
                  pure $ VContract (MkAddress addr) (T.SomeEpc epc)
            Nothing ->
              throwError $ TcFailedOnValue cv (demote @ty) "Contract literal unknown"
                instrPos (Just $ UnknownContract addr)
        SmartRollupAddress _ -> do
          -- TODO [#932]: Check that argument type matches
          Dict <- checkParameterScope
          pure $ VContract (MkAddress addr) $ T.SomeEpc T.unsafeEpcCallRoot

withComparable
  :: forall a (t :: T.T) ty op. Sing a
  -> U.Value' [] op
  -> Sing t
  -> (T.Comparable a => TypeCheckInstr op ty)
  -> TypeCheckInstr op ty
withComparable s uv t act = case T.comparabilityPresence s of
  Just Dict -> act
  Nothing -> do
    instrPos <- view tcieErrorPos
    liftEither . Left $
      TcFailedOnValue uv (fromSing t) "Require a comparable type here" instrPos Nothing

withBigMapAbsence
  :: forall a (t :: T.T) ty op. Sing a
  -> U.Value' [] op
  -> Sing t
  -> (T.ForbidBigMap a => TypeCheckInstr op ty)
  -> TypeCheckInstr op ty
withBigMapAbsence s uv t act = case T.tAbsence T.SPSBigMap s of
  Just Refl -> act
  Nothing -> do
    instrPos <- view tcieErrorPos
    liftEither . Left $
      TcFailedOnValue uv (fromSing t) "Require a type which doesn't contain `big_map` here"
        instrPos Nothing