{-# LANGUAGE UndecidableSuperClasses #-}

-- | Type annotations for Lorentz types.
module Lorentz.TypeAnns
  ( HasTypeAnn (..)
  , GHasTypeAnn (..)
  ) where

import Data.Singletons (SingI)
import qualified GHC.Generics as G
import Named (NamedF)

import Michelson.Text
import Michelson.Typed
  (BigMap, ContractRef(..), EpAddress, Notes(..), Operation, T(..), ToCT, ToT, starNotes)
import Michelson.Typed.Haskell.Value (GValueType)
import Michelson.Untyped (TypeAnn, ann, noAnn)
import Tezos.Address
import Tezos.Core
import Tezos.Crypto
import Util.TypeLits

-- For supporting type annotations of entrypoint arguments.
--
--At the botton of this infra is the HasTypeAnn class, which defines the type
--annotations for a given type. Right now the type annotations can only come
--from names in a named field. That is, we are not deriving names from, for
--example, record field names.

class HasTypeAnn a where
  getTypeAnn :: Notes (ToT a)

  default getTypeAnn
    :: (GHasTypeAnn (G.Rep a), GValueType (G.Rep a) ~ ToT a) => Notes (ToT a)
  getTypeAnn = gGetTypeAnn @(G.Rep a)

instance (HasTypeAnn a, KnownSymbol name)
  => HasTypeAnn (NamedF Identity a name) where
  getTypeAnn = insertNote (symbolAnn @name) $ getTypeAnn @a
    where
      symbolAnn :: forall s. KnownSymbol s => TypeAnn
      symbolAnn = ann $ symbolValT' @s
      -- Insert the provided type annotation into the provided note.
      insertNote :: forall (b :: T). TypeAnn -> Notes b -> Notes b
      insertNote nt s = case s of
        NTc _ -> NTc 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

instance (HasTypeAnn (Maybe a), KnownSymbol name)
  => HasTypeAnn (NamedF Maybe a name) where
  getTypeAnn = getTypeAnn @(NamedF Identity (Maybe a) name)

-- Primitive instances
instance (HasTypeAnn a) => HasTypeAnn (Maybe a) where
  getTypeAnn = NTOption noAnn (getTypeAnn @a)

instance HasTypeAnn ()

instance HasTypeAnn Integer where
  getTypeAnn = starNotes

instance HasTypeAnn Natural where
  getTypeAnn = starNotes

instance HasTypeAnn MText where
  getTypeAnn = starNotes

instance HasTypeAnn Bool where
  getTypeAnn = starNotes

instance HasTypeAnn ByteString where
  getTypeAnn = starNotes

instance HasTypeAnn Mutez where
  getTypeAnn = starNotes

instance HasTypeAnn Address where
  getTypeAnn = starNotes

instance HasTypeAnn EpAddress where
  getTypeAnn = starNotes

instance HasTypeAnn KeyHash where
  getTypeAnn = starNotes

instance HasTypeAnn Timestamp where
  getTypeAnn = starNotes

instance HasTypeAnn PublicKey where
  getTypeAnn = starNotes

instance HasTypeAnn Signature where
  getTypeAnn = starNotes

instance (HasTypeAnn a) => HasTypeAnn (ContractRef a) where
  getTypeAnn = NTContract noAnn (getTypeAnn @a)

instance (HasTypeAnn v) => HasTypeAnn (Map k v) where
  getTypeAnn = NTMap noAnn noAnn (getTypeAnn @v)

instance (HasTypeAnn v) => HasTypeAnn (BigMap k v) where
  getTypeAnn = NTBigMap noAnn noAnn (getTypeAnn @v)

instance (SingI (ToCT v), Typeable (ToCT v)) => HasTypeAnn (Set v) where
  getTypeAnn = starNotes

instance (HasTypeAnn a) => HasTypeAnn [a] where
  getTypeAnn = NTList noAnn (getTypeAnn @a)

instance HasTypeAnn Operation where
  getTypeAnn = starNotes

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

-- A Generic HasTypeAnn implementation
class GHasTypeAnn a where
  gGetTypeAnn :: Notes (GValueType a)

instance GHasTypeAnn G.U1 where
  gGetTypeAnn = starNotes

instance (GHasTypeAnn x) => GHasTypeAnn (G.M1 i0 i1 x) where
  gGetTypeAnn = gGetTypeAnn @x

instance
    ( GHasTypeAnn x
    , GHasTypeAnn y
    )
  =>
    GHasTypeAnn (x G.:+: y)
  where
  gGetTypeAnn = NTOr noAnn noAnn noAnn
    (gGetTypeAnn @x)
    (gGetTypeAnn @y)

instance
    ( GHasTypeAnn x
    , GHasTypeAnn y
    )
  =>
    GHasTypeAnn (x G.:*: y)
  where
  gGetTypeAnn = NTPair noAnn noAnn noAnn
    (gGetTypeAnn @x)
    (gGetTypeAnn @y)

instance (HasTypeAnn x) => GHasTypeAnn (G.Rec0 x) where
  gGetTypeAnn = getTypeAnn @x