-- SPDX-FileCopyrightText: 2020 Tocqueville Group
-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA
-- SPDX-License-Identifier: LicenseRef-MIT-obsidian-systems

{-# LANGUAGE DeriveLift #-}

-- | Module that defines Expression type, its related types
-- and its JSON instance.
module Morley.Micheline.Expression
  ( Exp
    ( ..
    , ExpPrim'
    )
  , expressionInt
  , expressionString
  , expressionBytes
  , expressionSeq
  , expressionPrim
  , expressionPrim'
  , RegularExp
  , Expression
  , MichelinePrimAp(..)
  , MichelinePrimitive(..)
  , michelsonPrimitive
  , ExpExtensionDescriptorKind
  , ExpExtensionDescriptor (..)

  , ExpExtrasConstrained
  , ExpAllExtrasConstrainted
  , ExpExtras (..)
  , mkUniformExpExtras
  , hoistExpExtras

  , Annotation (..)
  , annotToText
  , annotFromText
  , isAnnotationField
  , isAnnotationType
  , isAnnotationVariable
  , isNoAnn
  , mkAnns
  , toAnnSet
  , mkAnnsFromAny

  -- * Prisms
  , _ExpInt
  , _ExpString
  , _ExpBytes
  , _ExpSeq
  , _ExpPrim
  , _ExpressionInt
  , _ExpressionString
  , _ExpressionBytes
  , _ExpressionSeq
  , _ExpressionPrim
  , _AnnotationField
  , _AnnotationVariable
  , _AnnotationType

  -- * Lenses
  , mpaPrimL
  , mpaArgsL
  , mpaAnnotsL
  ) where

import Control.Lens (Iso', Plated, Prism', iso, prism')
import Control.Lens.TH (makeLensesWith, makePrisms)

import Data.Aeson
  (FromJSON, ToJSON, object, parseJSON, toEncoding, toJSON, withObject, withText, (.!=), (.:),
  (.:?), (.=))
import Data.Aeson.Encoding.Internal qualified as Aeson
import Data.Aeson.KeyMap qualified as KeyMap
import Data.Aeson.Types qualified as Aeson
import Data.Data (Data)
import Data.Sequence qualified as Seq
import Data.Text qualified as T (uncons)
import Fmt (Buildable(..), pretty, (+|), (|+))

import Language.Haskell.TH.Syntax (Lift)
import Morley.Micheline.Json (StringEncode(StringEncode, unStringEncode))
import Morley.Michelson.Untyped qualified as U
import Morley.Michelson.Untyped.Annotation
  (AnnotationSet(..), FieldAnn, FieldTag, KnownAnnTag(..), TypeAnn, TypeTag, VarAnn, VarTag,
  annPrefix, fullAnnSet, minimizeAnnSet, mkAnnotation)
import Morley.Tezos.Crypto (encodeBase58Check)
import Morley.Util.ByteString (HexJSONByteString(..))
import Morley.Util.Lens (postfixLFields)

newtype MichelinePrimitive = MichelinePrimitive Text
  deriving newtype (Eq, Ord, IsString, ToJSON, FromJSON)
  deriving stock (Show, Data, Lift)

michelsonPrimitive :: Seq Text
michelsonPrimitive = Seq.fromList [
  -- NOTE: The order of this list *matters*!
  --
  -- The position of each item in the list determines which binary code it gets packed to.
  -- E.g.
  --   * "parameter" is at index 0 on the list, so it gets packed to `0x0300`
  --   * "storage" is at index 1, so it gets packed to `0x0301`
  --
  -- You can ask `octez-client` which code corresponds to a given instruction/type/constructor.
  --
  -- > octez-client convert data 'storage' from michelson to binary
  -- > 0x0301
  --
  -- Whenever new instructions/types/constructors are added to the protocol,
  -- we can regenerate this list using this script:
  --
  -- > ./scripts/get-micheline-exprs.sh
  --
  -- or find the full primitives list in the <https://gitlab.com/tezos/tezos/blob/master/src/proto_alpha/lib_protocol/michelson_v1_primitives.ml | sources>,
  -- see "prim_encoding" variable.
  --
  "parameter", "storage", "code", "False", "Elt", "Left", "None", "Pair",
  "Right", "Some", "True", "Unit", "PACK", "UNPACK", "BLAKE2B", "SHA256",
  "SHA512", "ABS", "ADD", "AMOUNT", "AND", "BALANCE", "CAR", "CDR",
  "CHECK_SIGNATURE", "COMPARE", "CONCAT", "CONS", "CREATE_ACCOUNT", "CREATE_CONTRACT", "IMPLICIT_ACCOUNT", "DIP",
  "DROP", "DUP", "EDIV", "EMPTY_MAP", "EMPTY_SET", "EQ", "EXEC", "FAILWITH",
  "GE", "GET", "GT", "HASH_KEY", "IF", "IF_CONS", "IF_LEFT", "IF_NONE",
  "INT", "LAMBDA", "LE", "LEFT", "LOOP", "LSL", "LSR", "LT",
  "MAP", "MEM", "MUL", "NEG", "NEQ", "NIL", "NONE", "NOT",
  "NOW", "OR", "PAIR", "PUSH", "RIGHT", "SIZE", "SOME", "SOURCE",
  "SENDER", "SELF", "STEPS_TO_QUOTA", "SUB", "SWAP", "TRANSFER_TOKENS", "SET_DELEGATE", "UNIT",
  "UPDATE", "XOR", "ITER", "LOOP_LEFT", "ADDRESS", "CONTRACT", "ISNAT", "CAST",
  "RENAME", "bool", "contract", "int", "key", "key_hash", "lambda", "list",
  "map", "big_map", "nat", "option", "or", "pair", "set", "signature",
  "string", "bytes", "mutez", "timestamp", "unit", "operation", "address", "SLICE",
  "DIG", "DUG", "EMPTY_BIG_MAP", "APPLY", "chain_id", "CHAIN_ID", "LEVEL", "SELF_ADDRESS",
  "never", "NEVER", "UNPAIR", "VOTING_POWER", "TOTAL_VOTING_POWER", "KECCAK", "SHA3", "PAIRING_CHECK",
  "bls12_381_g1", "bls12_381_g2", "bls12_381_fr", "sapling_state", "sapling_transaction_deprecated", "SAPLING_EMPTY_STATE", "SAPLING_VERIFY_UPDATE", "ticket",
  "TICKET_DEPRECATED", "READ_TICKET", "SPLIT_TICKET", "JOIN_TICKETS", "GET_AND_UPDATE", "chest", "chest_key", "OPEN_CHEST",
  "VIEW", "view", "constant", "SUB_MUTEZ", "tx_rollup_l2_address", "MIN_BLOCK_TIME", "sapling_transaction", "EMIT", "Lambda_rec", "LAMBDA_REC", "TICKET"
  ]

-- | Type for Micheline Expression with extension points.
--
-- Following the Trees-that-Grow approach, this type provides the core set
-- of constructors used by Tezos accompanied with additional data (@XExp*@).
-- Plus additional constructors provided by @XExp@.
--
-- The type argument @x@ will be called /extension descriptor/ and it must have
-- @ExpExtensionDescriptor@ instance.
data Exp x
  = ExpInt (XExpInt x) Integer
  | ExpString (XExpString x) Text
  | ExpBytes (XExpBytes x) ByteString
  | ExpSeq (XExpSeq x) [Exp x]
  | ExpPrim (XExpPrim x) (MichelinePrimAp x)
  | ExpX (XExp x)

pattern ExpPrim' :: XExpPrim x -> MichelinePrimitive -> [Exp x] -> [Annotation] -> Exp x
pattern ExpPrim' x primAp exprs anns = ExpPrim x (MichelinePrimAp primAp exprs anns)

deriving stock instance ExpAllExtrasConstrainted Eq x => Eq (Exp x)
deriving stock instance ExpAllExtrasConstrainted Show x => Show (Exp x)
deriving stock instance (ExpAllExtrasConstrainted Data x, Typeable x) => Data (Exp x)
deriving stock instance ExpAllExtrasConstrainted Lift x => Lift (Exp x)

-- | Kind of extension descriptors.
--
-- We use a dedicated open type for this, not just @Type@, to notice earlier
-- when type arguments are mis-placed.
type ExpExtensionDescriptorKind = ExpExtensionTag -> Type
data ExpExtensionTag

-- | Defines details of extension descriptor.
class ExpExtensionDescriptor (x :: ExpExtensionDescriptorKind) where

  -- | Additional data in 'ExpInt' constructor.
  type XExpInt x :: Type
  type XExpInt _ = ()

  -- | Additional data in 'ExpString' constructor.
  type XExpString x :: Type
  type XExpString _ = ()

  -- | Additional data in 'ExpBytes' constructor.
  type XExpBytes x :: Type
  type XExpBytes _ = ()

  -- | Additional data in 'ExpSeq' constructor.
  type XExpSeq x :: Type
  type XExpSeq _ = ()

  -- | Additional data in 'ExpPrim' constructor.
  type XExpPrim x :: Type
  type XExpPrim _ = ()

  -- | Additional constructors.
  type XExp x :: Type
  type XExp _ = Void

-- | Constraint all the extra fields provided by this extension.
type ExpExtrasConstrained c x =
  Each '[c]
  [XExpInt x, XExpString x, XExpBytes x, XExpSeq x, XExpPrim x]

-- | Constraint all the extra fields and the constructor provided by
-- this extension.
type ExpAllExtrasConstrainted c x = (ExpExtrasConstrained c x, c (XExp x))

-- | A helper type that carries something for all extra fields.
--
-- Fields are carried in the given functor @f@ so that one could provide
-- a generator, parser or something else.
--
-- Extra constructor is not included here as it may need special treatment,
-- you have to carry it separately.
data ExpExtras f x = ExpExtras
  { eeInt :: f (XExpInt x)
  , eeString :: f (XExpString x)
  , eeBytes :: f (XExpBytes x)
  , eeSeq :: f (XExpSeq x)
  , eePrim :: f (XExpPrim x)
  }

-- | Fill 'ExpExtras' with the same data, assuming all types of extras are
-- the same.
mkUniformExpExtras
  :: ( extra ~ XExpInt x
     , extra ~ XExpString x
     , extra ~ XExpBytes x
     , extra ~ XExpSeq x
     , extra ~ XExpPrim x
     )
  => f extra -> ExpExtras f x
mkUniformExpExtras x = ExpExtras x x x x x

-- | Change the functor used in 'ExpExtras'.
hoistExpExtras
  :: (forall extra. f1 extra -> f2 extra)
  -> ExpExtras f1 x -> ExpExtras f2 x
hoistExpExtras f ExpExtras{..} = ExpExtras
  { eeInt = f eeInt
  , eeString = f eeString
  , eeBytes = f eeBytes
  , eeSeq = f eeSeq
  , eePrim = f eePrim
  }

-- | Extension descriptor for plain expressions without additional data.
data RegularExp :: ExpExtensionDescriptorKind
instance ExpExtensionDescriptor RegularExp

-- | Simple expression without any extras.
type Expression = Exp RegularExp

expressionInt :: Integer -> Expression
expressionInt a = ExpInt () a

expressionString :: Text -> Expression
expressionString a = ExpString () a

expressionBytes :: ByteString -> Expression
expressionBytes a = ExpBytes () a

expressionSeq :: [Expression] -> Expression
expressionSeq a = ExpSeq () a

expressionPrim :: MichelinePrimAp RegularExp -> Expression
expressionPrim a = ExpPrim () a

expressionPrim' :: Text -> [Expression] -> [Annotation] -> Expression
expressionPrim' primName args anns =
  ExpPrim () (MichelinePrimAp (MichelinePrimitive primName) args anns)

-- | Default instance that uses @uniplate@ as implementation.
--
-- If it tries to find expressions for polymorphic types too agressively
-- (requiring 'Data' where you don't what that), feel free to define an
-- overlapping manual instance.
instance ( Typeable x
         , ExpAllExtrasConstrainted Data x
         , ExpAllExtrasConstrainted Typeable x)
         => Plated (Exp x)

instance Buildable Expression where
  build = \case
    ExpInt () i -> build i
    ExpString () s -> build s
    ExpBytes () b ->
      build $ encodeBase58Check b
    ExpSeq () s -> "(" +| buildList build s |+ ")"
    ExpPrim () (MichelinePrimAp (MichelinePrimitive text) s annots) ->
      text <> " " |+ "(" +|
      buildList build s +| ") " +|
      buildList (build . annotToText) annots
    where
      buildList buildElem = mconcat . intersperse ", " . map buildElem

data Annotation
  = AnnotationType TypeAnn
  | AnnotationVariable VarAnn
  | AnnotationField FieldAnn
  deriving stock (Eq, Show, Data, Lift)

data MichelinePrimAp x = MichelinePrimAp
  { mpaPrim :: MichelinePrimitive
  , mpaArgs :: [Exp x]
  , mpaAnnots :: [Annotation]
  }

deriving stock instance Eq (Exp x) => Eq (MichelinePrimAp x)
deriving stock instance Show (Exp x) => Show (MichelinePrimAp x)
deriving stock instance (Data (Exp x), Typeable x) => Data (MichelinePrimAp x)
deriving stock instance Lift (Exp x) => Lift (MichelinePrimAp x)

instance FromJSON (Exp x) => FromJSON (MichelinePrimAp x) where
  parseJSON = withObject "Prim" $ \v -> MichelinePrimAp
    <$> v .: "prim"
    <*> v .:? "args" .!= []
    <*> v .:? "annots" .!= []

instance ToJSON (Exp x) => ToJSON (MichelinePrimAp x) where
  toJSON MichelinePrimAp {..} = object $ catMaybes
    [ Just ("prim" .= mpaPrim)
    , if null mpaArgs then Nothing else Just ("args" .= mpaArgs)
    , if null mpaAnnots then Nothing else Just ("annots" .= mpaAnnots)
    ]

annotFromText :: forall m. MonadFail m => Text -> m Annotation
annotFromText txt = do
  (n, t) <-
    maybe (fail $ "Annotation '" <> toString txt <> "' is missing an annotation prefix.") pure $
      T.uncons txt
  if | toText [n] == annPrefix @TypeTag  -> handleErr $ AnnotationType <$> mkAnnotation t
     | toText [n] == annPrefix @VarTag   -> handleErr $ AnnotationVariable <$> mkAnnotation t
     | toText [n] == annPrefix @FieldTag -> handleErr $ AnnotationField <$> mkAnnotation t
     | otherwise                         -> fail $ "Unknown annotation type: " <> toString txt

  where
    handleErr :: Either Text a -> m a
    handleErr = \case
      Left err -> fail $ "Failed to parse annotation '" <> toString txt <> "': " <> toString err
      Right a -> pure a

annotToText :: Annotation -> Text
annotToText = \case
  AnnotationType n -> pretty n
  AnnotationVariable n -> pretty n
  AnnotationField n -> pretty n

mkAnns :: [TypeAnn] -> [FieldAnn] -> [VarAnn] -> [Annotation]
mkAnns tas fas vas =
  let minAnnSet = minimizeAnnSet $ fullAnnSet tas fas vas
  in (AnnotationType <$> asTypes minAnnSet) <>
     (AnnotationField <$> asFields minAnnSet) <>
     (AnnotationVariable <$> asVars minAnnSet)

mkAnnsFromAny :: [U.AnyAnn] -> [Annotation]
mkAnnsFromAny xs = xs <&> \case
  U.AnyAnnType x -> AnnotationType x
  U.AnyAnnField x -> AnnotationField x
  U.AnyAnnVar x -> AnnotationVariable x

isAnnotationField :: Annotation -> Bool
isAnnotationField = \case
  AnnotationField _ -> True
  _                 -> False

isAnnotationVariable :: Annotation -> Bool
isAnnotationVariable = \case
  AnnotationVariable _ -> True
  _                    -> False

isAnnotationType :: Annotation -> Bool
isAnnotationType = \case
  AnnotationType _ -> True
  _                -> False

isNoAnn :: Annotation -> Bool
isNoAnn = \case
  AnnotationVariable (U.Annotation "") -> True
  AnnotationField (U.Annotation "")    -> True
  AnnotationType (U.Annotation "")     -> True
  _                                    -> False

toAnnSet :: [Annotation] -> AnnotationSet
toAnnSet = foldMap $ \case
  AnnotationType a     -> U.singleAnnSet a
  AnnotationField a    -> U.singleAnnSet a
  AnnotationVariable a -> U.singleAnnSet a

instance FromJSON Annotation where
  parseJSON = withText "Annotation" annotFromText

instance ToJSON Annotation where
  toJSON = toJSON . annotToText
  toEncoding = toEncoding . annotToText

instance FromJSON Expression where
  parseJSON v = ExpSeq () <$> parseJSON v
            <|> ExpPrim () <$> parseJSON v
            <|> ExpString () <$> withObject "ExpressionString" (.: "string") v
            <|> ExpInt () . unStringEncode <$> withObject "ExpressionInt" (.: "int") v
            <|> ExpBytes () . unHexJSONByteString <$> withObject "ExpressionBytes" (.: "bytes") v

instance ToJSON Expression where
  toJSON (ExpSeq () xs) = toJSON xs
  toJSON (ExpPrim () xs) = toJSON xs
  toJSON (ExpString () x) = Aeson.Object (KeyMap.singleton "string" $ toJSON x)
  toJSON (ExpInt () x) = Aeson.Object (KeyMap.singleton "int" $ toJSON $ StringEncode x)
  toJSON (ExpBytes () x) = Aeson.Object (KeyMap.singleton "bytes" $ toJSON $ HexJSONByteString x)

  toEncoding (ExpSeq () xs) = toEncoding xs
  toEncoding (ExpPrim () xs) = toEncoding xs
  toEncoding (ExpString () x) = Aeson.pairs (Aeson.pair "string" (toEncoding x))
  toEncoding (ExpInt () x) = Aeson.pairs (Aeson.pair "int" (toEncoding $ StringEncode x))
  toEncoding (ExpBytes () x) = Aeson.pairs (Aeson.pair "bytes" (toEncoding $ HexJSONByteString x))

--------------------------------------------------------------------------------
-- Optics
--------------------------------------------------------------------------------

_ExpInt :: Prism' (Exp d) (XExpInt d, Integer)
_ExpInt = prism' (uncurry ExpInt) \case
  ExpInt x a -> Just (x, a)
  _ -> Nothing

_ExpString :: Prism' (Exp d) (XExpString d, Text)
_ExpString = prism' (uncurry ExpString) \case
  ExpString x a -> Just (x, a)
  _ -> Nothing

_ExpBytes :: Prism' (Exp d) (XExpBytes d, ByteString)
_ExpBytes = prism' (uncurry ExpBytes) \case
  ExpBytes x a -> Just (x, a)
  _ -> Nothing

_ExpSeq :: Prism' (Exp d) (XExpSeq d, [Exp d])
_ExpSeq = prism' (uncurry ExpSeq) \case
  ExpSeq x a -> Just (x, a)
  _ -> Nothing

_ExpPrim :: Prism' (Exp d) (XExpPrim d, MichelinePrimAp d)
_ExpPrim = prism' (uncurry ExpPrim) \case
  ExpPrim x a -> Just (x, a)
  _ -> Nothing

neglecting :: Iso' ((), a) a
neglecting = iso snd pure

_ExpressionInt :: Prism' Expression Integer
_ExpressionInt = _ExpInt . neglecting

_ExpressionString :: Prism' Expression Text
_ExpressionString = _ExpString . neglecting

_ExpressionBytes :: Prism' Expression ByteString
_ExpressionBytes = _ExpBytes . neglecting

_ExpressionSeq :: Prism' Expression [Expression]
_ExpressionSeq = _ExpSeq . neglecting

_ExpressionPrim :: Prism' Expression (MichelinePrimAp RegularExp)
_ExpressionPrim = _ExpPrim . neglecting

makePrisms ''Annotation
makeLensesWith postfixLFields ''MichelinePrimAp