-- 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
  ( Expression(.., PrimExpr)
  , MichelinePrimAp(..)
  , MichelinePrimitive(..)
  , michelsonPrimitive

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

  -- * Prisms
  , _ExpressionInt
  , _ExpressionString
  , _ExpressionBytes
  , _ExpressionSeq
  , _ExpressionPrim
  , _AnnotationField
  , _AnnotationVariable
  , _AnnotationType

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

import Control.Lens (Plated)
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.Types qualified as Aeson
import Data.Data (Data)
import Data.HashMap.Strict qualified as HashMap
import Data.Sequence qualified as Seq
import Data.Text qualified as T (uncons)
import Fmt (Buildable(..), pretty, (+|), (|+))

import Language.Haskell.TH.Lift (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, 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 `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
  --
  -- 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", "SAPLING_EMPTY_STATE", "SAPLING_VERIFY_UPDATE", "ticket",
  "TICKET", "READ_TICKET", "SPLIT_TICKET", "JOIN_TICKETS", "GET_AND_UPDATE", "chest", "chest_key", "OPEN_CHEST",
  "VIEW", "view", "constant", "SUB_MUTEZ"
  ]

-- | 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 [Expression]
  | ExpressionPrim MichelinePrimAp
  deriving stock (Eq, Show, Data, Lift)

pattern PrimExpr :: Text -> [Expression] -> [Annotation] -> Expression
pattern PrimExpr primName args anns =
  ExpressionPrim (MichelinePrimAp (MichelinePrimitive primName) args anns)

instance Plated Expression

instance Buildable Expression where
  build = \case
    ExpressionInt i -> build $ i
    ExpressionString s -> build s
    ExpressionBytes b ->
      build $ encodeBase58Check b
    ExpressionSeq s -> "(" +| buildList build s |+ ")"
    ExpressionPrim (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 = MichelinePrimAp
  { mpaPrim :: MichelinePrimitive
  , mpaArgs :: [Expression]
  , mpaAnnots :: [Annotation]
  } deriving stock (Eq, Show, Data, Lift)

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

instance ToJSON MichelinePrimAp 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)

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

makePrisms ''Expression
makePrisms ''Annotation
makeLensesWith postfixLFields ''MichelinePrimAp