-- | Utilities for lightweight entrypoints support.
module Michelson.Typed.EntryPoints
  ( EpAddress (..)
  , ParseEpAddressError (..)
  , formatEpAddress
  , mformatEpAddress
  , parseEpAddress
  , unsafeParseEpAddress

  , ParamNotes (..)
  , pattern ParamNotes
  , ArmCoord (..)
  , ArmCoords
  , ParamEpError (..)
  , mkParamNotes

  , EpLiftSequence (..)
  , EntryPointCallT (..)
  , SomeEntryPointCallT (..)
  , sepcName
  , MkEntryPointCallRes (..)
  , mkEntryPointCall

  , tyImplicitAccountParam

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

import Data.Default (Default(..))
import qualified Data.List.NonEmpty as NE
import Data.Singletons (Sing, SingI(..))
import qualified Data.Text as T
import Data.Typeable ((:~:)(..))
import Fmt (Buildable(..), pretty, (+|), (|+))
import Test.QuickCheck (Arbitrary(..))

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 (Address, ParseAddressError, formatAddress, parseAddress)
import Util.Typeable

----------------------------------------------------------------------------
-- 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)

instance Buildable EpAddress where
  build = build . formatEpAddress

formatEpAddress :: EpAddress -> Text
formatEpAddress (EpAddress addr ep)
  | ep == def = 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
  | ParseEpAddressBadRefAnn Text
  | ParseEpAddressRefAnnError EpNameFromRefAnnError
  deriving stock (Show, Eq)

instance Buildable ParseEpAddressError where
  build = \case
    ParseEpAddressBadAddress err -> build err
    ParseEpAddressBadRefAnn txt -> pretty $ "Invalid reference annotation: " <> txt
    ParseEpAddressRefAnnError err -> build err

-- | 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 def
    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

instance Arbitrary FieldAnn => Arbitrary EpAddress where
  arbitrary = EpAddress <$> arbitrary <*> arbitrary

-- 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.
newtype ParamNotes (t :: T) = ParamNotesUnsafe
  { unParamNotes :: Notes t
  } deriving stock (Show, Eq)

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

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

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)

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 -> Either ParamEpError ()
verifyParamNotes notes = do
  let allEps = appEndo (gatherEntrypoints notes) []
      duplicatedEps = mapMaybe (safeHead . tail) . NE.group $ sort allEps
  whenJust (nonEmpty duplicatedEps) $ \dups ->
    Left $ ParamEpDuplicatedNames dups

  void $ ensureAllCallable notes
    & first ParamEpUncallableArm
  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, epNameL == Just (def @EpName)]
        let haveDefR = or [haveDefRR, epNameR == Just (def @EpName)]

        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
          _ -> Left []

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

----------------------------------------------------------------------------
-- 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 :: EpLiftSequence arg subparam -> EpLiftSequence arg ('TOr subparam r)
  EplWrapRight :: EpLiftSequence arg subparam -> EpLiftSequence arg ('TOr l subparam)

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

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 Buildable (EntryPointCallT param arg) where
  build EntryPointCall{..} =
    "Call " +| epcName |+ ": " +| epcLiftSequence |+ ""

-- | Calls the default entrypoint.
-- TODO [TM-280]: This implementation always refers parameter root, fix it.
instance (param ~ arg, ParameterScope arg) => Default (EntryPointCallT param arg) where
  def = EntryPointCall
    { epcName = def
    , epcParamProxy = Proxy
    , epcLiftSequence = EplArgHere
    }

-- | '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 Buildable (SomeEntryPointCallT arg) where
  build (SomeEpc epc) = build epc

instance ParameterScope arg => Default (SomeEntryPointCallT arg) where
  def = SomeEpc def

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
  :: (ParameterScope param)
  => EpName
  -> (Sing param, Notes param)
  -> (forall arg. (ParameterScope arg) => (Notes arg, EpLiftSequence arg param) -> r)
  -> Maybe r
withEpLiftSequence epName@(epNameToParamAnn -> epAnn) param cont =
  case param of
    (STOr lSing rSing, 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 (lSing, lNotes) (cont . fmap @((,) _) EplWrapLeft)
          , withEpLiftSequence epName (rSing, 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
  -> (Sing param, ParamNotes param)
  -> Maybe (MkEntryPointCallRes param)
mkEntryPointCall epName (paramSing, ParamNotes paramNotes) =
  asum
  [ withEpLiftSequence epName (paramSing, paramNotes) $ \(argInfo, liftSeq) ->
      MkEntryPointCallRes argInfo $ EntryPointCall
        { epcName = epName
        , epcParamProxy = Proxy
        , epcLiftSequence = liftSeq
        }
  , guard (epName == def) $>
      MkEntryPointCallRes paramNotes EntryPointCall
        { epcName = def
        , epcParamProxy = Proxy
        , epcLiftSequence = EplArgHere
        }
  ]

-- | "Parameter" type of implicit account.
tyImplicitAccountParam :: (Sing 'TUnit, ParamNotes 'TUnit)
tyImplicitAccountParam = (sing, ParamNotesUnsafe starNotes)

-- TODO [TM-280]: We also need to be able to handle field annotation in root
-- of parameter's @or@ tree.
-- Currently we don't even support field annotation at such position.
--
-- Also it would be nice to automatically add @%root@ annotation in each parameter
-- declaration when compiling Lorentz to Michelson.