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

-- | Utilities for lightweight entrypoints support.
module Michelson.Typed.Entrypoints
  ( EpAddress (..)
  , ParseEpAddressError (..)
  , formatEpAddress
  , mformatEpAddress
  , parseEpAddress
  , unsafeParseEpAddress
  , parseEpAddressRaw
  , unsafeParseEpAddressRaw
  , ParamNotes (..)
  , pattern ParamNotes
  , starParamNotes
  , ArmCoord (..)
  , ArmCoords
  , ParamEpError (..)
  , mkParamNotes

  , EpLiftSequence (..)
  , EntrypointCallT (..)
  , epcPrimitive
  , epcCallRootUnsafe
  , SomeEntrypointCallT (..)
  , sepcCallRootUnsafe
  , sepcPrimitive
  , sepcName
  , ForbidOr
  , MkEntrypointCallRes (..)
  , mkEntrypointCall

  , tyImplicitAccountParam

    -- * Re-exports
  , EpName (..)
  , pattern DefEpName
  , epNameFromParamAnn
  , epNameToParamAnn
  , epNameFromRefAnn
  , epNameToRefAnn
  , EpNameFromRefAnnError (..)
  ) where

import Control.Monad.Except (throwError)
import qualified Data.ByteString as BS
import Data.Constraint (Dict(..))
import qualified Data.List.NonEmpty as NE
import qualified Data.Text as T
import Fmt (Buildable(..), hexF, pretty, (+|), (|+))

import Michelson.Text
import Michelson.Typed.Annotation
import Michelson.Typed.Scope
import Michelson.Typed.Sing
import Michelson.Typed.T
import Michelson.Untyped.Annotation
import Michelson.Untyped.Entrypoints
import Tezos.Address
import Tezos.Crypto (keyHashLengthBytes)
import Util.TH
import Util.Typeable
import Util.TypeLits

----------------------------------------------------------------------------
-- Primitives
----------------------------------------------------------------------------
--
-- EpAddress
----------------------------------------------------------------------------

-- | Address with optional entrypoint name attached to it.
-- TODO: come up with better name?
data EpAddress = EpAddress
  { eaAddress :: Address
    -- ^ Address itself
  , eaEntrypoint :: EpName
    -- ^ Entrypoint name (might be empty)
  } deriving stock (Show, Eq, Ord, Generic)

instance Buildable EpAddress where
  build = build . formatEpAddress

instance NFData EpAddress

formatEpAddress :: EpAddress -> Text
formatEpAddress (EpAddress addr ep)
  | isDefEpName ep = formatAddress addr
  | otherwise = formatAddress addr <> "%" <> pretty ep

mformatEpAddress :: EpAddress -> MText
mformatEpAddress ea =
  let t = formatEpAddress ea
     -- Should be safe because set of characters allowed in annotations
     -- (and thus in 'EpName') is subset of characters allowed in Michelson strings.
  in mkMTextUnsafe t

data ParseEpAddressError
  = ParseEpAddressBadAddress ParseAddressError
  | ParseEpAddressRawBadAddress ParseAddressRawError
  | ParseEpAddressBadEntryopint ByteString UnicodeException
  | ParseEpAddressBadRefAnn Text
  | ParseEpAddressRefAnnError EpNameFromRefAnnError
  | ParseEpAddressInvalidLength Int
  deriving stock (Show, Eq, Generic)

instance NFData ParseEpAddressError

instance Buildable ParseEpAddressError where
  build = \case
    ParseEpAddressBadAddress err -> build err
    ParseEpAddressRawBadAddress err -> build err
    ParseEpAddressBadEntryopint addr exception ->
      "Invalid entrypoint given for raw adddress " <> hexF addr <>
      " and failed with " <> build (show @Text exception)
    ParseEpAddressBadRefAnn txt -> pretty $ "Invalid reference annotation: " <> txt
    ParseEpAddressRefAnnError err -> build err
    ParseEpAddressInvalidLength len ->
      "Given raw entrypoint address has invalid length: " <> build len

-- | Parse an address which can be suffixed with entrypoint name
-- (e.g. "tz1faswCTDciRzE4oJ9jn2Vm2dvjeyA9fUzU%entrypoint").
parseEpAddress :: Text -> Either ParseEpAddressError EpAddress
parseEpAddress txt =
  let (addrTxt, mannotTxt) = T.breakOn "%" txt
  in case mannotTxt of
    "" -> do
      addr <- first ParseEpAddressBadAddress $ parseAddress addrTxt
      return $ EpAddress addr DefEpName
    annotTxt' -> do
      addr <- first ParseEpAddressBadAddress $ parseAddress addrTxt
      annot <- first ParseEpAddressBadRefAnn $ case T.stripPrefix "%" annotTxt' of
        Nothing -> error "impossible"
        Just a -> mkAnnotation a
      epName <- first ParseEpAddressRefAnnError $ epNameFromRefAnn annot
      return $ EpAddress addr epName

