-- 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, ann, annPrefix)
import qualified Michelson.Untyped.Annotation as MUA (Annotation)
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 [
  "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", "SHA3",
  "KECCAK", "LEVEL"
  ]

-- | 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 :: MonadFail m => Text -> m Annotation
annotFromText txt = case result of
    Just a -> pure a
    Nothing -> fail "Unknown annotation type"
  where
    result = (AnnotationType <$> stripPrefix @TypeTag txt)
         <|> (AnnotationVariable <$> stripPrefix @VarTag txt)
         <|> (AnnotationField <$> stripPrefix @FieldTag txt)

stripPrefix :: forall tag . KnownAnnTag tag => Text -> Maybe (MUA.Annotation tag)
stripPrefix txt = do
  (n, t) <- T.uncons txt
  guard (toText [n] == prefix)
  Just $ ann t
  where
    prefix = annPrefix @tag

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