-- SPDX-FileCopyrightText: 2020 Tocqueville Group
--
-- SPDX-License-Identifier: LicenseRef-MIT-TQ

-- | Module, providing @Notes t@ data type, which holds annotations for a
-- given type @t@.
--
-- Annotation type @Notes t@ is a tree, each leaf is either a star (@*@) or a
-- constructor holding some annotation data for a given type @t@.
-- Star corresponds to the case when given Michelson type contains no
-- annotations.
--
-- This module also provides type class 'Converge' along with some
-- utility functions which are used to combine two annotations trees
-- `a` and `b` into a new one `c` in such a way that `c` can be obtained from
-- both `a` and `b` by replacing some @*@ leafs with type or/and field
-- annotations.

module Michelson.Typed.Annotation
  ( Notes (..)
  , AnnConvergeError(..)
  , converge
  , convergeAnns
  , insertTypeAnn
  , orAnn
  , isStar
  , starNotes
  , notesSing
  , notesT
  ) where

import Data.Default (Default(..))
import qualified Data.Kind as Kind
import Data.Singletons (Sing, SingI(..))
import Fmt (Buildable(..), (+|), (|+))
import Text.PrettyPrint.Leijen.Text (Doc, (<+>))
import qualified Text.Show

import Michelson.Printer.Util
  (RenderDoc(..), addParens, buildRenderDoc, doesntNeedParens, needsParens, printDocS)
import Michelson.Typed.Sing
import Michelson.Typed.T (T(..))
import Michelson.Untyped.Annotation
  (Annotation, FieldAnn, TypeAnn, fullAnnSet, noAnn, singleAnnSet, unifyAnn)
import Util.TH
import Util.Typeable

