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

{-# LANGUAGE UndecidableSuperClasses #-}

-- | Type and field annotations for Lorentz types.
module Lorentz.Annotation
  ( FollowEntrypointFlag (..)
  , GenerateFieldAnnFlag (..)
  , HasAnnotation (..)
  , GHasAnnotation (..)
  , gGetAnnotationNoField
  , insertTypeAnn
  ) where

import qualified GHC.Generics as G
import Named (NamedF)

import Lorentz.Entrypoints.Helpers (ctorNameToAnn)
import Michelson.Text
import Michelson.Typed
  (BigMap, ContractRef(..), EpAddress, KnownIsoT, Notes(..), Operation, ToT, insertTypeAnn,
  starNotes)
import Michelson.Typed.Haskell.Value (GValueType)
import Michelson.Untyped (FieldAnn, TypeAnn, ann, noAnn)
import Tezos.Address
import Tezos.Core
import Tezos.Crypto
import Util.TypeLits

-- | Used in `GHasAnnotation` and `HasAnnotation` as a flag to track
-- whether or not it directly follows an entrypoint to avoid introducing
-- extra entrypoints.
data FollowEntrypointFlag = FollowEntrypoint | NotFollowEntrypoint

-- | Used in `GHasAnnotation` as a flag to track whether or not field/constructor
-- annotations should be generated.
data GenerateFieldAnnFlag = GenerateFieldAnn | NotGenerateFieldAnn

-- | Use this in the instance of @HasAnnotation@ when field annotations
-- should not be generated.
gGetAnnotationNoField
    :: forall a. (GHasAnnotation (G.Rep a), GValueType (G.Rep a) ~ ToT a)
    => FollowEntrypointFlag -> Notes (ToT a)
gGetAnnotationNoField = \_ -> fst $ gGetAnnotation @(G.Rep a) NotFollowEntrypoint NotGenerateFieldAnn

-- | This class defines the type and field annotations for a given type. Right now
-- the type annotations come from names in a named field, and field annotations are
-- generated from the record fields.
class HasAnnotation a where
  getAnnotation :: FollowEntrypointFlag -> Notes (ToT a)
  default getAnnotation
    :: (GHasAnnotation (G.Rep a), GValueType (G.Rep a) ~ ToT a)
    => FollowEntrypointFlag
    -> Notes (ToT a)
  getAnnotation b = fst $ gGetAnnotation @(G.Rep a) b GenerateFieldAnn

instance (HasAnnotation a, KnownSymbol name)
  => HasAnnotation (NamedF Identity a name) where
  getAnnotation b = insertTypeAnn (symbolAnn @name) $
    getAnnotation @a b
    where
      symbolAnn :: forall s. KnownSymbol s => TypeAnn
      symbolAnn = ann $ symbolValT' @s

instance (HasAnnotation (Maybe a), KnownSymbol name)
  => HasAnnotation (NamedF Maybe a name) where
  getAnnotation b = getAnnotation @(NamedF Identity (Maybe a) name) b

-- Primitive instances
instance (HasAnnotation a) => HasAnnotation (Maybe a) where
  getAnnotation _ = NTOption noAnn (getAnnotation @a NotFollowEntrypoint)

instance HasAnnotation ()

instance HasAnnotation Integer where
  getAnnotation _ = starNotes

instance HasAnnotation Natural where
  getAnnotation _ = starNotes

instance HasAnnotation MText where
  getAnnotation _ = starNotes

instance HasAnnotation Bool where
  getAnnotation _ = starNotes

instance HasAnnotation ByteString where
  getAnnotation _ = starNotes

instance HasAnnotation Mutez where
  getAnnotation _ = starNotes

instance HasAnnotation Address where
  getAnnotation _ = starNotes

instance HasAnnotation EpAddress where
  getAnnotation _ = starNotes

instance HasAnnotation KeyHash where
  getAnnotation _ = starNotes

instance HasAnnotation Timestamp where
  getAnnotation _ = starNotes

instance HasAnnotation PublicKey where
  getAnnotation _ = starNotes

instance HasAnnotation Signature where
  getAnnotation _ = starNotes

instance (HasAnnotation a) => HasAnnotation (ContractRef a) where
  getAnnotation _ = NTContract noAnn (getAnnotation @a NotFollowEntrypoint)

instance (HasAnnotation k, HasAnnotation v) => HasAnnotation (Map k v) where
  getAnnotation _ = NTMap noAnn (getAnnotation @k NotFollowEntrypoint) (getAnnotation @v NotFollowEntrypoint)

instance (HasAnnotation k, HasAnnotation v) => HasAnnotation (BigMap k v) where
  getAnnotation _ = NTBigMap noAnn (getAnnotation @k NotFollowEntrypoint) (getAnnotation @v NotFollowEntrypoint)

instance (KnownIsoT v) => HasAnnotation (Set v) where
  getAnnotation _ = starNotes

instance (HasAnnotation a) => HasAnnotation [a] where
  getAnnotation _ = NTList noAnn (getAnnotation @a NotFollowEntrypoint)

instance HasAnnotation Operation where
  getAnnotation _ = starNotes

instance (HasAnnotation a, HasAnnotation b) => HasAnnotation (a, b)
instance (HasAnnotation a, HasAnnotation b, HasAnnotation c) => HasAnnotation (a, b, c)
instance (HasAnnotation a, HasAnnotation b, HasAnnotation c, HasAnnotation d) => HasAnnotation (a, b, c, d)
instance (HasAnnotation a, HasAnnotation b, HasAnnotation c, HasAnnotation d, HasAnnotation e)
  => HasAnnotation (a, b, c, d, e)
instance (HasAnnotation a, HasAnnotation b, HasAnnotation c, HasAnnotation d, HasAnnotation e, HasAnnotation f)
  => HasAnnotation (a, b, c, d, e, f)
instance
  ( HasAnnotation a, HasAnnotation b, HasAnnotation c, HasAnnotation d, HasAnnotation e
  , HasAnnotation f, HasAnnotation g)
  => HasAnnotation (a, b, c, d, e, f, g)

-- | A Generic @HasAnnotation@ implementation
class GHasAnnotation a where
  gGetAnnotation :: FollowEntrypointFlag -> GenerateFieldAnnFlag -> (Notes (GValueType a), FieldAnn)

instance GHasAnnotation G.U1 where
  gGetAnnotation _ _ = (starNotes, noAnn)

instance (GHasAnnotation x)
  => GHasAnnotation (G.M1 G.S ('G.MetaSel 'Nothing b c d) x)
  where
  gGetAnnotation b b2 = gGetAnnotation @x b b2

instance (GHasAnnotation x, KnownSymbol a)
  => GHasAnnotation (G.M1 G.S ('G.MetaSel ('Just a) b c d) x)
  where
  gGetAnnotation b b2 = case b2 of
    GenerateFieldAnn -> (fst $ gGetAnnotation @x b b2, ctorNameToAnn @a)
    NotGenerateFieldAnn -> (fst $ gGetAnnotation @x b b2, noAnn)

