-- | Common implementations of entrypoints.
module Lorentz.EntryPoints.Impl
  ( EpdPlain
  , EpdRecursive
  , EpdDelegate
  ) where

import qualified Data.Kind as Kind
import Data.Singletons (SingI (..))
import Data.Singletons.Prelude (Sing (STrue, SFalse))
import Data.Singletons.Prelude.Eq ((%==))
import Data.Type.Bool (If)
import qualified GHC.Generics as G
import Util.TypeLits
import Fcf (Exp, Eval)
import qualified Fcf

import Michelson.Typed
import Michelson.Typed.Haskell.Instr.Sum (IsPrimitiveValue)
import Michelson.Typed.Haskell.Value (GValueType, GenericIsoValue)
import Michelson.Untyped (FieldAnn, noAnn)
import Lorentz.Value
import Util.Type
import Util.Fcf (Over2, type (<|>), TyEqSing)

import Lorentz.EntryPoints.Helpers
import Lorentz.EntryPoints.Core
import Lorentz.TypeAnns

-- | Implementation of 'ParameterHasEntryPoints' which fits for case when
-- your contract exposes multiple entrypoints via having sum type as its
-- parameter.
--
-- In particular, each constructor would produce a homonymous entrypoint with
-- argument type equal to type of constructor field (each constructor should
-- have only one field).
-- Constructor called 'Default' will designate the default entrypoint.
data EpdPlain
instance PlainEntryPointsC EpdPlain cp => EntryPointsDerivation EpdPlain cp where
  type EpdAllEntryPoints EpdPlain cp = PlainAllEntryPointsExt EpdPlain cp
  type EpdLookupEntryPoint EpdPlain cp = PlainLookupEntryPointExt EpdPlain cp
  epdNotes = plainEpdNotesExt @EpdPlain @cp
  epdCall = plainEpdCallExt @EpdPlain @cp

-- | Extension of 'EpdPlain' on parameters 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.
--
-- It does not assign names to intermediate nodes of 'Or' tree, only to the very
-- leaves.
--
-- If some entrypoint arguments have custom 'IsoValue' instance, this
-- derivation way will not work. As a workaround, you can wrap your
-- argument into some primitive (e.g. ':!').
data EpdRecursive
instance PlainEntryPointsC EpdRecursive cp => EntryPointsDerivation EpdRecursive cp where
  type EpdAllEntryPoints EpdRecursive cp = PlainAllEntryPointsExt EpdRecursive cp
  type EpdLookupEntryPoint EpdRecursive cp = PlainLookupEntryPointExt EpdRecursive cp
  epdNotes = plainEpdNotesExt @EpdRecursive @cp
  epdCall = plainEpdCallExt @EpdRecursive @cp

-- | Extension of 'EpdPlain' on parameters being defined as several nested
-- datatypes.
--
-- In particular, it will traverse the immediate sum type, and require another
-- 'ParameterHasEntryPoints' for the inner complex datatypes. Only those
-- inner types are considered which are the only fields in their respective
-- constructors.
-- Inner types should not themselves declare default entrypoint, we enforce
-- this for better modularity.
-- Each top-level constructor will be treated as entrypoint even if it contains
-- a complex datatype within, in such case that would be an entrypoint
-- corresponding to intermediate node in @or@ tree.
--
-- Comparing to 'EpdRecursive' this gives you more control over where and how
-- entrypoints will be derived.
data EpdDelegate
instance (PlainEntryPointsC EpdDelegate cp) => EntryPointsDerivation EpdDelegate cp where
  type EpdAllEntryPoints EpdDelegate cp = PlainAllEntryPointsExt EpdDelegate cp
  type EpdLookupEntryPoint EpdDelegate cp = PlainLookupEntryPointExt EpdDelegate cp
  epdNotes = plainEpdNotesExt @EpdDelegate @cp
  epdCall = plainEpdCallExt @EpdDelegate @cp

type PlainAllEntryPointsExt mode cp = AllEntryPoints mode (BuildEPTree mode cp) cp

type PlainLookupEntryPointExt mode cp = LookupEntryPoint mode (BuildEPTree mode cp) cp

plainEpdNotesExt
  :: forall mode cp.
     (PlainEntryPointsC mode cp, HasCallStack)
  => Notes (ToT cp)
plainEpdNotesExt = mkEntryPointsNotes @mode @(BuildEPTree mode cp) @cp

plainEpdCallExt
  :: forall mode cp name.
     (PlainEntryPointsC mode cp, ParameterScope (ToT cp), KnownSymbol name)
  => Label name
  -> EpConstructionRes (ToT cp) (Eval (LookupEntryPoint mode (BuildEPTree mode cp) cp name))
plainEpdCallExt = mkEpLiftSequence @mode @(BuildEPTree mode cp) @cp

type PlainEntryPointsC mode cp =
  ( GenericIsoValue cp
  , EntryPointsNotes mode (BuildEPTree mode cp) cp
  , RequireSumType cp
  )

-- | 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
    -- ^ We are in the intermediate node and need to go deeper.
  | EPLeaf
    -- ^ We reached entrypoint argument.
  | EPDelegate
    -- ^ We reached complex parameter part and will need to ask how to process it.

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

