-- SPDX-FileCopyrightText: 2020 Tocqueville Group
-- SPDX-FileCopyrightText: 2018 obsidian.systems
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ
-- SPDX-License-Identifier: LicenseRef-MIT-obsidian-systems

-- | Module that defines Expression type, its related types
-- and its JSON instance.
module Morley.Micheline.Expression
  ( Annotation(..)
  , Expression(..)
  , MichelinePrimAp(..)
  , MichelinePrimitive(..)
  , michelsonPrimitive
  , annotToText
  , annotFromText
  ) where

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

import Michelson.Untyped.Annotation
  (FieldAnn, FieldTag, KnownAnnTag(..), TypeAnn, TypeTag, VarAnn, VarTag, annPrefix, mkAnnotation)
import Morley.Micheline.Json (StringEncode(StringEncode, unStringEncode))
import Tezos.Crypto (encodeBase58Check)
import Util.ByteString (HexJSONByteString(..))

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

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 `tezos-client` which code corresponds to a given instruction/type/constructor.
  --
  -- > tezos-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
  --
  "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", "SAPLING_EMPTY_STATE", "SAPLING_VERIFY_UPDATE", "ticket",
  "TICKET", "READ_TICKET", "SPLIT_TICKET", "JOIN_TICKETS", "GET_AND_UPDATE"
  ]

-- | Type for Micheline Expression
data Expression
  = ExpressionInt Integer
    -- ^ Micheline represents both nats and ints using the same decimal format.
    -- The Haskell Integer type spans all possible values that the final
    -- (Michelson) type could end up being, and then some, so we use
    -- (StringEncode Integer) to represent all integral values here for easy
    -- JSON encoding compatibility.
  | ExpressionString Text
  | ExpressionBytes ByteString
  | ExpressionSeq (Seq Expression)
  | ExpressionPrim MichelinePrimAp
  deriving stock (Eq, Show, Data)

instance Plated Expression

instance Buildable Expression where
  build = \case
    ExpressionInt i -> build $ i
    ExpressionString s -> build s
    ExpressionBytes b ->
      build $ encodeBase58Check b
    ExpressionSeq s -> "(" +| buildSeq build s |+ ")"
    ExpressionPrim (MichelinePrimAp (MichelinePrimitive text) s annots) ->
      text <> " " |+ "(" +|
      buildSeq build s +| ") " +|
      buildSeq (build . annotToText) annots
    where
      buildSeq buildElem =
        mconcat . intersperse ", " . map
        buildElem . toList

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

data MichelinePrimAp = MichelinePrimAp
  { mpaPrim :: MichelinePrimitive
  , mpaArgs :: Seq Expression
  , mpaAnnots :: Seq Annotation
  } deriving stock (Eq, Show, Data)

instance FromJSON MichelinePrimAp where
  parseJSON = withObject "Prim" $ \v -> MichelinePrimAp
    <$> v .: "prim"
    <*> v .:? "args" .!= mempty
    <*> v .:? "annots" .!= mempty

instance ToJSON MichelinePrimAp where
  toJSON MichelinePrimAp {..} = object $ catMaybes
    [ Just ("prim" .= mpaPrim)
    , if mpaArgs == mempty then Nothing else Just ("args" .= mpaArgs)
    , if mpaAnnots == mempty 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

instance FromJSON Annotation where
  parseJSON = withText "Annotation" annotFromText

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

instance FromJSON Expression where
  parseJSON v = ExpressionSeq <$> parseJSON v
            <|> ExpressionPrim <$> parseJSON v
            <|> ExpressionString <$> withObject "ExpressionString" (.: "string") v
            <|> ExpressionInt . unStringEncode <$> withObject "ExpressionInt" (.: "int") v
            <|> ExpressionBytes . unHexJSONByteString <$> withObject "ExpressionBytes" (.: "bytes") v

instance ToJSON Expression where
  toJSON (ExpressionSeq xs) = toJSON xs
  toJSON (ExpressionPrim xs) = toJSON xs
  toJSON (ExpressionString x) = Aeson.Object (HashMap.singleton "string" $ toJSON x)
  toJSON (ExpressionInt x) = Aeson.Object (HashMap.singleton "int" $ toJSON $ StringEncode x)
  toJSON (ExpressionBytes x) = Aeson.Object (HashMap.singleton "bytes" $ toJSON $ HexJSONByteString x)

  toEncoding (ExpressionSeq xs) = toEncoding xs
  toEncoding (ExpressionPrim xs) = toEncoding xs
  toEncoding (ExpressionString x) = Aeson.pairs (Aeson.pair "string" (toEncoding x))
  toEncoding (ExpressionInt x) = Aeson.pairs (Aeson.pair "int" (toEncoding $ StringEncode x))
  toEncoding (ExpressionBytes x) = Aeson.pairs (Aeson.pair "bytes" (toEncoding $ HexJSONByteString x))