-- | Data type, holding annotation data for a given Michelson type @t@.
--
-- Each constructor corresponds to exactly one constructor of 'T'
-- and holds all type and field annotations that can be attributed to a
-- Michelson type corresponding to @t@.
data Notes t where
  NTKey       :: TypeAnn -> Notes 'TKey
  NTUnit      :: TypeAnn -> Notes 'TUnit
  NTSignature :: TypeAnn -> Notes 'TSignature
  NTChainId   :: TypeAnn -> Notes 'TChainId
  NTOption    :: TypeAnn -> Notes t -> Notes ('TOption t)
  NTList      :: TypeAnn -> Notes t -> Notes ('TList t)
  NTSet       :: TypeAnn -> Notes t -> Notes ('TSet t)
  NTOperation :: TypeAnn -> Notes 'TOperation
  NTContract  :: TypeAnn -> Notes t -> Notes ('TContract t)
  NTPair      :: TypeAnn -> FieldAnn -> FieldAnn
              -> Notes p -> Notes q -> Notes ('TPair p q)
  NTOr        :: TypeAnn -> FieldAnn -> FieldAnn
              -> Notes p -> Notes q -> Notes ('TOr p q)
  NTLambda    :: TypeAnn -> Notes p -> Notes q -> Notes ('TLambda p q)
  NTMap       :: TypeAnn -> Notes k -> Notes v -> Notes ('TMap k v)
  NTBigMap    :: TypeAnn -> Notes k -> Notes v -> Notes ('TBigMap k v)
  NTInt       :: TypeAnn -> Notes 'TInt
  NTNat       :: TypeAnn -> Notes 'TNat
  NTString    :: TypeAnn -> Notes 'TString
  NTBytes     :: TypeAnn -> Notes 'TBytes
  NTMutez     :: TypeAnn -> Notes 'TMutez
  NTBool      :: TypeAnn -> Notes 'TBool
  NTKeyHash   :: TypeAnn -> Notes 'TKeyHash
  NTTimestamp :: TypeAnn -> Notes 'TTimestamp
  NTAddress   :: TypeAnn -> Notes 'TAddress

deriving stock instance Eq (Notes t)
$(deriveGADTNFData ''Notes)


instance Show (Notes t) where
  show = printDocS True . renderDoc doesntNeedParens

instance Buildable (Notes t) where
  build = buildRenderDoc

instance RenderDoc (Notes t) where
  renderDoc pn n = addParens pn $ case n of
    NTInt ta                -> "NTInt" <+> rendT ta
    NTNat ta                -> "NTNat" <+> rendT ta
    NTString ta             -> "NTString" <+> rendT ta
    NTBytes ta              -> "NTBytes" <+> rendT ta
    NTMutez ta              -> "NTMutez" <+> rendT ta
    NTBool ta               -> "NTBool" <+> rendT ta
    NTKeyHash ta            -> "NTKeyHash" <+> rendT ta
    NTTimestamp ta          -> "NTTimestamp" <+> rendT ta
    NTAddress ta            -> "NTAddress" <+> rendT ta
    NTKey ta                -> "NTKey" <+> rendT ta
    NTUnit ta               -> "NTUnit" <+> rendT ta
    NTSignature ta          -> "NTSignature" <+> rendT ta
    NTChainId ta            -> "NTChainId" <+> rendT ta
    NTOption ta nt          -> "NTOption" <+> rendT ta <+> rendN nt
    NTList ta nt            -> "NTList" <+> rendT ta <+> rendN nt
    NTSet ta nt             -> "NTSet" <+> rendT ta <+> rendN nt
    NTOperation ta          -> "NTOperation" <+> rendT ta
    NTContract ta nt        -> "NTContract" <+> rendT ta <+> rendN nt
    NTPair ta fa1 fa2 np nq -> "NTPair" <+> rendTF2 ta fa1 fa2 <+> rendN np <+> rendN nq
    NTOr ta fa1 fa2 np nq   -> "NTOr" <+> rendTF2 ta fa1 fa2 <+> rendN np <+> rendN nq
    NTLambda ta np nq       -> "NTLambda" <+> rendT ta <+> rendN np <+> rendN nq
    NTMap ta np nq          -> "NTMap" <+> rendT ta <+> rendN np <+> rendN nq
    NTBigMap ta np nq       -> "NTBigMap" <+> rendT ta <+> rendN np <+> rendN nq
    where
      rendN :: Notes n -> Doc
      rendN = renderDoc needsParens

      rendT :: TypeAnn -> Doc
      rendT = renderDoc doesntNeedParens . singleAnnSet

      rendTF2 :: TypeAnn -> FieldAnn -> FieldAnn -> Doc
      rendTF2 t f1 f2 = renderDoc doesntNeedParens $ fullAnnSet [t] [f1, f2] []

-- | Forget information about annotations, pick singleton with the same type.
--
-- Note: currently we cannot derive 'Sing' from 'Notes' without 'SingI' because
-- for comparable types notes do not remember which exact comparable was used.
notesSing :: SingI t => Notes t -> Sing t
notesSing _ = sing

-- | Get term-level type of notes.
notesT :: SingI t => Notes t -> T
notesT = fromSingT . notesSing

-- | In memory of `NStar` constructor.
--   Generates notes with no annotations.
starNotes :: forall t. SingI t => Notes t
starNotes = case sing @t of
  STInt -> NTInt noAnn
  STNat -> NTNat noAnn
  STString -> NTString noAnn
  STBytes -> NTBytes noAnn
  STMutez -> NTMutez noAnn
  STBool -> NTBool noAnn
  STKeyHash -> NTKeyHash noAnn
  STTimestamp -> NTTimestamp noAnn
  STAddress -> NTAddress noAnn
  STKey -> NTKey noAnn
  STUnit -> NTUnit noAnn
  STSignature -> NTSignature noAnn
  STChainId -> NTChainId noAnn
  STOperation -> NTOperation noAnn
  STSet _ -> NTSet noAnn starNotes
  STList _ -> NTList noAnn starNotes
  STOption _ -> NTOption noAnn starNotes
  STContract _ -> NTContract noAnn starNotes
  STMap _ _ -> NTMap noAnn starNotes starNotes
  STBigMap _ _ -> NTBigMap noAnn starNotes starNotes
  STPair _ _ -> NTPair noAnn noAnn noAnn starNotes starNotes
  STOr _ _ -> NTOr noAnn noAnn noAnn starNotes starNotes
  STLambda _ _ -> NTLambda noAnn starNotes starNotes

-- | Checks if no annotations are present.
isStar :: SingI t => Notes t -> Bool
isStar = (== starNotes)

orAnn :: Annotation t -> Annotation t -> Annotation t
orAnn a b = bool a b (a == def)

-- | Combines two annotations trees `a` and `b` into a new one `c`
-- in such a way that `c` can be obtained from both `a` and `b` by replacing
-- some empty leaves with type or/and field annotations.
converge :: Notes t -> Notes t -> Either AnnConvergeError (Notes t)
converge n1 n2 = case (n1, n2) of
  (NTInt a, NTInt b) -> NTInt <$> convergeAnns a b
  (NTNat a, NTNat b) -> NTNat <$> convergeAnns a b
  (NTString a, NTString b) -> NTString <$> convergeAnns a b
  (NTBytes a, NTBytes b) -> NTBytes <$> convergeAnns a b
  (NTMutez a, NTMutez b) -> NTMutez <$> convergeAnns a b
  (NTBool a, NTBool b) -> NTBool <$> convergeAnns a b
  (NTKeyHash a, NTKeyHash b) -> NTKeyHash <$> convergeAnns a b
  (NTTimestamp a, NTTimestamp b) -> NTTimestamp <$> convergeAnns a b
  (NTAddress a, NTAddress b) -> NTAddress <$> convergeAnns a b
  (NTKey a, NTKey b) -> NTKey <$> convergeAnns a b
  (NTUnit a, NTUnit b) -> NTUnit <$> convergeAnns a b
  (NTSignature a, NTSignature b) ->
    NTSignature <$> convergeAnns a b
  (NTChainId a, NTChainId b) ->
    NTChainId <$> convergeAnns a b
  (NTOption a n, NTOption b m) ->
    NTOption <$> convergeAnns a b <*> converge n m
  (NTList a n, NTList b m) ->
    NTList <$> convergeAnns a b <*> converge n m
  (NTSet a n, NTSet b m) ->
    NTSet <$> convergeAnns a b <*> converge n m
  (NTOperation a, NTOperation b) ->
    NTOperation <$> convergeAnns a b
  (NTContract a n, NTContract b m) ->
    NTContract <$> convergeAnns a b <*> converge n m
  (NTPair a pF qF pN qN, NTPair b pG qG pM qM) ->
    NTPair <$> convergeAnns a b <*> convergeAnns pF pG
           <*> convergeAnns qF qG <*> converge pN pM <*> converge qN qM
  (NTOr a pF qF pN qN, NTOr b pG qG pM qM) ->
    NTOr <$> convergeAnns a b <*> convergeAnns pF pG <*> convergeAnns qF qG
         <*> converge pN pM <*> converge qN qM
  (NTLambda a pN qN, NTLambda b pM qM) ->
    NTLambda <$> convergeAnns a b <*> converge pN pM <*> converge qN qM
  (NTMap a kN vN, NTMap b kM vM) ->
    NTMap <$> convergeAnns a b <*> converge kN kM <*> converge vN vM
  (NTBigMap a kN vN, NTBigMap b kM vM) ->
    NTBigMap <$> convergeAnns a b <*> converge kN kM <*> converge vN vM

-- | Insert the provided type annotation into the provided notes.
insertTypeAnn :: forall (b :: T). TypeAnn -> Notes b -> Notes b
insertTypeAnn nt s = case s of
  NTInt _ -> NTInt nt
  NTNat _ -> NTNat nt
  NTString _ -> NTString nt
  NTBytes _ -> NTBytes nt
  NTMutez _ -> NTMutez nt
  NTBool _ -> NTBool nt
  NTKeyHash _ -> NTKeyHash nt
  NTTimestamp _ -> NTTimestamp nt
  NTAddress _ -> NTAddress nt
  NTKey _ -> NTKey nt
  NTUnit _ -> NTUnit nt
  NTSignature _ -> NTSignature nt
  NTOption _ n1  -> NTOption nt n1
  NTList _ n1 -> NTList nt n1
  NTSet _ n1 -> NTSet nt n1
  NTOperation _ -> NTOperation nt
  NTContract _ n1 -> NTContract nt n1
  NTPair _ n1 n2 n3 n4 -> NTPair nt n1 n2 n3 n4
  NTOr _ n1 n2 n3 n4 -> NTOr nt n1 n2 n3 n4
  NTLambda _ n1 n2 -> NTLambda nt n1 n2
  NTMap _ n1 n2 -> NTMap nt n1 n2
  NTBigMap _ n1 n2 -> NTBigMap nt n1 n2
  NTChainId _ -> NTChainId nt

data AnnConvergeError where
  AnnConvergeError
    :: forall (tag :: Kind.Type).
       (Buildable (Annotation tag), Show (Annotation tag), Typeable tag)
    => Annotation tag -> Annotation tag -> AnnConvergeError

deriving stock instance Show AnnConvergeError

instance Eq AnnConvergeError where
  AnnConvergeError ann1 ann2 == AnnConvergeError ann1' ann2' =
    (ann1 `eqParam1` ann1') && (ann2 `eqParam1` ann2')

instance Buildable AnnConvergeError where
  build (AnnConvergeError ann1 ann2) =
    "Annotations do not converge: " +| ann1 |+ " /= " +| ann2 |+ ""

-- | Converge two type or field notes (which may be wildcards).
convergeAnns
  :: forall (tag :: Kind.Type).
     (Buildable (Annotation tag), Show (Annotation tag), Typeable tag)
  => Annotation tag -> Annotation tag -> Either AnnConvergeError (Annotation tag)
convergeAnns a b = maybe (Left $ AnnConvergeError a b)
                          pure $ unifyAnn a b

$(deriveGADTNFData ''AnnConvergeError)