-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

-- | Michelson types represented in untyped model.

{-# LANGUAGE DeriveLift #-}

module Morley.Michelson.Untyped.Type
  ( Ty (..)
  , T (..)
  , ParameterType (..)
  , toption
  , tpair
  , tor
  , tyint
  , tynat
  , tyunit
  , tybool
  , typair
  , tyor
  , tyImplicitAccountParam
  , isAtomicType
  , isKey
  , isSignature
  , isComparable
  , isMutez
  , isKeyHash
  , isBool
  , isString
  , isInteger
  , isTimestamp
  , isNat
  , isInt
  , isBytes
  , renderType
  , unwrapT
  ) where

import Data.Aeson.TH (deriveJSON)
import Data.Data (Data(..))
import Data.List.NonEmpty ((<|))
import Fmt (Buildable(..), Doc, isEmpty, (<+>))
import Language.Haskell.TH.Syntax (Lift)
import Prelude hiding ((<$>))
import Prettyprinter qualified as PP

import Morley.Michelson.Printer.Util
import Morley.Michelson.Untyped.Annotation
  (AnnotationSet, FieldAnn, RootAnn, TypeAnn, VarAnn, emptyAnnSet, fullAnnSet, noAnn,
  pattern Annotation, singleAnnSet)
import Morley.Util.Aeson
import Morley.Util.MismatchError

-- | Annotated Michelson type.
-- We don't name it 'Type' to avoid conflicts with 'Data.Kind.Type'.
data Ty
  = Ty ~T TypeAnn
  deriving stock (Eq, Show, Data, Generic, Lift)
  deriving anyclass NFData

-- | Extract a raw Michelson type from an annotated one
unwrapT :: Ty -> T
unwrapT (Ty t _) = t

instance RenderDoc Ty where
  renderDoc pn (Ty t ta) = renderType t pn (singleAnnSet ta)

instance Buildable (MismatchError Ty) where
  build = buildDocDiff

instance RenderDoc T where
  renderDoc pn t = renderType t pn emptyAnnSet

-- | Since Babylon parameter type can have special root annotation.
data ParameterType = ParameterType Ty RootAnn
  deriving stock (Eq, Show, Data, Generic, Lift)

instance NFData ParameterType

instance RenderDoc ParameterType where
  renderDoc pn (ParameterType (Ty t ta) ra) =
    renderType t pn (fullAnnSet [ta] [ra] [])

-- | Render a type representation
--
-- Ordering between different kinds of annotations is not significant,
-- but ordering among annotations of the same kind is. Annotations
-- of a same kind must be grouped together.
--
-- > (prim @v :t %x arg1 arg2 ...)
--
-- these are equivalent
--
-- > PAIR :t @my_pair %x %y
-- > PAIR %x %y :t @my_pair
renderType
  :: T
  -> RenderContext
  -> AnnotationSet
  -> Doc
renderType t pn annSet =
  let annDoc = renderDoc doesntNeedParens annSet
      recRenderer t' annSet' = renderType t' needsParens annSet'
      renderParametric :: Text -> [Doc] -> Doc
      renderParametric name args = PP.group $
        if isEmpty annDoc && length name < 3
        then addParens pn $ build name <+> PP.align (PP.vsep args)
        else PP.hang 2 $ addParens pn $ PP.vsep $ build name <+> annDoc : args

      collectBranches :: T -> TypeAnn -> FieldAnn -> NonEmpty Doc
      collectBranches (TPair fa1 fa2 va1 _ (Ty t1 ta1) (Ty t2 ta2)) _ (Annotation "")
        = (recRenderer t1 $ (fullAnnSet [ta1] [fa1] [va1])) <| collectBranches t2 ta2 fa2
      collectBranches t' ta fa = (recRenderer t' $ fullAnnSet [ta] [fa] []) :| []

  in
  case t of
    TInt              -> wrapInParens pn $ "int" :| [annDoc]
    TNat              -> wrapInParens pn $ "nat" :| [annDoc]
    TString           -> wrapInParens pn $ "string" :| [annDoc]
    TMutez            -> wrapInParens pn $ "mutez" :| [annDoc]
    TBool             -> wrapInParens pn $ "bool" :| [annDoc]
    TKeyHash          -> wrapInParens pn $ "key_hash" :| [annDoc]
    TTimestamp        -> wrapInParens pn $ "timestamp" :| [annDoc]
    TBytes            -> wrapInParens pn $ "bytes" :| [annDoc]
    TAddress          -> wrapInParens pn $ "address" :| [annDoc]
    TKey              -> wrapInParens pn $ "key" :| [annDoc]
    TBls12381Fr       -> wrapInParens pn $ "bls12_381_fr" :| [annDoc]
    TBls12381G1       -> wrapInParens pn $ "bls12_381_g1" :| [annDoc]
    TBls12381G2       -> wrapInParens pn $ "bls12_381_g2" :| [annDoc]
    TUnit             -> wrapInParens pn $ "unit" :| [annDoc]
    TSignature        -> wrapInParens pn $ "signature" :| [annDoc]
    TChainId          -> wrapInParens pn $ "chain_id" :| [annDoc]
    TOperation        -> wrapInParens pn $ "operation" :| [annDoc]
    TChest            -> wrapInParens pn $ "chest" :| [annDoc]
    TChestKey         -> wrapInParens pn $ "chest_key" :| [annDoc]
    TNever            -> wrapInParens pn $ "never" :| [annDoc]
    TSaplingState n   -> addParens pn $ "sapling_state" <+> annDoc <+> build n
    TSaplingTransaction n -> addParens pn $ "sapling_transaction" <+> annDoc <+> build n

    TOption (Ty t1 ta1) ->
      renderParametric "option" [recRenderer t1 (singleAnnSet ta1)]

    TList (Ty t1 ta1) ->
      renderParametric "list" [recRenderer t1 (singleAnnSet ta1)]

    TSet (Ty t1 ta1) ->
      renderParametric "set" [recRenderer t1 (singleAnnSet ta1)]

    TContract (Ty t1 ta1) ->
      renderParametric "contract" [recRenderer t1 (singleAnnSet ta1)]

    TTicket (Ty t1 ta1) ->
      renderParametric "ticket" [recRenderer t1 (singleAnnSet ta1)]

    -- Optimize in comb pair rendering: `pair x y z` instead of `pair x (pair y z)`
    -- Works only if there is no field annotation on nested pair
    p@(TPair _ (Annotation "") _ _ (Ty _ _) (Ty (TPair {}) _)) ->
      renderParametric "pair" $ toList $ collectBranches p noAnn noAnn

    TPair fa1 fa2 va1 va2 (Ty t1 ta1) (Ty t2 ta2) ->
      renderParametric "pair"
        [ recRenderer t1 $ fullAnnSet [ta1] [fa1] [va1]
        , recRenderer t2 $ fullAnnSet [ta2] [fa2] [va2]
        ]

    TOr fa1 fa2 (Ty t1 ta1) (Ty t2 ta2) ->
      renderParametric "or"
        [ recRenderer t1 $ fullAnnSet [ta1] [fa1] []
        , recRenderer t2 $ fullAnnSet [ta2] [fa2] []
        ]

    TLambda (Ty t1 ta1) (Ty t2 ta2) ->
      renderParametric "lambda"
        [ recRenderer t1 $ singleAnnSet ta1
        , recRenderer t2 $ singleAnnSet ta2
        ]

    TMap (Ty t1 ta1) (Ty t2 ta2) ->
      renderParametric "map"
        [ recRenderer t1 $ singleAnnSet ta1
        , recRenderer t2 $ singleAnnSet ta2
        ]

    TBigMap (Ty t1 ta1) (Ty t2 ta2) ->
      renderParametric "big_map"
        [ recRenderer t1 $ singleAnnSet ta1
        , recRenderer t2 $ singleAnnSet ta2
        ]