unsafeParseEpAddress :: HasCallStack => Text -> EpAddress
unsafeParseEpAddress = either (error . pretty) id . parseEpAddress

-- | Parses byte representation of entrypoint address.
--
-- For every address
--
-- @
-- KT1QbdJ7M7uAQZwLpvzerUyk7LYkJWDL7eDh%foo%bar
-- @
--
-- we get the following byte representation
--
-- @
-- 01afab866e7f1e74f9bba388d66b246276ce50bf4700666f6f25626172
-- \________________________________________/\/\____/\/\____/
--               address                     %   ep1  % ep2
-- @
--
parseEpAddressRaw :: ByteString -> Either ParseEpAddressError EpAddress
parseEpAddressRaw raw = do
  let (bytes, eps) = BS.splitAt (keyHashLengthBytes + 2) raw
  eaAddress <- first ParseEpAddressRawBadAddress $ parseAddressRaw bytes
  decodedEntrypoint <- first (ParseEpAddressBadEntryopint raw) $ decodeUtf8' eps
  decodedAnnotation <- first ParseEpAddressBadRefAnn $ mkAnnotation decodedEntrypoint
  eaEntrypoint <- first ParseEpAddressRefAnnError $ epNameFromRefAnn decodedAnnotation
  pure $ EpAddress {..}

unsafeParseEpAddressRaw :: ByteString -> EpAddress
unsafeParseEpAddressRaw = either (error . pretty) id . parseEpAddressRaw

-- ParamNotes
----------------------------------------------------------------------------

-- | Annotations for contract parameter declaration.
--
-- Following the Michelson specification, this type has the following invariants:
-- 1. No entrypoint name is duplicated.
-- 2. If @default@ entrypoint is explicitly assigned, no "arm" remains uncallable.
data ParamNotes (t :: T) = ParamNotesUnsafe
  { pnNotes   :: Notes t
  , pnRootAnn :: RootAnn
  } deriving stock (Show, Eq, Generic)
    deriving anyclass (NFData)

