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

module Morley.Michelson.Untyped.Entrypoints
  ( EpName (..)
  , pattern DefEpName
  , isDefEpName
  , epNameFromParamAnn
  , epNameToParamAnn
  , epNameFromRefAnn
  , epNameFromSelfAnn
  , epNameToRefAnn
  , EpNameFromRefAnnError (..)
  , buildEpName
  , mkEntrypointsMap
  , HandleImplicitDefaultEp(..)
  ) where

import Data.Aeson.TH (deriveJSON)
import Data.Map qualified as Map
import Fmt (Buildable(..), pretty, (<+>))
import Prettyprinter (enclose)

import Morley.Michelson.Untyped.Annotation
import Morley.Michelson.Untyped.Type
import Morley.Util.Aeson
import Morley.Util.CLI

-- | Entrypoint name.
--
-- There are two properties we care about:
--
-- 1. Special treatment of the @default@ entrypoint name.
-- @default@ is prohibited in the @CONTRACT@ instruction and in
-- values of @address@ and @contract@ types.
-- However, it is not prohibited in the @SELF@ instruction.
-- Hence, the value inside @EpName@ __can__ be @"default"@, so that
-- we can distinguish @SELF@ and @SELF %default@. It is important
-- to distinguish them because their binary representation that is
-- inserted into blockchain is different. For example, typechecking
-- @SELF %default@ consumes more gas than @SELF@.
-- In this module, we provide several smart constructors with different
-- handling of @default@, please use the appropriate one for your use case.
-- 2. The set of permitted characters. Intuitively, an entrypoint name should
-- be valid only if it is a valid annotation (because entrypoints are defined
-- using field annotations). However, it is not enforced in Tezos.
-- It is not clear whether this behavior is intended. There is an upstream
-- [issue](https://gitlab.com/tezos/tezos/-/issues/851) which received @bug@
-- label, so probably it is considered a bug. Currently we treat it as a bug
-- and deviate from upstream implementation by probiting entrypoint names that
-- are not valid annotations. If Tezos developers fix it soon, we will be happy.
-- If they don't, we should (maybe temporarily) remove this limitation from our
-- code. There is an
-- [issue](https://gitlab.com/morley-framework/morley/-/issues/275) in our
-- repo as well.
newtype EpName = UnsafeEpName { unEpName :: Text }
  deriving stock (Show, Eq, Ord, Generic)

instance NFData EpName

deriveJSON morleyAesonOptions ''EpName

-- | This is a bidirectional pattern that can be used for two purposes:
--
-- 1. Construct an 'EpName' referring to the default entrypoint.
-- 2. Use it in pattern-matching or in equality comparison to check whether
-- 'EpName' refers to the default entrypoint. This is trickier because there
-- are two possible 'EpName' values referring to the default entrypoints.
-- 'DefEpName' will match only the most common one (no entrypoint).
-- However, there is a special case: @SELF@ instruction can have explicit
-- @%default@ reference. For this reason, it is recommended to use
-- 'isDefEpName' instead. Pattern-matching on 'DefEpName' is still permitted
-- for backwards compatibility and for the cases when you are sure that
-- 'EpName' does not come from the @SELF@ instruction.
pattern DefEpName :: EpName
pattern DefEpName = UnsafeEpName ""

-- | Check whether given 'EpName' refers to the default entrypoint.
-- Unlike 'DefEpName' pattern, this function correctly handles all cases,
-- including the @SELF@ instruction.
isDefEpName :: EpName -> Bool
isDefEpName epName = epName == DefEpName || epName == UnsafeEpName "default"

instance Buildable EpName where
  build = \case
    DefEpName -> "<default>"
    UnsafeEpName name -> build name

-- | Make up 'EpName' from annotation in parameter type declaration.
--
-- Returns 'Nothing' if no entrypoint is assigned here.
epNameFromParamAnn :: FieldAnn -> Maybe EpName
epNameFromParamAnn an@(Annotation a)
  | an == noAnn = Nothing
  | a == "default" = Just (UnsafeEpName "")
  | otherwise = Just $ UnsafeEpName a

-- | Turn entrypoint name into annotation for contract parameter declaration.
epNameToParamAnn :: EpName -> FieldAnn
epNameToParamAnn (UnsafeEpName name)
  | name == "" = [annQ|default|]
  | otherwise = unsafe . mkAnnotation $ name

data EpNameFromRefAnnError
  = InEpNameBadAnnotation FieldAnn
  deriving stock (Show, Eq, Generic)

instance NFData EpNameFromRefAnnError

instance Buildable EpNameFromRefAnnError where
  build = \case
    InEpNameBadAnnotation (Annotation an) ->
      "Invalid entrypoint reference" <+> enclose "`" "`" (build an)

-- | Make up 'EpName' from annotation which is reference to an entrypoint.
-- Note that it's more common for Michelson to prohibit explicit @default@
-- entrypoint reference.
--
-- Specifically, @%default@ annotation is probitited in values of @address@
-- and @contract@ types. It's also prohibited in the @CONTRACT@ instruction.
-- However, there is an exception: @SELF %default@ is a perfectly valid
-- instruction. Hence, when you construct an 'EpName' from an annotation
-- that's part of @SELF@, you should use 'epNameFromSelfAnn' instead.
epNameFromRefAnn :: FieldAnn -> Either EpNameFromRefAnnError EpName
epNameFromRefAnn an@(Annotation a)
  | a == "default" = Left $ InEpNameBadAnnotation an
  | otherwise = Right $ UnsafeEpName a

-- | Make up an 'EpName' from an annotation which is part of the
-- @SELF@ instruction.
epNameFromSelfAnn :: FieldAnn -> EpName
epNameFromSelfAnn (Annotation a) = UnsafeEpName a

-- | Turn entrypoint name into annotation used as reference to entrypoint.
epNameToRefAnn :: EpName -> FieldAnn
epNameToRefAnn (UnsafeEpName name) = unsafe . mkAnnotation $ name

-- | Make a valid entrypoint name from an arbitrary text. This
-- function prohibits explicit @default@ entrypoint name which is
-- permitted by Michelson inside the @SELF@ instruction. This
-- limitation shouldn't be restrictive because @SELF@ is equivalent to
-- @SELF %default@.
buildEpName :: Text -> Either String EpName
buildEpName txt = do
  annotation <-
    mkAnnotation txt
    & first (mappend "Failed to parse entrypoint: " . pretty)
  epNameFromRefAnn annotation
    & first pretty

instance HasCLReader EpName where
  getReader = eitherReader (buildEpName . toText)
  getMetavar = "ENTRYPOINT"

-- | A simple enum flag to choose how to handle implicit default entrypoint in
-- 'mkEntrypointsMap'.
data HandleImplicitDefaultEp
  = WithoutImplicitDefaultEp
    -- ^ Omit the default entrypoint unless it's specified explicitly.
  | WithImplicitDefaultEp
    -- ^ Include the default entrypoint even if it's not specified explicitly.
    -- This produces exactly the full set of entrypoints as per Michelson spec.
  deriving stock (Show, Eq, Ord, Enum)

-- | Given an untyped parameter type, extract a map that maps entrypoint names
-- to the their parameter types. If there are duplicate entrypoints in the
-- given Type then the duplicate entrypoints at a deeper nesting level will get
-- overwritten with the ones that are on top.
mkEntrypointsMap :: HandleImplicitDefaultEp -> ParameterType -> Map EpName Ty
mkEntrypointsMap hide (ParameterType ty rootAnn) = mkEntrypointsMapRec rootAnn ty & case hide of
  WithImplicitDefaultEp -> Map.insertWith (flip const) DefEpName ty
  WithoutImplicitDefaultEp -> id

-- | Version of 'mkEntrypointsMap' for plain untyped type.
--
-- Note that this function does not handle implicit default entrypoint, in all
-- likelihood you want to use 'mkEntrypointsMap'.
mkEntrypointsMapRec :: FieldAnn -> Ty -> Map EpName Ty
mkEntrypointsMapRec curRootAnn ty =
  accountRoot curRootAnn <> accountTree ty
  where
    accountRoot rootAnn = Map.fromList $ do
      Just rootEp <- pure $ epNameFromParamAnn rootAnn
      return (rootEp, ty)

    accountTree (Ty t _) = case t of
      -- We are only interested in `Or` branches to extract entrypoint
      -- annotations.
      TOr f1 f2 t1 t2 -> mkEntrypointsMapRec f1 t1 <> mkEntrypointsMapRec f2 t2
      _ -> mempty