type family GBuildEntryPointsTree (mode :: Kind.Type) (x :: Kind.Type -> Kind.Type)
             :: EPTree where
  GBuildEntryPointsTree mode (G.D1 _ x) =
    GBuildEntryPointsTree mode x
  GBuildEntryPointsTree mode (x G.:+: y) =
    'EPNode (GBuildEntryPointsTree mode x) (GBuildEntryPointsTree mode y)

  GBuildEntryPointsTree EpdPlain (G.C1 _ _) =
    'EPLeaf
  GBuildEntryPointsTree EpdRecursive (G.C1 _ x) =
    GBuildEntryPointsTree EpdRecursive x
  GBuildEntryPointsTree EpdDelegate (G.C1 _ (G.S1 _ (G.Rec0 _))) =
    'EPDelegate
  GBuildEntryPointsTree EpdDelegate (G.C1 _ _) =
    'EPLeaf
  GBuildEntryPointsTree mode (G.S1 _ x) =
    GBuildEntryPointsTree mode x
  GBuildEntryPointsTree _ G.U1 =
    'EPLeaf
  GBuildEntryPointsTree _ (_ G.:*: _) =
    'EPLeaf
  GBuildEntryPointsTree mode (G.Rec0 a) =
    If (IsPrimitiveValue a)
       'EPLeaf
       (BuildEPTree mode a)

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

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

-- | Makes up a way to lift entrypoint argument to full parameter.
mkEpLiftSequence
  :: forall mode ep a name.
      ( EntryPointsNotes mode ep a, ParameterScope (ToT a)
      , GenericIsoValue a, KnownSymbol name
      )
  => Label name
  -> EpConstructionRes (ToT a) (Eval (LookupEntryPoint mode ep a name))
mkEpLiftSequence = gMkEpLiftSequence @mode @ep @(G.Rep a)

-- | Fetches information about all entrypoints - leaves of 'Or' tree.
type AllEntryPoints mode ep a = GAllEntryPoints mode ep (G.Rep a)

-- | Fetches information about all entrypoints - leaves of 'Or' tree.
type LookupEntryPoint mode ep a = GLookupEntryPoint mode ep (G.Rep a)

-- | Generic traversal for 'EntryPointsNotes'.
class GEntryPointsNotes (mode :: Kind.Type) (ep :: EPTree) (x :: Kind.Type -> Kind.Type) where
  type GAllEntryPoints mode ep x :: [(Symbol, Kind.Type)]
  type GLookupEntryPoint mode ep x :: Symbol -> Exp (Maybe Kind.Type)

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

  gMkEpLiftSequence
    :: (KnownSymbol name, ParameterScope (GValueType x))
    => Label name
    -> EpConstructionRes (GValueType x) (Eval (GLookupEntryPoint mode ep x name))

instance GEntryPointsNotes mode ep x => GEntryPointsNotes mode ep (G.D1 i x) where
  type GAllEntryPoints mode ep (G.D1 i x) = GAllEntryPoints mode ep x
  type GLookupEntryPoint mode ep (G.D1 i x) = GLookupEntryPoint mode ep x
  gMkEntryPointsNotes = gMkEntryPointsNotes @mode @ep @x
  gMkEpLiftSequence = gMkEpLiftSequence @mode @ep @x

instance (GEntryPointsNotes mode epx x, GEntryPointsNotes mode epy y) =>
         GEntryPointsNotes mode ('EPNode epx epy) (x G.:+: y) where
  type GAllEntryPoints mode ('EPNode epx epy) (x G.:+: y) =
    GAllEntryPoints mode epx x ++ GAllEntryPoints mode epy y
  type GLookupEntryPoint mode ('EPNode epx epy) (x G.:+: y) =
    Over2 (<|>) (GLookupEntryPoint mode epx x) (GLookupEntryPoint mode epy y)
  gMkEntryPointsNotes =
    let (xnotes, xann) = gMkEntryPointsNotes @mode @epx @x
        (ynotes, yann) = gMkEntryPointsNotes @mode @epy @y
    in (NTOr noAnn xann yann xnotes ynotes, noAnn)
  gMkEpLiftSequence label =
    case sing @(GValueType (x G.:+: y)) of
      STOr sl _ -> case (checkOpPresence sl, checkNestedBigMapsPresence sl) of
        (OpAbsent, NestedBigMapsAbsent) ->
          case gMkEpLiftSequence @mode @epx @x label of
            EpConstructed liftSeq -> EpConstructed (EplWrapLeft liftSeq)
            EpConstructionFailed ->
              case gMkEpLiftSequence @mode @epy @y label of
                EpConstructed liftSeq -> EpConstructed (EplWrapRight liftSeq)
                EpConstructionFailed -> EpConstructionFailed

instance ( GHasTypeAnn x, KnownSymbol ctor
         , ToT (GExtractField x) ~ GValueType x
         ) =>
         GEntryPointsNotes mode 'EPLeaf (G.C1 ('G.MetaCons ctor _1 _2) x) where
  type GAllEntryPoints mode 'EPLeaf (G.C1 ('G.MetaCons ctor _1 _2) x) =
    '[ '(ctor, GExtractField x) ]
  type GLookupEntryPoint mode 'EPLeaf (G.C1 ('G.MetaCons ctor _1 _2) x) =
    JustOnEq ctor (GExtractField x)
  gMkEntryPointsNotes =
    (gGetTypeAnn @x, ctorNameToAnn @ctor)
  gMkEpLiftSequence (_ :: Label name) =
    case sing @ctor %== sing @name of
      STrue -> EpConstructed EplArgHere
      SFalse -> EpConstructionFailed

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

