-- | Michelson types represented in untyped model.

module Michelson.Untyped.Type
  ( Type (..)
  , Comparable (..)
  , compToType
  , typeToComp
  , T (..)
  , CT (..)
  , pattern Tint
  , pattern Tnat
  , pattern Tstring
  , pattern Tbytes
  , pattern Tmutez
  , pattern Tbool
  , pattern Tkey_hash
  , pattern Ttimestamp
  , pattern Taddress
  , tint
  , tnat
  , tstring
  , tbytes
  , tmutez
  , tbool
  , tkeyHash
  , ttimestamp
  , taddress
  , toption
  , tpair
  , tor
  , tyint
  , tynat
  , tyunit
  , tybool
  , typair
  , tyor
  , tyImplicitAccountParam
  , isAtomicType
  , isKey
  , isSignature
  , isComparable
  , isMutez
  , isKeyHash
  , isBool
  , isString
  , isInteger
  , isTimestamp
  , isNat
  , isInt
  , isBytes
  ) where

import Data.Aeson.TH (defaultOptions, deriveJSON)
import Data.Data (Data(..))
import Formatting.Buildable (Buildable(build))
import Language.Haskell.TH.Lift (deriveLift)
import Prelude hiding ((<$>))
import Text.PrettyPrint.Leijen.Text (Doc, align, softbreak, (<$>), (<+>))

import Michelson.Printer.Util
  (Prettier(..), RenderContext, RenderDoc(..), addParens, buildRenderDoc, doesntNeedParens,
  needsParens, wrapInParens)
import Michelson.Untyped.Annotation
  (AnnotationSet, FieldAnn, TypeAnn, emptyAnnSet, fullAnnSet, noAnn, singleAnnSet)

-- Annotated type
data Type
  = Type ~T TypeAnn
  deriving stock (Eq, Show, Data, Generic)

instance NFData Type

instance RenderDoc Comparable where
  renderDoc np (Comparable ct ta) =
    addParens np $
    renderCT ct <+> renderDoc doesntNeedParens (singleAnnSet ta)

instance RenderDoc (Prettier Type) where
  renderDoc pn (Prettier w) = case w of
    (Type t ta) -> renderType t False pn (singleAnnSet ta)

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

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

-- 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
  -> Bool
  -> RenderContext
  -> AnnotationSet
  -> Doc
renderType t forceSingleLine pn annSet =
  let annDoc = renderDoc doesntNeedParens annSet
      recRenderer t' annSet' = renderType t' forceSingleLine needsParens annSet'
      renderBranches d1 d2 =
        if forceSingleLine
        then (d1 <+> d2)
        else align $ softbreak <> (d1 <$> d2)
  in
  case t of
    Tc ct             -> wrapInParens pn $ renderCT ct :| [annDoc]
    TKey              -> wrapInParens pn $ "key"  :| [annDoc]
    TUnit             -> wrapInParens pn $ "unit" :| [annDoc]
    TSignature        -> wrapInParens pn $ "signature" :| [annDoc]
    TChainId          -> wrapInParens pn $ "chain_id" :| [annDoc]
    TOperation        -> wrapInParens pn $ "operation" :| [annDoc]

    TOption (Type t1 ta1) ->
      addParens pn $
      "option" <+> annDoc <+> recRenderer t1 (singleAnnSet ta1)

    TList (Type t1 ta1)       ->
      addParens pn $
      "list" <+> annDoc <+> recRenderer t1 (singleAnnSet ta1)

    TSet (Comparable ct1 ta1) ->
      addParens pn $
      "set" <+> annDoc <+> recRenderer (Tc ct1) (singleAnnSet ta1)

    TContract (Type t1 ta1)   ->
      addParens pn $
      "contract" <+> annDoc <+> recRenderer t1 (singleAnnSet ta1)

    TPair fa1 fa2 (Type t1 ta1) (Type t2 ta2) ->
      addParens pn $
        "pair" <+> annDoc <+>
          renderBranches
            (recRenderer t1 $ fullAnnSet [ta1] [fa1] [])
            (recRenderer t2 $ fullAnnSet [ta2] [fa2] [])

    TOr fa1 fa2 (Type t1 ta1) (Type t2 ta2) ->
      addParens pn $
        "or" <+> annDoc <+>
          renderBranches
            (recRenderer t1 $ fullAnnSet [ta1] [fa1] [])
            (recRenderer t2 $ fullAnnSet [ta2] [fa2] [])

    TLambda (Type t1 ta1) (Type t2 ta2) ->
      addParens pn $
        "lambda" <+> annDoc <+>
          renderBranches
            (recRenderer t1 $ singleAnnSet ta1)
            (recRenderer t2 $ singleAnnSet ta2)

    TMap (Comparable ct1 ta1) (Type t2 ta2) ->
      addParens pn $
        "map" <+> annDoc <+>
          renderBranches
            (recRenderer (Tc ct1) $ singleAnnSet ta1)
            (recRenderer t2 $ singleAnnSet ta2)

    TBigMap (Comparable ct1 ta1) (Type t2 ta2) ->
      addParens pn $
        "big_map" <+> annDoc <+>
          renderBranches
            (recRenderer (Tc ct1) $ singleAnnSet ta1)
            (recRenderer t2 $ singleAnnSet ta2)

