-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA
{-# LANGUAGE DeriveLift #-}

{-# OPTIONS_GHC -Wno-orphans #-}

-- | 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 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 @*@ leaves with type or/and field
-- annotations.

module Morley.Michelson.Typed.Annotation
  ( Notes (..)
  , insertTypeAnn
  , isStar
  , starNotes
  , mkUType
  , notesSing
  , notesT
    -- * Helpers
  , AnnVar
  , Anns(.., Anns1, Anns2, Anns2', Anns3, Anns3', Anns3'', Anns4, Anns4'', Anns5')
  , AnnotateInstr(..)
  ) where

import Data.Default (Default(def))
import Data.Singletons (Sing, SingI(..), fromSing)
import Language.Haskell.TH.Syntax (Lift)

import Morley.Michelson.Printer.Util (RenderDoc(..))
import Morley.Michelson.Typed.Sing
import Morley.Michelson.Typed.T (T(..))
import Morley.Michelson.Untyped qualified as Un
import Morley.Michelson.Untyped.Annotation (Annotation, FieldAnn, TypeAnn, VarAnn, noAnn)
import Morley.Util.Peano qualified as Peano
import Morley.Util.PeanoNatural (singPeanoVal)
import Morley.Util.TH

{-# ANN module ("HLint: ignore Avoid lambda using `infix`" :: Text) #-}

-- | 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)
  NTTicket    :: TypeAnn -> Notes t -> Notes ('TTicket 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
  NTChest     :: TypeAnn -> Notes 'TChest
  NTChestKey  :: TypeAnn -> Notes 'TChestKey
  NTNever     :: TypeAnn -> Notes 'TNever
  NTSaplingState :: forall (n :: Peano.Peano). TypeAnn -> Sing n -> Notes ('TSaplingState n)
  NTSaplingTransaction :: forall (n :: Peano.Peano). TypeAnn -> Sing n -> Notes ('TSaplingTransaction n)

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

instance RenderDoc (Notes t) where
  renderDoc pn = renderDoc pn . mkUType

-- | Forget information about annotations, pick singleton with the same type.
notesSing :: Notes t -> Sing t
notesSing = \case
  NTInt _        -> sing
  NTNat _        -> sing
  NTString _     -> sing
  NTBytes _      -> sing
  NTMutez _      -> sing
  NTBool _       -> sing
  NTKeyHash _    -> sing
  NTBls12381Fr _ -> sing
  NTBls12381G1 _ -> sing
  NTBls12381G2 _ -> sing
  NTTimestamp _  -> sing
  NTAddress _    -> sing
  NTKey _        -> sing
  NTUnit _       -> sing
  NTSignature _  -> sing
  NTChainId _    -> sing
  NTOption _ n   -> STOption (notesSing n)
  NTList _ n     -> STList (notesSing n)
  NTSet _ n      -> STSet (notesSing n)
  NTOperation _  -> sing
  NTChest _      -> sing
  NTChestKey _   -> sing
  NTNever _      -> sing
  NTSaplingState _ s -> STSaplingState s
  NTSaplingTransaction _ s -> STSaplingTransaction s
  NTContract _ n -> STContract (notesSing n)
  NTTicket _ n   -> STTicket (notesSing n)
  NTPair _ _ _ _ _ l r -> STPair   (notesSing l) (notesSing r)
  NTOr _ _ _ l r       -> STOr     (notesSing l) (notesSing r)
  NTLambda _ l r       -> STLambda (notesSing l) (notesSing r)
  NTMap _ l r          -> STMap    (notesSing l) (notesSing r)
  NTBigMap _ l r       -> STBigMap (notesSing l) (notesSing r)

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

-- | Get the term-level type of notes, preserving annotations.
mkUType :: Notes x -> Un.Ty
mkUType = \case
  NTInt tn        -> Un.Ty Un.TInt tn
  NTNat tn        -> Un.Ty Un.TNat tn
  NTString tn     -> Un.Ty Un.TString tn
  NTBytes tn      -> Un.Ty Un.TBytes tn
  NTMutez tn      -> Un.Ty Un.TMutez tn
  NTBool tn       -> Un.Ty Un.TBool tn
  NTKeyHash tn    -> Un.Ty Un.TKeyHash tn
  NTBls12381Fr tn -> Un.Ty Un.TBls12381Fr tn
  NTBls12381G1 tn -> Un.Ty Un.TBls12381G1 tn
  NTBls12381G2 tn -> Un.Ty Un.TBls12381G2 tn
  NTTimestamp tn  -> Un.Ty Un.TTimestamp tn
  NTAddress tn    -> Un.Ty Un.TAddress tn
  NTKey tn        -> Un.Ty Un.TKey tn
  NTUnit tn       -> Un.Ty Un.TUnit tn
  NTSignature tn  -> Un.Ty Un.TSignature tn
  NTChainId tn    -> Un.Ty Un.TChainId tn
  NTOption tn n   -> Un.Ty (Un.TOption $ mkUType n) tn
  NTList tn n     -> Un.Ty (Un.TList $ mkUType n) tn
  NTSet tn n      -> Un.Ty (Un.TSet $ mkUType n) tn
  NTOperation tn  -> Un.Ty Un.TOperation tn
  NTChest tn      -> Un.Ty Un.TChest tn
  NTChestKey tn   -> Un.Ty Un.TChestKey tn
  NTNever tn      -> Un.Ty Un.TNever tn
  NTSaplingState tn s -> Un.Ty (Un.TSaplingState (singPeanoVal s)) tn
  NTSaplingTransaction tn s -> Un.Ty (Un.TSaplingTransaction (singPeanoVal s)) tn
  NTContract tn n -> Un.Ty (Un.TContract $ mkUType n) tn
  NTTicket tn n   -> Un.Ty (Un.TTicket $ mkUType n) tn
  NTPair tn fl fr vl vr nl nr ->
    Un.Ty (Un.TPair fl fr vl vr (mkUType nl) (mkUType nr)) tn
  NTOr tn fl fr nl nr -> Un.Ty (Un.TOr fl fr (mkUType nl) (mkUType nr)) tn
  NTLambda tn np nq   -> Un.Ty (Un.TLambda (mkUType np) (mkUType nq)) tn
  NTMap tn nk nv      -> Un.Ty (Un.TMap (mkUType nk) (mkUType nv)) tn
  NTBigMap tn nk nv   -> Un.Ty (Un.TBigMap (mkUType nk) (mkUType nv)) tn


-- | In memory of @NStar@ constructor.
--   Generates notes with no annotations.
starNotes :: forall t. SingI t => Notes t
starNotes = starNotes' sing

starNotes' :: Sing t -> Notes t
starNotes' = \case
  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
  STSaplingState s -> NTSaplingState noAnn s
  STSaplingTransaction s -> NTSaplingTransaction noAnn s
  STSignature -> NTSignature noAnn
  STChainId -> NTChainId noAnn
  STOperation -> NTOperation noAnn
  STSet t -> NTSet noAnn $ starNotes' t
  STList t -> NTList noAnn $ starNotes' t
  STOption t -> NTOption noAnn $ starNotes' t
  STContract t -> NTContract noAnn $ starNotes' t
  STTicket t -> NTTicket noAnn $ starNotes' t
  STMap t t' -> NTMap noAnn (starNotes' t) (starNotes' t')
  STBigMap t t' -> NTBigMap noAnn (starNotes' t) (starNotes' t')
  STPair t t' ->
    NTPair noAnn noAnn noAnn noAnn noAnn (starNotes' t) (starNotes' t')
  STOr t t' -> NTOr noAnn noAnn noAnn (starNotes' t) (starNotes' t')
  STLambda t t' -> NTLambda noAnn (starNotes' t) (starNotes' t')
  STChest -> NTChest noAnn
  STChestKey -> NTChestKey noAnn

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

-- | 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
  NTTicket _ n1 -> NTTicket 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
  NTChest _ -> NTChest nt
  NTChestKey _ -> NTChestKey nt
  NTNever _ -> NTNever nt
  NTSaplingState _ n -> NTSaplingState nt n
  NTSaplingTransaction _ n -> NTSaplingTransaction nt n

-- | A typed heterogenous list of annotations. Simplified pattern synonyms for
-- common use cases are provided.
data Anns xs where
  AnnsCons :: Typeable tag => !(Annotation tag) -> Anns xs -> Anns (Annotation tag ': xs)
  AnnsTyCons :: SingI t => !(Notes t) -> Anns xs -> Anns (Notes t ': xs)
  AnnsNil :: Anns '[]

infixr 5 `AnnsCons`
infixr 5 `AnnsTyCons`

deriveGADTNFData ''Anns

deriving stock instance (Each '[Show] rs) => Show (Anns rs)
deriving stock instance (Eq r, Eq (Anns rs)) => Eq (Anns (r ': rs))

instance Eq (Anns '[]) where
  AnnsNil == AnnsNil = True

-- | 'Anns' only containing a single 'VarAnn'.
type AnnVar = Anns '[VarAnn]

instance Default (Anns '[]) where
  def = AnnsNil

instance (Typeable tag, Default (Anns xs)) => Default (Anns (Annotation tag ': xs)) where
  def = noAnn `AnnsCons` def

instance (SingI t, Default (Anns xs)) => Default (Anns (Notes t ': xs)) where
  def = starNotes `AnnsTyCons` def

type family AnnotateInstrArg (xs :: [Type]) r where
  AnnotateInstrArg (Notes _ ': xs) r = Un.Ty -> AnnotateInstrArg xs r
  AnnotateInstrArg (x ': xs) r = x -> AnnotateInstrArg xs r
  AnnotateInstrArg '[] r = r

-- | Utility typeclass to simplify extracting annotations from 'Anns' and
-- passing those as arguments to an untyped instruction data constructor.
class AnnotateInstr (xs :: [Type]) r where
  annotateInstr :: Anns xs -> AnnotateInstrArg xs r -> r

instance AnnotateInstr '[] r where
  annotateInstr (AnnsNil) c = c

instance AnnotateInstr xs r => AnnotateInstr (Annotation tag ': xs) r where
  annotateInstr (AnnsCons va xs) c = annotateInstr @_ @r xs $ c va

instance AnnotateInstr xs r => AnnotateInstr (Notes t ': xs) r where
  annotateInstr (AnnsTyCons va xs) c = annotateInstr @_ @r xs $ c (mkUType va)

-- | Convenience pattern synonym matching a single simple annotation.
pattern Anns1 :: Typeable a => Annotation a -> Anns '[Annotation a]
pattern Anns1 x = x `AnnsCons` AnnsNil

-- | Convenience pattern synonym matching two simple annotations.
pattern Anns2
  :: Each '[Typeable] '[a, b]
  => Annotation a
  -> Annotation b
  -> Anns '[Annotation a, Annotation b]
pattern Anns2 x y = x `AnnsCons` y `AnnsCons` AnnsNil

-- | Convenience pattern synonym matching two annotations, first being
-- a simple one, the second being 'Notes', corresponding to an annotated
-- type.
pattern Anns2' :: (Typeable a, SingI t) => Annotation a -> Notes t -> Anns '[Annotation a, Notes t]
pattern Anns2' x y = x `AnnsCons` y `AnnsTyCons` AnnsNil

-- | Convenience pattern synonym matching three simple annotations.
pattern Anns3
  :: Each '[Typeable] '[a, b, c]
  => Annotation a
  -> Annotation b
  -> Annotation c
  -> Anns '[Annotation a, Annotation b, Annotation c]
pattern Anns3 x y z = x `AnnsCons` y `AnnsCons` z `AnnsCons` AnnsNil

-- | Convenience pattern synonym matching three annotations, first two being
-- simple, the last one being 'Notes', corresponding to an annotated
-- type.
pattern Anns3'
  :: (Each '[Typeable] '[a, b], SingI t)
  => Annotation a
  -> Annotation b
  -> Notes t
  -> Anns '[Annotation a, Annotation b, Notes t]
pattern Anns3' x y z = x `AnnsCons` y `AnnsCons` z `AnnsTyCons` AnnsNil

-- | Convenience pattern synonym matching three annotations, first being
-- a simple one, the last two being 'Notes', corresponding to annotated
-- types.
pattern Anns3''
  :: (Typeable a, SingI t, SingI u)
  => Annotation a
  -> Notes t
  -> Notes u
  -> Anns '[Annotation a, Notes t, Notes u]
pattern Anns3'' x y z = x `AnnsCons` y `AnnsTyCons` z `AnnsTyCons` AnnsNil

-- | Convenience pattern synonym matching four simple annotations.
pattern Anns4
  :: Each '[Typeable] '[a, b, c, d]
  => Annotation a
  -> Annotation b
  -> Annotation c
  -> Annotation d
  -> Anns '[Annotation a, Annotation b, Annotation c, Annotation d]
pattern Anns4 x y z w = x `AnnsCons` y `AnnsCons` z `AnnsCons` w `AnnsCons` AnnsNil

-- | Convenience pattern synonym matching four annotations, first two being
-- simple, the last two being 'Notes', corresponding to annotated
-- types.
pattern Anns4''
  :: (Each '[Typeable] '[a, b], SingI t, SingI u)
  => Annotation a
  -> Annotation b
  -> Notes t
  -> Notes u
  -> Anns '[Annotation a, Annotation b, Notes t, Notes u]
pattern Anns4'' x y z w
  = x `AnnsCons` y `AnnsCons` z `AnnsTyCons` w `AnnsTyCons` AnnsNil

-- | Convenience pattern synonym matching five annotations, first four being
-- simple, the last one being 'Notes', corresponding to an annotated
-- type.
pattern Anns5'
  :: (Each '[Typeable] '[a, b, c, d], SingI t)
  => Annotation a
  -> Annotation b
  -> Annotation c
  -> Annotation d
  -> Notes t
  -> Anns '[Annotation a, Annotation b, Annotation c, Annotation d, Notes t]
pattern Anns5' x y z v w
  = x `AnnsCons` y `AnnsCons` z `AnnsCons` v `AnnsCons` w `AnnsTyCons` AnnsNil

{-# COMPLETE Anns1 #-}
{-# COMPLETE Anns2 #-}
{-# COMPLETE Anns2' #-}
{-# COMPLETE Anns3 #-}
{-# COMPLETE Anns3' #-}
{-# COMPLETE Anns3'' #-}
{-# COMPLETE Anns4 #-}
{-# COMPLETE Anns4'' #-}
{-# COMPLETE Anns5' #-}