-- | Michelson Type
data T =
    TKey
  | TUnit
  | TSignature
  | TChainId
  | TOption Ty
  | TList Ty
  | TSet Ty
  | TOperation
  | TContract Ty
  | TTicket Ty
  | TPair FieldAnn FieldAnn VarAnn VarAnn Ty Ty
  | TOr FieldAnn FieldAnn Ty Ty
  | TLambda Ty Ty
  | TMap Ty Ty
  | TBigMap Ty Ty
  | TInt
  | TNat
  | TString
  | TBytes
  | TMutez
  | TBool
  | TKeyHash
  | TBls12381Fr
  | TBls12381G1
  | TBls12381G2
  | TTimestamp
  | TAddress
  | TChest
  | TChestKey
  | TSaplingState Natural
  | TSaplingTransaction Natural
  | TNever
  deriving stock (Eq, Show, Data, Generic, Lift)
  deriving anyclass NFData

-- | Construct non-annotated @option@ type from an annotated type
toption :: Ty -> T
toption t = TOption t

-- | Construct non-annotated @pair@ type from two annotated types
tpair :: Ty -> Ty -> T
tpair l r = TPair noAnn noAnn noAnn noAnn l r

-- | Construct non-annotated @or@ type from two annotated types
tor :: Ty -> Ty -> T
tor l r = TOr noAnn noAnn l r