renderCT :: CT -> Doc
renderCT = \case
  CInt       -> "int"
  CNat       -> "nat"
  CString    -> "string"
  CMutez     -> "mutez"
  CBool      -> "bool"
  CKeyHash   -> "key_hash"
  CTimestamp -> "timestamp"
  CBytes     -> "bytes"
  CAddress   -> "address"

instance RenderDoc CT where
  renderDoc _ = renderCT

instance Buildable Type where
  build = buildRenderDoc

-- Annotated Comparable Sub-type
data Comparable = Comparable CT TypeAnn
  deriving stock (Eq, Show, Data, Generic)

instance Buildable Comparable where
  build = buildRenderDoc

instance NFData Comparable

compToType :: Comparable -> Type
compToType (Comparable ct tn) = Type (Tc ct) tn

typeToComp :: Type -> Maybe Comparable
typeToComp (Type (Tc ct) tn) = Just $ Comparable ct tn
typeToComp _ = Nothing

-- Michelson Type
data T =
    Tc CT
  | TKey
  | TUnit
  | TSignature
  | TChainId
  | TOption Type
  | TList Type
  | TSet Comparable
  | TOperation
  | TContract Type
  | TPair FieldAnn FieldAnn Type Type
  | TOr FieldAnn FieldAnn Type Type
  | TLambda Type Type
  | TMap Comparable Type
  | TBigMap Comparable Type
  deriving stock (Eq, Show, Data, Generic)

instance Buildable T where
  build = buildRenderDoc

instance NFData T

-- Comparable Sub-Type
data CT =
    CInt
  | CNat
  | CString
  | CBytes
  | CMutez
  | CBool
  | CKeyHash
  | CTimestamp
  | CAddress
  deriving stock (Eq, Ord, Show, Data, Enum, Bounded, Generic)

instance Buildable CT where
  build = buildRenderDoc

instance NFData CT

pattern Tint :: T
pattern Tint = Tc CInt

pattern Tnat :: T
pattern Tnat = Tc CNat

pattern Tstring :: T
pattern Tstring = Tc CString

pattern Tbytes :: T
pattern Tbytes = Tc CBytes

pattern Tmutez :: T
pattern Tmutez = Tc CMutez

pattern Tbool :: T
pattern Tbool = Tc CBool

pattern Tkey_hash :: T
pattern Tkey_hash = Tc CKeyHash

pattern Ttimestamp :: T
pattern Ttimestamp = Tc CTimestamp

pattern Taddress :: T
pattern Taddress = Tc CAddress

tint :: T
tint = Tc CInt

tnat :: T
tnat = Tc CNat

tstring :: T
tstring = Tc CString

tbytes :: T
tbytes = Tc CBytes

tmutez :: T
tmutez = Tc CMutez

tbool :: T
tbool = Tc CBool

tkeyHash :: T
tkeyHash = Tc CKeyHash

ttimestamp :: T
ttimestamp = Tc CTimestamp

taddress :: T
taddress = Tc CAddress

toption :: Type -> T
toption t = TOption t

tpair :: Type -> Type -> T
tpair l r = TPair noAnn noAnn l r

tor :: Type -> Type -> T
tor l r = TOr noAnn noAnn l r

tyint :: Type
tyint = Type tint noAnn

tynat :: Type
tynat = Type tnat noAnn

tyunit :: Type
tyunit = Type TUnit noAnn

tybool :: Type
tybool = Type tbool noAnn

typair :: Type -> Type -> Type
typair l r = Type (tpair l r) noAnn

tyor :: Type -> Type -> Type
tyor l r = Type (tor l r) noAnn

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

isAtomicType :: Type -> Bool
isAtomicType t@(Type _ tAnn) | tAnn == noAnn =
    isComparable t || isKey t || isUnit t || isSignature t || isOperation t
isAtomicType _ = False

isKey :: Type -> Bool
isKey (Type TKey _) = True
isKey _              = False

isUnit :: Type -> Bool
isUnit (Type TUnit _) = True
isUnit _               = False

isSignature :: Type -> Bool
isSignature (Type TSignature _) = True
isSignature _                    = False

isOperation :: Type -> Bool
isOperation (Type TOperation _) = True
isOperation _                    = False

isComparable :: Type -> Bool
isComparable (Type (Tc _) _) = True
isComparable _ = False

isMutez :: Type -> Bool
isMutez (Type (Tc CMutez) _) = True
isMutez _ = False

isTimestamp :: Type -> Bool
isTimestamp (Type (Tc CTimestamp) _) = True
isTimestamp _ = False

isKeyHash :: Type -> Bool
isKeyHash (Type (Tc CKeyHash) _) = True
isKeyHash _ = False

isBool  :: Type -> Bool
isBool (Type (Tc CBool) _) = True
isBool _ = False

isString  :: Type -> Bool
isString (Type (Tc CString) _) = True
isString _ = False

isInteger :: Type -> Bool
isInteger a = isNat a || isInt a || isMutez a || isTimestamp a

isNat  :: Type -> Bool
isNat (Type (Tc CNat) _) = True
isNat _ = False

isInt  :: Type -> Bool
isInt (Type (Tc CInt) _) = True
isInt _ = False

isBytes :: Type -> Bool
isBytes (Type (Tc CBytes) _) = True
isBytes _ = False

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

deriveJSON defaultOptions ''Type
deriveJSON defaultOptions ''Comparable
deriveJSON defaultOptions ''T
deriveJSON defaultOptions ''CT

deriveLift ''Type
deriveLift ''Comparable
deriveLift ''T
deriveLift ''CT