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

module Lorentz.Lambda
  ( WrappedLambda(..)
  , Lambda
  , mkLambda
  , mkLambdaRec
  ) where

import Data.Singletons (demote)

import Lorentz.Annotation
import Lorentz.Base
import Lorentz.Coercions
import Lorentz.Instr.Framed
import Lorentz.Value
import Lorentz.Zip
import Morley.AsRPC
import Morley.Michelson.Doc
import Morley.Michelson.Typed hiding (Contract, pattern S)
import Morley.Michelson.Typed.Contract (giveNotInView)
import Morley.Michelson.Untyped (noAnn)
import Morley.Util.Markdown
import Morley.Util.Type

-- | A helper type to construct Lorentz lambda values; Use this for lambda
-- values outside of Lorentz contracts or with @push@.
data WrappedLambda i o
  = WrappedLambda (i :-> o)
  | RecLambda (i ++ '[WrappedLambda i o] :-> o)
  deriving stock (Show, Eq, Generic)

instance (KnownList i, ZipInstr i, ZipInstr o) => IsoValue (WrappedLambda i o) where
  type ToT (WrappedLambda i o) = 'TLambda (ToT (ZippedStack i)) (ToT (ZippedStack o))
  toVal (WrappedLambda i) = mkVLam $ unLorentzInstr $ zippingStack i
  toVal (RecLambda i) = mkVLamRec $ unLorentzInstr $
    framed @'[WrappedLambda i o] (unzipInstr @i)
    ## i
    ## zipInstr
  fromVal (VLam (LambdaCode i)) = WrappedLambda $ unzippingStack $ LorentzInstr i
  fromVal (VLam (LambdaCodeRec i)) = RecLambda $
    framed @'[WrappedLambda i o] (zipInstr @i)
    ## LorentzInstr i
    ## unzipInstr

instance MapLorentzInstr (WrappedLambda inp out) where
  mapLorentzInstr
    :: (forall i o. (i :-> o) -> (i :-> o))
    -> WrappedLambda inp out
    -> WrappedLambda inp out
  mapLorentzInstr f = \case
    WrappedLambda i -> WrappedLambda $ f i
    RecLambda i -> RecLambda $ f i

instance (Each '[HasAnnotation] '[ZippedStack i, ZippedStack o])
  => HasAnnotation (WrappedLambda i o) where
  getAnnotation b = NTLambda noAnn
    (getAnnotation @(ZippedStack i) b)
    (getAnnotation @(ZippedStack o) b)

instance HasRPCRepr (WrappedLambda i o) where type AsRPC (WrappedLambda i o) = WrappedLambda i o

-- | A constructor providing the required constraint for 'WrappedLambda'. This is
-- the only way to construct a lambda that uses operations forbidden in views.
mkLambda :: (IsNotInView => i :-> o) -> WrappedLambda i o
mkLambda i = WrappedLambda $ giveNotInView i

-- | A constructor providing the required constraint for 'WrappedLambda'. This is
-- the only way to construct a lambda that uses operations forbidden in views.
mkLambdaRec :: (IsNotInView => i ++ '[WrappedLambda i o] :-> o) -> WrappedLambda i o
mkLambdaRec i = RecLambda $ giveNotInView i

-- | A type synonym representing Michelson lambdas.
type Lambda i o = WrappedLambda '[i] '[o]

instance (Each [Typeable, ReifyList TypeHasDoc] [i, o])
  => TypeHasDoc (WrappedLambda i o) where
  typeDocName _ = "WrappedLambda (extended lambda)"
  typeDocMdReference tp wp =
    let DocItemRef ctorDocItemId = docItemRef (DType tp)
        refToThis = mdLocalRef (mdTicked "WrappedLambda") ctorDocItemId
    in applyWithinParens wp $
      mconcat $ intersperse " " [refToThis, refToStack @i, refToStack @o]
    where
    refToStack :: forall s. ReifyList TypeHasDoc s => Markdown
    refToStack =
      let stack = reifyList @_ @TypeHasDoc @s (\p -> typeDocMdReference p (WithinParens False))
      in mconcat
          [ mdBold "["
          , case stack of
              [] -> " "
              st -> mconcat $ intersperse (mdBold "," <> " ") st
          , mdBold "]"
          ]

  typeDocMdDescription =
    "`WrappedLambda i o` stands for a sequence of instructions which accepts stack \
    \of type `i` and returns stack of type `o`.\n\n\
    \When both `i` and `o` are of length 1, this primitive corresponds to \
    \the Michelson lambda. In more complex cases code is surrounded with `pair`\
    \and `unpair` instructions until fits into mentioned restriction.\
    \"
  typeDocDependencies _ = mconcat
    [ reifyList @_ @TypeHasDoc @i dTypeDepP
    , reifyList @_ @TypeHasDoc @o dTypeDepP
    , [ dTypeDep @Integer
      , dTypeDep @Natural
      , dTypeDep @MText
      ]
    ]
  typeDocHaskellRep _ _ = Nothing
  typeDocMichelsonRep _ =
    ( Just "WrappedLambda [Integer, Natural, MText, ()] [ByteString]"
    , demote @(ToT (WrappedLambda [Integer, Natural, MText, ()] '[ByteString]))
    )

instance ( CanCastTo (ZippedStack inp1) (ZippedStack inp2)
         , CanCastTo (ZippedStack out1) (ZippedStack out2)
         , CanCastTo (ZippedStack (inp1 ++ '[WrappedLambda inp1 out1]))
                     (ZippedStack (inp2 ++ '[WrappedLambda inp2 out2]))
         )
  => WrappedLambda inp1 out1 `CanCastTo` WrappedLambda inp2 out2 where
  castDummy = castDummyG