-- 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
  , convergeDestrAnns
  , insertTypeAnn
  , isStar
  , starNotes
  , notesSing
  , notesT
  ) where

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, VarAnn, convergeVarAnns, fullAnnSet, noAnn, singleAnnSet, unifyAnn, unifyPairFieldAnn)
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
              -> VarAnn -> VarAnn
              -> 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
  NTBls12381Fr :: TypeAnn -> Notes 'TBls12381Fr
  NTBls12381G1 :: TypeAnn -> Notes 'TBls12381G1
  NTBls12381G2 :: TypeAnn -> Notes 'TBls12381G2
  NTTimestamp :: TypeAnn -> Notes 'TTimestamp
  NTAddress   :: TypeAnn -> Notes 'TAddress
  NTNever     :: TypeAnn -> Notes 'TNever

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
    NTBls12381Fr ta         -> "NTBls12381Fr" <+> rendT ta
    NTBls12381G1 ta         -> "NTBls12381G1" <+> rendT ta
    NTBls12381G2 ta         -> "NTBls12381G2" <+> 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
    NTNever ta              -> "NTNever" <+> 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 va1 va2 np nq -> "NTPair" <+> rendTFV2 ta fa1 fa2 va1 va2 <+> 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] []

      rendTFV2 :: TypeAnn -> FieldAnn -> FieldAnn -> VarAnn -> VarAnn -> Doc
      rendTFV2 t f1 f2 v1 v2 = renderDoc doesntNeedParens $ fullAnnSet [t] [f1, f2] [v1, v2]

-- | 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
  STBls12381Fr -> NTBls12381Fr noAnn
  STBls12381G1 -> NTBls12381G1 noAnn
  STBls12381G2 -> NTBls12381G2 noAnn
  STTimestamp -> NTTimestamp noAnn
  STAddress -> NTAddress noAnn
  STKey -> NTKey noAnn
  STUnit -> NTUnit noAnn
  STNever -> NTNever 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 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)

-- | 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
  (NTBls12381Fr a, NTBls12381Fr b) -> NTBls12381Fr <$> convergeAnns a b
  (NTBls12381G1 a, NTBls12381G1 b) -> NTBls12381G1 <$> convergeAnns a b
  (NTBls12381G2 a, NTBls12381G2 b) -> NTBls12381G2 <$> 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
  (NTNever a, NTNever b) -> NTNever <$> 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 pV qV pN qN, NTPair b pG qG pW qW pM qM) ->
    NTPair <$> convergeAnns a b
           <*> convergeAnns pF pG <*> convergeAnns qF qG
           <*> pure (convergeVarAnns pV pW) <*> pure (convergeVarAnns qV qW)
           <*> 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
  NTBls12381Fr _ -> NTBls12381Fr nt
  NTBls12381G1 _ -> NTBls12381G1 nt
  NTBls12381G2 _ -> NTBls12381G2 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 n5 n6 -> NTPair nt n1 n2 n3 n4 n5 n6
  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
  NTNever _ -> NTNever nt

data AnnConvergeError where
  AnnConvergeError
    :: forall (tag :: 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 |+ ""

convergeAnnsImpl
  :: forall (tag :: Type).
     (Buildable (Annotation tag), Show (Annotation tag), Typeable tag)
  => (Annotation tag -> Annotation tag -> Maybe (Annotation tag))
  -> Annotation tag -> Annotation tag -> Either AnnConvergeError (Annotation tag)
convergeAnnsImpl unify a b = maybe (Left $ AnnConvergeError a b) pure $ unify a b

-- | Converge two type or field notes (which may be wildcards).
convergeAnns
  :: forall (tag :: Type).
     (Buildable (Annotation tag), Show (Annotation tag), Typeable tag)
  => Annotation tag -> Annotation tag -> Either AnnConvergeError (Annotation tag)
convergeAnns = convergeAnnsImpl unifyAnn

-- | Converge two field notes in CAR or CDR, given that one of them may be a
-- special annotation.
convergeDestrAnns :: FieldAnn -> FieldAnn -> Either AnnConvergeError FieldAnn
convergeDestrAnns = convergeAnnsImpl unifyPairFieldAnn

$(deriveGADTNFData ''AnnConvergeError)