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

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

  , EpLiftSequence (..)
  , EntrypointCallT (..)
  , epcPrimitive
  , unsafeEpcCallRoot
  , SomeEntrypointCallT (..)
  , unsafeSepcCallRoot
  , sepcPrimitive
  , sepcName
  , ForbidOr
  , MkEntrypointCallRes (..)
  , mkEntrypointCall
  , mkDefEntrypointCall

  , tyImplicitAccountParam

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

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

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

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

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

pattern EpAddress :: KindedAddress kind -> EpName -> EpAddress
pattern EpAddress addr name = EpAddress' (MkAddress addr) name

{-# COMPLETE EpAddress #-}

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 unsafe . mkMText $ 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 (displayException exception)
    ParseEpAddressBadRefAnn txt -> "Invalid reference annotation:" <+> build 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

-- | 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 (hashLengthBytes + 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' {..}

-- 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) = UnsafeParamNotes
  { pnNotes   :: Notes t
  , pnRootAnn :: RootAnn
  } deriving stock (Show, Eq, Generic)
    deriving anyclass (NFData)

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

-- | Parameter without annotations.
starParamNotes :: SingI t => ParamNotes t
starParamNotes = UnsafeParamNotes 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 $ punctuate ", " $ squotes . 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 ""
      ]

-- | 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 t'ParamNotes' performing all necessary checks.
mkParamNotes :: Notes t -> RootAnn -> Either ParamEpError (ParamNotes t)
mkParamNotes nt fa = verifyParamNotes nt fa $> UnsafeParamNotes 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
    :: (SingI subparam, SingI r)
    => EpLiftSequence arg subparam -> EpLiftSequence arg ('TOr subparam r)
  EplWrapRight
    :: (SingI l, SingI 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.
unsafeEpcCallRoot :: ParameterScope param => EntrypointCallT param param
unsafeEpcCallRoot = 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 = unsafeEpcCallRoot
  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.
unsafeSepcCallRoot :: ParameterScope param => SomeEntrypointCallT param
unsafeSepcCallRoot = SomeEpc unsafeEpcCallRoot

-- | 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 param of
  NTOr _ lFieldAnn rFieldAnn (lNotes :: Notes a) (rNotes :: Notes b) ->
    withDeMorganScope @ForbidOp @'TOr @a @b $
    withDeMorganScope @ForbidNestedBigMaps @'TOr @a @b $
      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.
--
-- Prefer using 'mkDefEntrypointCall' for the default entrypoint.
mkEntrypointCall
  :: (ParameterScope param)
  => EpName
  -> ParamNotes param
  -> Maybe (MkEntrypointCallRes param)
mkEntrypointCall epName
  | isDefEpName epName = Just . mkDefEntrypointCall
  | otherwise = mkNamedEntrypointCall epName

-- | Build 'EntrypointCallT' calling the default entrypoint. Unlike
-- 'mkEntrypointCall', always succeeds.
mkDefEntrypointCall :: ParameterScope param => ParamNotes param -> MkEntrypointCallRes param
mkDefEntrypointCall notes@(ParamNotes paramNotes _) = fromMaybe rootCall $
  mkNamedEntrypointCall DefEpName notes
  where
    rootCall = MkEntrypointCallRes paramNotes EntrypointCall
      { epcName = DefEpName
      , epcParamProxy = Proxy
      , epcLiftSequence = EplArgHere
      }

-- | Make entrypoint call explicitly for the named entrypoint, without
-- defaulting to default entrypoint.
--
-- Intended to be internal.
mkNamedEntrypointCall
  :: (ParameterScope param)
  => EpName
  -> ParamNotes param
  -> Maybe (MkEntrypointCallRes param)
mkNamedEntrypointCall 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
        }
  ]

-- | @parameter@ type of implicit account.
tyImplicitAccountParam :: ParamNotes 'TUnit
tyImplicitAccountParam = UnsafeParamNotes starNotes noAnn