-- | Construct annotated @int@ type with an empty annotation
tyint :: Ty
tyint = Ty TInt noAnn

-- | Construct annotated @nat@ type with an empty annotation
tynat :: Ty
tynat = Ty TNat noAnn

-- | Construct annotated @unit@ type with an empty annotation
tyunit :: Ty
tyunit = Ty TUnit noAnn

-- | Construct annotated @bool@ type with an empty annotation
tybool :: Ty
tybool = Ty TBool noAnn

-- | Construct annotated @pair@ type with an empty annotation
typair :: Ty -> Ty -> Ty
typair l r = Ty (tpair l r) noAnn

-- | Construct annotated @or@ type with an empty annotation
tyor :: Ty -> Ty -> Ty
tyor l r = Ty (tor l r) noAnn

-- | For implicit account, which Ty its parameter seems to have
-- from outside.
tyImplicitAccountParam :: Ty
tyImplicitAccountParam = Ty TUnit noAnn

-- | Check if type is atomic.
isAtomicType :: Ty -> Bool
isAtomicType t@(Ty _ tAnn) | tAnn == noAnn =
    isComparable t || isKey t || isUnit t || isSignature t || isOperation t
isAtomicType _ = False

-- | Predicate checking if type is @key@
isKey :: Ty -> Bool
isKey (Ty TKey _) = True
isKey _ = False

-- | Predicate checking if type is @unit@
isUnit :: Ty -> Bool
isUnit (Ty TUnit _) = True
isUnit _ = False

-- | Predicate checking if type is @signature@
isSignature :: Ty -> Bool
isSignature (Ty TSignature _) = True
isSignature _ = False

-- | Predicate checking if type is @operation@
isOperation :: Ty -> Bool
isOperation (Ty TOperation _) = True
isOperation _ = False

-- | Predicate checking if type is comparable, i.e. true for @int@, @nat@, @string@, etc.
-- see <https://tezos.gitlab.io/michelson-reference/> for a complete list of comparable types.
isComparable :: Ty -> Bool
isComparable (Ty t _) = case t of
  TInt -> True
  TNat -> True
  TString -> True
  TBytes -> True
  TMutez -> True
  TBool -> True
  TKeyHash -> True
  TTimestamp -> True
  TAddress -> True
  _ -> False

-- | Predicate checking if type is @mutez@
isMutez :: Ty -> Bool
isMutez (Ty TMutez _) = True
isMutez _ = False

-- | Predicate checking if type is @timestamp@
isTimestamp :: Ty -> Bool
isTimestamp (Ty TTimestamp _) = True
isTimestamp _ = False

-- | Predicate checking if type is @keyhash@
isKeyHash :: Ty -> Bool
isKeyHash (Ty TKeyHash _) = True
isKeyHash _ = False

-- | Predicate checking if type is @bool@
isBool :: Ty -> Bool
isBool (Ty TBool _) = True
isBool _ = False

-- | Predicate checking if type is @string@
isString :: Ty -> Bool
isString (Ty TString _) = True
isString _ = False

-- | Predicate checking if type is integral, i.e. @nat@, @int@, @mutez@, or @timestamp@
isInteger :: Ty -> Bool
isInteger a = isNat a || isInt a || isMutez a || isTimestamp a

-- | Predicate checking if type is @nat@
isNat :: Ty -> Bool
isNat (Ty TNat _) = True
isNat _ = False

-- | Predicate checking if type is @int@
isInt :: Ty -> Bool
isInt (Ty TInt _) = True
isInt _ = False

-- | Predicate checking if type is @bytes@
isBytes :: Ty -> Bool
isBytes (Ty TBytes _) = True
isBytes _ = False

----------------------------------------------------------------------------
-- TH derivations
----------------------------------------------------------------------------

-- T and Ty are mutually recursive, so their derivation quotes need to be aware of each other
$(mconcat
  [ deriveJSON morleyAesonOptions ''Ty
  , deriveJSON morleyAesonOptions ''T
  ])

deriveJSON morleyAesonOptions ''ParameterType