instance ( ep ~ 'EPDelegate, GEntryPointsNotes mode ep x
         , KnownSymbol ctor, ToT (GExtractField x) ~ GValueType x
         ) =>
         GEntryPointsNotes mode 'EPDelegate (G.C1 ('G.MetaCons ctor _1 _2) x) where
  type GAllEntryPoints mode 'EPDelegate (G.C1 ('G.MetaCons ctor _1 _2) x) =
    '(ctor, GExtractField x) ': GAllEntryPoints mode 'EPDelegate x
  type GLookupEntryPoint mode 'EPDelegate (G.C1 ('G.MetaCons ctor _1 _2) x) =
    Over2 (<|>) (JustOnEq ctor (GExtractField x)) (GLookupEntryPoint mode 'EPDelegate x)
  gMkEntryPointsNotes =
    let (notes, _rootAnn) = gMkEntryPointsNotes @mode @ep @x
    in (notes, ctorNameToAnn @ctor)
  gMkEpLiftSequence (label :: Label name) =
    case sing @ctor %== sing @name of
      STrue -> EpConstructed EplArgHere
      SFalse -> gMkEpLiftSequence @mode @ep @x label

instance GEntryPointsNotes mode ep x => GEntryPointsNotes mode ep (G.S1 i x) where
  type GAllEntryPoints mode ep (G.S1 i x) = GAllEntryPoints mode ep x
  type GLookupEntryPoint mode ep (G.S1 i x) = GLookupEntryPoint mode ep x
  gMkEntryPointsNotes = gMkEntryPointsNotes @mode @ep @x
  gMkEpLiftSequence = gMkEpLiftSequence @mode @ep @x

instance (EntryPointsNotes EpdRecursive ep a, GenericIsoValue a) =>
         GEntryPointsNotes EpdRecursive ep (G.Rec0 a) where
  type GAllEntryPoints EpdRecursive ep (G.Rec0 a) = AllEntryPoints EpdRecursive ep a
  type GLookupEntryPoint EpdRecursive ep (G.Rec0 a) = LookupEntryPoint EpdRecursive ep a
  gMkEntryPointsNotes = (mkEntryPointsNotes @EpdRecursive @ep @a, noAnn)
  gMkEpLiftSequence = mkEpLiftSequence @EpdRecursive @ep @a

instance (ParameterDeclaresEntryPoints a) =>
         GEntryPointsNotes EpdDelegate 'EPDelegate (G.Rec0 a) where
  type GAllEntryPoints EpdDelegate 'EPDelegate (G.Rec0 a) = AllParameterEntryPoints a
  type GLookupEntryPoint EpdDelegate 'EPDelegate (G.Rec0 a) = LookupParameterEntryPoint a
  gMkEntryPointsNotes = (pepNotes @a, noAnn)
  -- TODO [#35]: should use field ann ^^^^^ returned by 'epdNotes'
  gMkEpLiftSequence = pepCall @a

instance GEntryPointsNotes mode 'EPLeaf G.U1 where
  type GAllEntryPoints mode 'EPLeaf G.U1 = '[]
  type GLookupEntryPoint mode 'EPLeaf G.U1 = Fcf.ConstFn 'Nothing
  gMkEntryPointsNotes = (starNotes, noAnn)
  gMkEpLiftSequence _ = EpConstructionFailed

instance Each [Typeable, SingI] [GValueType x, GValueType y] =>
         GEntryPointsNotes mode 'EPLeaf (x G.:*: y) where
  type GAllEntryPoints mode 'EPLeaf (x G.:*: y) = '[]
  type GLookupEntryPoint mode 'EPLeaf (x G.:*: y) = Fcf.ConstFn 'Nothing
  gMkEntryPointsNotes = (starNotes, noAnn)
  gMkEpLiftSequence _ = EpConstructionFailed

-- Return 'Just' iff given entries of type @k1@ are equal.
type family JustOnEq (a :: k1) (b :: k2) :: k1 -> Exp (Maybe k2) where
  JustOnEq a b =
    Fcf.Flip Fcf.Guarded
      '[ TyEqSing a 'Fcf.:= Fcf.Pure ('Just b)
       , Fcf.Otherwise 'Fcf.:= Fcf.Pure 'Nothing
       ]

-- Get field type under 'G.C1'.
type family GExtractField (x :: Kind.Type -> Kind.Type) where
  GExtractField (G.S1 _ x) = GExtractField x
  GExtractField (G.Rec0 a) = a
  GExtractField G.U1 = ()