pattern ParamNotes :: Notes t -> RootAnn -> ParamNotes t
pattern ParamNotes t f <- ParamNotesUnsafe t f
{-# COMPLETE ParamNotes #-}

-- | Parameter without annotations.
starParamNotes :: SingI t => ParamNotes t
starParamNotes = ParamNotesUnsafe starNotes noAnn

-- | Coordinates of "arm" in Or tree, used solely in error messages.
type ArmCoords = [ArmCoord]
data ArmCoord = AcLeft | AcRight
  deriving stock (Show, Eq, Generic)

instance NFData ArmCoord

instance Buildable ArmCoord where
  build = \case
    AcLeft -> "left"
    AcRight -> "right"

-- | Errors specific to parameter type declaration (entrypoints).
data ParamEpError
  = ParamEpDuplicatedNames (NonEmpty EpName)
  | ParamEpUncallableArm ArmCoords
  deriving stock (Show, Eq, Generic)

instance NFData ParamEpError

instance Buildable ParamEpError where
  build = \case
    ParamEpDuplicatedNames names -> mconcat
      [ "Duplicated entrypoint names: "
      , mconcat . intersperse ", " $ map (surround "'" "'" . build) (toList names)
      ]
    ParamEpUncallableArm arm -> mconcat
      [ "Due to presence of 'default' entrypoint, one of contract \"arms\" \
        \cannot be called: \""
      , mconcat . intersperse " - " $ map build arm
      , "\""
      , if length arm > 1 then " (in top-to-bottom order)" else ""
      ]
    where
    surround pre post builder = pre <> builder <> post

-- | Check whether given notes are valid parameter notes.
verifyParamNotes :: Notes t -> RootAnn -> Either ParamEpError ()
verifyParamNotes notes ra = do
  let allEps = appEndo (gatherEntrypoints notes) []
      duplicatedEps
        = mapMaybe (safeHead . tail)
        . NE.group
        . sort
        $ maybe allEps (: allEps) (epNameFromParamAnn ra)

  whenJust (nonEmpty duplicatedEps) $ \dups ->
    throwError $ ParamEpDuplicatedNames dups

  -- In case contract have explicit root entrypoint, we assume that everything is
  -- callable.
  when (ra == noAnn) $ void
    $ first ParamEpUncallableArm
    $ ensureAllCallable notes
  where
    gatherEntrypoints :: Notes t -> Endo [EpName]
    gatherEntrypoints = \case
      NTOr _ fn1 fn2 l r -> mconcat
        [ Endo $ maybe id (:) (epNameFromParamAnn fn1)
        , Endo $ maybe id (:) (epNameFromParamAnn fn2)
        , gatherEntrypoints l
        , gatherEntrypoints r
        ]
      _ -> mempty

    -- Here we can assume that there is no more than one @default@ entrypoint,
    -- because duplications check occurs earlier.
    --
    -- In case when multiple entrypoints are uncallable, the reference
    -- implementation prefers displaying (in error message) arms which are
    -- closer to the root, but here we don't do this because that would be
    -- some more complex to implement and not sure how much does it worth that.
    ensureAllCallable :: Notes t -> Either ArmCoords Bool
    ensureAllCallable = \case
      NTOr _ fnL fnR l r -> do
        let epNameL = epNameFromParamAnn fnL
        let epNameR = epNameFromParamAnn fnR
        haveDefLL <- first (AcLeft  :) $ ensureAllCallable l
        haveDefRR <- first (AcRight :) $ ensureAllCallable r

        let haveDefL = or [haveDefLL, maybe False isDefEpName epNameL]
        let haveDefR = or [haveDefRR, maybe False isDefEpName epNameR]

        when haveDefL $ first (AcRight :) $ checkAllEpsNamed epNameR r
        when haveDefR $ first (AcLeft  :) $ checkAllEpsNamed epNameL l

        return $ or [haveDefL, haveDefR]

      _ -> return False

    checkAllEpsNamed :: Maybe EpName -> Notes t -> Either ArmCoords ()
    checkAllEpsNamed epNameRoot
      | isJust epNameRoot = \_ -> pass
      | otherwise = \case
          NTOr _ fnL fnR l r -> do
            let epNameL = epNameFromParamAnn fnL
                epNameR = epNameFromParamAnn fnR

            first (AcLeft  :) $ checkAllEpsNamed epNameL l
            first (AcRight :) $ checkAllEpsNamed epNameR r

          _ -> throwError []

-- | Construct 'ParamNotes' performing all necessary checks.
mkParamNotes :: Notes t -> RootAnn -> Either ParamEpError (ParamNotes t)
mkParamNotes nt fa = verifyParamNotes nt fa $> ParamNotesUnsafe nt fa

----------------------------------------------------------------------------
-- Entrypoint logic
----------------------------------------------------------------------------

-- | Describes how to construct full contract parameter from given entrypoint
-- argument.
--
-- This could be just wrapper over @Value arg -> Value param@, but we cannot
-- use @Value@ type in this module easily.
data EpLiftSequence (arg :: T) (param :: T) where
  EplArgHere :: EpLiftSequence arg arg
  EplWrapLeft
    :: (KnownT subparam, KnownT r)
    => EpLiftSequence arg subparam -> EpLiftSequence arg ('TOr subparam r)
  EplWrapRight
    :: (KnownT l, KnownT subparam)
    => EpLiftSequence arg subparam -> EpLiftSequence arg ('TOr l subparam)

deriving stock instance Eq (EpLiftSequence arg param)
deriving stock instance Show (EpLiftSequence arg param)

$(deriveGADTNFData ''EpLiftSequence)

instance Buildable (EpLiftSequence arg param) where
  build = \case
    EplArgHere -> "×"
    EplWrapLeft es -> "Left (" <> build es <> ")"
    EplWrapRight es -> "Right (" <> build es <> ")"

-- | Reference for calling a specific entrypoint of type @arg@.
data EntrypointCallT (param :: T) (arg :: T) =
  ParameterScope arg => EntrypointCall
  { epcName :: EpName
    -- ^ Name of entrypoint.
  , epcParamProxy :: Proxy param
    -- ^ Proxy of parameter, to make parameter type more easily fetchable.
  , epcLiftSequence :: EpLiftSequence arg param
    -- ^ How to call this entrypoint in the corresponding contract.
  }

deriving stock instance Eq (EntrypointCallT param arg)
deriving stock instance Show (EntrypointCallT param arg)
instance NFData (EntrypointCallT param arg) where
  rnf (EntrypointCall name Proxy s) = rnf (name, s)

instance Buildable (EntrypointCallT param arg) where
  build EntrypointCall{..} =
    "Call " +| epcName |+ ": " +| epcLiftSequence |+ ""

-- | Construct 'EntrypointCallT' which calls no entrypoint and assumes that
-- there is no explicit "default" one.
--
-- Validity of such operation is not ensured.
epcCallRootUnsafe :: ParameterScope param => EntrypointCallT param param
epcCallRootUnsafe = EntrypointCall
  { epcName = DefEpName
  , epcParamProxy = Proxy
  , epcLiftSequence = EplArgHere
  }

-- | Call parameter which has no entrypoints, always safe.
epcPrimitive
  :: forall p.
     (ParameterScope p, ForbidOr p)
  => EntrypointCallT p p
epcPrimitive = epcCallRootUnsafe
  where
  _requireNoOr = Dict @(ForbidOr p)

type family ForbidOr (t :: T) :: Constraint where
  ForbidOr ('TOr l r) =
    TypeError
    ('Text "Cannot apply to sum type parameter " ':<>: 'ShowType ('TOr l r))
  ForbidOr _ = ()

-- | 'EntrypointCallT' with hidden parameter type.
--
-- This requires argument to satisfy 'ParameterScope' constraint.
-- Strictly speaking, entrypoint argument may one day start having different
-- set of constraints comparing to ones applied to parameter, but this seems
-- unlikely.
data SomeEntrypointCallT (arg :: T) =
  forall param. (ParameterScope param) =>
  SomeEpc (EntrypointCallT param arg)

instance Eq (SomeEntrypointCallT arg) where
  SomeEpc epc1 == SomeEpc epc2 = isJust @() $ do
    Refl <- eqP (epcParamProxy epc1) (epcParamProxy epc2)
    guard (epc1 == epc2)

deriving stock instance Show (SomeEntrypointCallT arg)
instance NFData (SomeEntrypointCallT arg) where
  rnf (SomeEpc epc) = rnf epc

instance Buildable (SomeEntrypointCallT arg) where
  build (SomeEpc epc) = build epc

-- | Construct 'SomeEntrypointCallT' which calls no entrypoint and assumes that
-- there is no explicit "default" one.
--
-- Validity of such operation is not ensured.
sepcCallRootUnsafe :: ParameterScope param => SomeEntrypointCallT param
sepcCallRootUnsafe = SomeEpc epcCallRootUnsafe

-- | Call parameter which has no entrypoints, always safe.
sepcPrimitive
  :: forall t.
     (ParameterScope t, ForbidOr t)
  => SomeEntrypointCallT t
sepcPrimitive = SomeEpc epcPrimitive

sepcName :: SomeEntrypointCallT arg -> EpName
sepcName (SomeEpc epc) = epcName epc

-- | Build 'EpLiftSequence'.
--
-- Here we accept entrypoint name and type information for the parameter
-- of target contract.
--
-- Returns 'Nothing' if entrypoint is not found.
-- Does not treat default entrypoints specially.
withEpLiftSequence
  :: forall param r.
     (ParameterScope param)
  => EpName
  -> Notes param
  -> (forall arg. (ParameterScope arg) => (Notes arg, EpLiftSequence arg param) -> r)
  -> Maybe r
withEpLiftSequence epName@(epNameToParamAnn -> epAnn) param cont =
  case (sing @param, param) of
    (STOr lSing _, NTOr _ lFieldAnn rFieldAnn lNotes rNotes) ->
      case (checkOpPresence lSing, checkNestedBigMapsPresence lSing) of
        (OpAbsent, NestedBigMapsAbsent) -> asum
          [ guard (lFieldAnn == epAnn) $> cont (lNotes, EplWrapLeft EplArgHere)
          , guard (rFieldAnn == epAnn) $> cont (rNotes, EplWrapRight EplArgHere)
          , withEpLiftSequence epName lNotes (cont . fmap @((,) _) EplWrapLeft)
          , withEpLiftSequence epName rNotes (cont . fmap @((,) _) EplWrapRight)
          ]
    _ -> Nothing

-- Helper datatype for 'mkEntrypointCall'.
data MkEntrypointCallRes param where
  MkEntrypointCallRes
    :: ParameterScope arg
    => Notes arg
    -> EntrypointCallT param arg
    -> MkEntrypointCallRes param

-- | Build 'EntrypointCallT'.
--
-- Here we accept entrypoint name and type information for the parameter
-- of target contract.
--
-- Returns 'Nothing' if entrypoint is not found.
mkEntrypointCall
  :: (ParameterScope param)
  => EpName
  -> ParamNotes param
  -> Maybe (MkEntrypointCallRes param)
mkEntrypointCall epName (ParamNotes paramNotes root) =
  asum
  [ do
      epName' <- epNameFromParamAnn root
      guard (epName == epName')
      return $ MkEntrypointCallRes
        paramNotes
        EntrypointCall
          { epcName = epName
          , epcParamProxy = Proxy
          , epcLiftSequence = EplArgHere
          }
  , withEpLiftSequence epName paramNotes $ \(argInfo, liftSeq) ->
      MkEntrypointCallRes argInfo $ EntrypointCall
        { epcName = epName
        , epcParamProxy = Proxy
        , epcLiftSequence = liftSeq
        }
  , guard (isDefEpName epName) $>
      MkEntrypointCallRes paramNotes EntrypointCall
        { epcName = epName
        , epcParamProxy = Proxy
        , epcLiftSequence = EplArgHere
        }
  ]

-- | "Parameter" type of implicit account.
tyImplicitAccountParam :: ParamNotes 'TUnit
tyImplicitAccountParam = ParamNotesUnsafe starNotes noAnn