-- 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 (..)
  , MichelinePrimitiveTag (..)
  , SingMichelinePrimitiveTag (..)
  , ClassifiedMichelinePrimitive (..)
  , withClassifiedPrim
  , 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)
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.Singletons (Sing, sing)
import Data.Text qualified as T
import Fmt (Buildable(..), listF, pretty, tupleF, unwordsF)
import Language.Haskell.TH (Dec(..), caseE, conE, conP, conT, gadtC, match, normalB)

import Language.Haskell.TH.Syntax (Lift)
import Morley.Micheline.Expression.Internal.MichelinePrimitive
import Morley.Micheline.Expression.Internal.TH
import Morley.Micheline.Json (StringEncode(StringEncode, unStringEncode))
import Morley.Michelson.Typed.ClassifiedInstr.Internal.TH (promote)
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)

-- | GADT that has the same shape as 'MichelinePrimitive', but each constructor
-- carries a 'MichelinePrimitiveTag' tag with classification.
do
  [DataD cxt' name tvb mk _ ders] <-
    [d|data ClassifiedMichelinePrimitive (tag :: MichelinePrimitiveTag) where|]
  cons <- withMichelinePrimitiveCons \nm classifiedName ->
    gadtC [classifiedName] [] [t|$(conT name) $(promote $ primClassification $ primFromName nm)|]
  pure [DataD cxt' name tvb mk cons ders]

-- | Classify a 'MichelinePrimitive'. Intended to be used with @LambdaCase@,
-- similar to @withClassifiedInstr@.
withClassifiedPrim
  :: MichelinePrimitive
  -> (forall tag. Sing tag -> ClassifiedMichelinePrimitive tag -> r)
  -> r
withClassifiedPrim prim f = $(do
  matches <- withMichelinePrimitiveCons \nm classifiedName ->
    match (conP nm []) (normalB [|f sing $(conE classifiedName)|]) []
  caseE [|prim|] $ pure <$> matches
  )

-- | 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' :: MichelinePrimitive -> [Expression] -> [Annotation] -> Expression
expressionPrim' prim args anns =
  ExpPrim () (MichelinePrimAp prim 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 -> listF s
    ExpPrim () (MichelinePrimAp prim s annots) ->
      unwordsF [build prim, listF s, tupleF $ annotToText <$> annots]

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

makePrisms ''Exp

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