{-# LANGUAGE UndecidableSuperClasses #-}

module Lorentz.EntryPoints
  ( ParameterEntryPoints (..)
  , ParameterEntryPointsSplit (..)
  , mapParameterEntryPoints
  , pepNone
  , pepPlain
  , pepRecursive
  ) where

import qualified Data.Kind as Kind
import Data.Singletons (SingI)
import Data.Type.Bool (If)
import qualified GHC.Generics as G
import Named (NamedF)

import Lorentz.Base ((:->))
import Lorentz.Constraints
import Lorentz.Zip
import Michelson.Text
import Michelson.Typed
  (BigMap, ContractRef(..), EpAddress, Notes(..), Operation, T(..), ToT, ToCT,
  starNotes)
import Michelson.Typed.Haskell.Instr.Sum (IsPrimitiveValue)
import Michelson.Typed.Haskell.Value (GValueType, IsGenericIsoValue)
import Michelson.Untyped (Annotation(..), FieldAnn, TypeAnn, ann, noAnn)
import Tezos.Address
import Tezos.Core
import Tezos.Crypto
import Util.Text
import Util.TypeLits

-- | Which entrypoints given parameter declares.
class NiceParameter p => ParameterEntryPoints p where
  parameterEntryPoints :: ParameterEntryPointsSplit p

-- | Implementation of 'parameterEntryPoints'.
newtype ParameterEntryPointsSplit p = ParameterEntryPointsSplit
  { pesNotes :: Notes (ToT p)
    -- ^ Parameter annotations which declare necessary entrypoints.
  }

mapParameterEntryPoints
  :: (ToT a ~ ToT b)
  => (a -> b) -> ParameterEntryPointsSplit a -> ParameterEntryPointsSplit b
mapParameterEntryPoints _ ParameterEntryPointsSplit{..} =
  ParameterEntryPointsSplit{..}

-- | No entrypoints declared, parameter type will serve as argument type
-- of the only existing entrypoint.
pepNone :: SingI (ToT p) => ParameterEntryPointsSplit p
pepNone = ParameterEntryPointsSplit starNotes

instance ParameterEntryPoints () where
  parameterEntryPoints = pepNone

-- Common implementations
----------------------------------------------------------------------------

-- | Fits for case when your contract exposes multiple entrypoints
-- via having sum type as its parameter.
--
-- In particular, this will attach field annotations to immediate parameter "arms"
-- which will be named as corresponding constructor names.
pepPlain :: PesEntryPointsC 'False cp st => ParameterEntryPointsSplit cp
pepPlain = pepEntryPointsExt @'False

-- | Similar to 'pesEntryPoints', but for case of parameter being defined as
-- several nested datatypes.
--
-- In particular, this will traverse sum types recursively, stopping at
-- Michelson primitives (like 'Natural') and constructors with number of
-- fields different from one.
pepRecursive :: PesEntryPointsC 'True cp st => ParameterEntryPointsSplit cp
pepRecursive = pepEntryPointsExt @'True

pepEntryPointsExt
  :: forall deep cp st.
     (PesEntryPointsC deep cp st)
  => ParameterEntryPointsSplit cp
pepEntryPointsExt = ParameterEntryPointsSplit $
  mkEntryPointsNotes @deep @(BuildEPTree deep cp) @cp

type PesEntryPointsC deep cp st =
  ( IsGenericIsoValue cp
  , EntryPointsNotes deep (BuildEPTree deep cp) cp
  , GRequireSumType (G.Rep cp)
  , HasCallStack
  )

-- | Turn constructor name into an appropriate field annotation.
ctorNameToAnn :: HasCallStack => Text -> FieldAnn
ctorNameToAnn = Annotation . headToLower

-- | Ensure that given type is a sum type.
--
-- This helps to prevent attempts to apply a function to, for instance, a pair.
type family GRequireSumType (x :: Kind.Type -> Kind.Type) :: Constraint where
  GRequireSumType (G.M1 _ _ x) = GRequireSumType x
  GRequireSumType (_ G.:+: _) = ()
  GRequireSumType _ = TypeError ('Text "Expected sum type")

-- | Entrypoints tree - skeleton on 'TOr' tree later used to distinguish
-- between constructors-entrypoints and constructors which consolidate
-- a whole pack of entrypoints.
data EPTree
  = EPNode EPTree EPTree
  | EPLeaf

-- | Build 'EPTree' by parameter type.
type BuildEPTree deep a = GBuildEntryPointsTree deep (G.Rep a)

type family GBuildEntryPointsTree (deep :: Bool) (x :: Kind.Type -> Kind.Type)
             :: EPTree where
  GBuildEntryPointsTree deep (G.D1 _ x) =
    GBuildEntryPointsTree deep x
  GBuildEntryPointsTree deep (x G.:+: y) =
    'EPNode (GBuildEntryPointsTree deep x) (GBuildEntryPointsTree deep y)
  GBuildEntryPointsTree 'False (G.C1 _ _) =
    'EPLeaf
  GBuildEntryPointsTree 'True (G.C1 _ x) =
    GBuildEntryPointsTree 'True x
  GBuildEntryPointsTree deep (G.S1 _ x) =
    GBuildEntryPointsTree deep x
  GBuildEntryPointsTree _ G.U1 =
    'EPLeaf
  GBuildEntryPointsTree _ (_ G.:*: _) =
    'EPLeaf
  GBuildEntryPointsTree deep (G.Rec0 a) =
    If (IsPrimitiveValue a)
       'EPLeaf
       (BuildEPTree deep a)

-- | Traverses sum type and constructs 'Notes' which report
-- constructor names via field annotations.
type EntryPointsNotes deep ep a = (Generic a, GEntryPointsNotes deep ep (G.Rep a))

-- | Makes up notes with proper field annotations for given parameter.
mkEntryPointsNotes
  :: forall deep ep a.
      (EntryPointsNotes deep ep a, IsGenericIsoValue a, HasCallStack)
  => Notes (ToT a)
mkEntryPointsNotes = fst $ gMkEntryPointsNotes @deep @ep @(G.Rep a)

-- | Generic traversal for 'EntryPointsNotes'.
class GEntryPointsNotes (deep :: Bool) (ep :: EPTree) (x :: Kind.Type -> Kind.Type) where
  {- | Returns:
    1. Notes corresponding to this level;
    2. Field annotation for this level (and which should be used one level above).
  -}
  gMkEntryPointsNotes :: HasCallStack => (Notes (GValueType x), FieldAnn)

instance GEntryPointsNotes deep ep x => GEntryPointsNotes deep ep (G.D1 i x) where
  gMkEntryPointsNotes = gMkEntryPointsNotes @deep @ep @x

instance
    ( GEntryPointsNotes deep epx x
    , GEntryPointsNotes deep epy y
    )
  =>
    GEntryPointsNotes deep ('EPNode epx epy) (x G.:+: y)
  where
  gMkEntryPointsNotes =
    let (xnotes, xann) = gMkEntryPointsNotes @deep @epx @x
        (ynotes, yann) = gMkEntryPointsNotes @deep @epy @y
    in (NTOr noAnn xann yann xnotes ynotes, noAnn)

instance (GHasTypeAnn x, KnownSymbol ctor) =>
         GEntryPointsNotes 'False ep (G.C1 ('G.MetaCons ctor _1 _2) x) where
  gMkEntryPointsNotes =
    (gGetTypeAnn @x, ctorNameToAnn (symbolValT' @ctor))

instance (GHasTypeAnn x, KnownSymbol ctor) =>
         GEntryPointsNotes 'True 'EPLeaf (G.C1 ('G.MetaCons ctor _1 _2) x) where
  gMkEntryPointsNotes =
    (gGetTypeAnn @x, ctorNameToAnn (symbolValT' @ctor))

instance (ep ~ 'EPNode epx epy, GEntryPointsNotes 'True ep x) =>
         GEntryPointsNotes 'True ('EPNode epx epy) (G.C1 ('G.MetaCons ctor _1 _2) x) where
  gMkEntryPointsNotes = gMkEntryPointsNotes @'True @ep @x

instance GEntryPointsNotes deep ep x => GEntryPointsNotes deep ep (G.S1 i x) where
  gMkEntryPointsNotes = gMkEntryPointsNotes @deep @ep @x

instance (EntryPointsNotes deep ep a, IsGenericIsoValue a) =>
         GEntryPointsNotes deep ep (G.Rec0 a) where
  gMkEntryPointsNotes = (mkEntryPointsNotes @deep @ep @a, noAnn)

instance GEntryPointsNotes deep 'EPLeaf G.U1 where
  gMkEntryPointsNotes = (starNotes, noAnn)

instance
    ( Typeable (GValueType x)
    , Typeable (GValueType y)
    , SingI (GValueType x)
    , SingI (GValueType y)
    )
  =>
    GEntryPointsNotes deep 'EPLeaf (x G.:*: y)
  where
  gMkEntryPointsNotes = (starNotes, noAnn)

-- 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)

instance {-# OVERLAPPABLE #-} (GHasTypeAnn (G.Rep a), GValueType (G.Rep a) ~ ToT a)
  => HasTypeAnn a where
  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 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 (ZippedStack i)
    , HasTypeAnn (ZippedStack o)
    )
  =>
    HasTypeAnn (i :-> o)
  where
  getTypeAnn = NTLambda noAnn
    (getTypeAnn @(ZippedStack i))
    (getTypeAnn @(ZippedStack o))

-- 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