-- 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 Data.Aeson
  (FromJSON, ToJSON, parseJSON, toEncoding, toJSON, withObject, withText, (.!=), (.:), (.:?))
import Data.Aeson.Casing (aesonPrefix, snakeCase)
import qualified Data.Aeson.Encoding.Internal as Aeson
import Data.Aeson.TH (deriveToJSON)
import qualified Data.Aeson.Types as Aeson
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
import Tezos.Crypto (encodeBase58Check)
import Util.ByteString (HexJSONByteString(..))

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

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"
  ]

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

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 MichelinePrimAp = MichelinePrimAp
  { mpaPrim :: MichelinePrimitive
  , mpaArgs :: Seq Expression
  , mpaAnnots :: Seq Annotation
  } deriving stock (Eq, Show)

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

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

deriveToJSON (aesonPrefix snakeCase) ''MichelinePrimAp