instance (GHasAnnotation x, KnownSymbol a) => GHasAnnotation (G.M1 G.C ('G.MetaCons a _p _f) x) where
  gGetAnnotation b b2 = (fst $ gGetAnnotation @x b b2, ctorNameToAnn @a)

instance (GHasAnnotation x) => GHasAnnotation (G.M1 G.D i1 x) where
  gGetAnnotation b b2 = gGetAnnotation @x b b2

instance
    ( GHasAnnotation x
    , GHasAnnotation y
    )
  =>
    GHasAnnotation (x G.:+: y)
  where

  gGetAnnotation followEntrypointFlag generateAnnFlag =
    let (xTypeAnn, xFieldAnn) = gGetAnnotation @x followEntrypointFlag generateAnnFlag
        (yTypeAnn, yFieldAnn) = gGetAnnotation @y followEntrypointFlag generateAnnFlag
    in case (followEntrypointFlag, generateAnnFlag) of
          (NotFollowEntrypoint, GenerateFieldAnn) ->
            ( NTOr noAnn
                xFieldAnn yFieldAnn
                xTypeAnn yTypeAnn
            , noAnn
            )
          _ ->
            ( NTOr noAnn
                noAnn noAnn
                xTypeAnn yTypeAnn
            , noAnn
            )

instance
    ( GHasAnnotation x
    , GHasAnnotation y
    )
  =>
    GHasAnnotation (x G.:*: y)
  where
  gGetAnnotation _ b2 =
    let  (xTypeAnn, xFieldAnn) = gGetAnnotation @x NotFollowEntrypoint b2
         (yTypeAnn, yFieldAnn) = gGetAnnotation @y NotFollowEntrypoint b2
    in
      ( NTPair noAnn
          xFieldAnn yFieldAnn
          xTypeAnn yTypeAnn
      , noAnn
      )

instance (HasAnnotation x) => GHasAnnotation (G.Rec0 x) where
  gGetAnnotation b _ = (getAnnotation @x b, noAnn)