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

{-# LANGUAGE FunctionalDependencies #-}

module Lorentz.ADT
  ( HasField
  , HasFieldOfType
  , HasFieldsOfType
  , NamedField (..)
  , (:=)
  , toField
  , toFieldNamed
  , getField
  , getFieldNamed
  , setField
  , modifyField
  , construct
  , constructT
  , constructStack
  , deconstruct
  , fieldCtor
  , wrap_
  , wrapOne
  , case_
  , caseT
  , unwrapUnsafe_
  , CaseTC
  , CaseArrow (..)
  , CaseClauseL (..)
  , InstrConstructC
  , ConstructorFieldTypes

    -- * Useful re-exports
  , Rec (..)
  , (:!)
  , (:?)
  , arg
  , argDef
  , argF
  ) where

import Data.Constraint ((\\))
import qualified Data.Kind as Kind
import Data.Vinyl.Core (RMap(..), Rec(..))
import GHC.TypeLits (AppendSymbol, Symbol)
import Named ((:!), (:?), arg, argDef, argF)

import Lorentz.Base
import Lorentz.Coercions
import Lorentz.Instr
import Michelson.Typed.Haskell.Instr
import Michelson.Typed.Haskell.Value
import Util.TypeTuple
import Util.Label (Label)
import Util.Type (type (++), KnownList)

-- | Allows field access and modification.
type HasField dt fname =
  ( InstrGetFieldC dt fname
  , InstrSetFieldC dt fname
  )

-- | Like 'HasField', but allows constrainting field type.
type HasFieldOfType dt fname fieldTy =
  ( HasField dt fname
  , GetFieldType dt fname ~ fieldTy
  )

-- | A pair of field name and type.
data NamedField = NamedField Symbol Kind.Type
type n := ty = 'NamedField n ty
infixr 0 :=

-- | Shortcut for multiple 'HasFieldOfType' constraints.
type family HasFieldsOfType (dt :: Kind.Type) (fs :: [NamedField])
             :: Constraint where
  HasFieldsOfType _ '[] = ()
  HasFieldsOfType dt ((n := ty) ': fs) =
    (HasFieldOfType dt n ty, HasFieldsOfType dt fs)

-- | Extract a field of a datatype replacing the value of this
-- datatype with the extracted field.
--
-- For this and the following functions you have to specify field name
-- which is either record name or name attached with @(:!)@ operator.
toField
  :: forall dt name st.
     InstrGetFieldC dt name
  => Label name -> dt & st :-> GetFieldType dt name & st
toField = I . instrGetField @dt

-- | Like 'toField', but leaves field named.
toFieldNamed
  :: forall dt name st.
     InstrGetFieldC dt name
  => Label name -> dt & st :-> (name :! GetFieldType dt name) & st
toFieldNamed l = toField l # forcedCoerce_

-- | Extract a field of a datatype, leaving the original datatype on stack.
getField
  :: forall dt name st.
     InstrGetFieldC dt name
  => Label name -> dt & st :-> GetFieldType dt name & dt ': st
getField l = dup # toField @dt l

-- | Like 'getField', but leaves field named.
getFieldNamed
  :: forall dt name st.
     InstrGetFieldC dt name
  => Label name -> dt & st :-> (name :! GetFieldType dt name) & dt ': st
getFieldNamed l = getField l # coerceWrap

-- | Set a field of a datatype.
setField
  :: forall dt name st.
     InstrSetFieldC dt name
  => Label name -> (GetFieldType dt name ': dt ': st) :-> (dt ': st)
setField = I . instrSetField @dt

-- | Apply given modifier to a datatype field.
modifyField
  :: forall dt name st.
     ( InstrGetFieldC dt name
     , InstrSetFieldC dt name
     )
  => Label name
  -> (forall st0. (GetFieldType dt name ': st0) :-> (GetFieldType dt name ': st0))
  -> dt & st :-> dt & st
modifyField l i = getField @dt l # i # setField @dt l

-- | Make up a datatype. You provide a pack of individual fields constructors.
--
-- Each element of the accepted record should be an instruction wrapped with
-- 'fieldCtor' function. This instruction will have access to the stack at
-- the moment of calling @construct@.
-- Instructions have to output fields of the built datatype, one per instruction;
-- instructions order is expected to correspond to the order of fields in the
-- datatype.
construct
  :: forall dt st.
     ( InstrConstructC dt
     , RMap (ConstructorFieldTypes dt)
     )
  => Rec (FieldConstructor st) (ConstructorFieldTypes dt)
  -> st :-> dt & st
construct fctors =
  I $ instrConstruct @dt $
  rmap (\(FieldConstructor i) -> FieldConstructor i) fctors

-- | Version of 'construct' which accepts tuple of field constructors.
constructT
  :: forall dt fctors st.
     ( InstrConstructC dt
     , RMap (ConstructorFieldTypes dt)
     , fctors ~ Rec (FieldConstructor st) (ConstructorFieldTypes dt)
     , RecFromTuple fctors
     )
  => IsoRecTuple fctors
  -> st :-> dt & st
constructT = construct . recFromTuple

-- | Construct an object from fields on the stack.
constructStack
  :: forall dt fields st .
  ( InstrConstructC dt
  , ToTs fields ~ ToTs (ConstructorFieldTypes dt)
  , KnownList fields
  )
  => (fields ++ st) :-> dt & st
constructStack =
  I (instrConstructStack @dt @(ToTs fields))
      \\ totsKnownLemma @fields
      \\ totsAppendLemma @fields @st

-- | Decompose a complex object into its fields
deconstruct
  :: forall dt fields st .
  ( InstrDeconstructC dt
  , KnownList fields
  , ToTs fields ~ ToTs (ConstructorFieldTypes dt)
  )
  => dt & st :-> (fields ++ st)
deconstruct =
  I (instrDeconstruct @dt @(ToTs fields))
    \\ totsKnownLemma @fields
    \\ totsAppendLemma @fields @st

-- | Lift an instruction to field constructor.
fieldCtor :: HasCallStack => (st :-> f & st) -> FieldConstructor st f
fieldCtor = \case
  I i -> FieldConstructor i
  FI _ -> error "Field constructor always fails"

-- | Wrap entry in constructor. Useful for sum types.
wrap_
  :: forall dt name st.
     InstrWrapC dt name
  => Label name -> (AppendCtorField (GetCtorField dt name) st) :-> dt & st
wrap_ =
  case appendCtorFieldAxiom @(GetCtorField dt name) @st of
    Dict -> I . instrWrap @dt

-- | Wrap entry in single-field constructor. Useful for sum types.
wrapOne
  :: forall dt name st.
     InstrWrapOneC dt name
  => Label name -> (CtorOnlyField name dt ': st) :-> dt & st
wrapOne = case appendCtorFieldAxiom @(GetCtorField dt name) @st of
  Dict -> I . instrWrapOne @dt

-- | Lorentz analogy of 'CaseClause', it works on plain 'Kind.Type' types.
data CaseClauseL (inp :: [Kind.Type]) (out :: [Kind.Type]) (param :: CaseClauseParam) where
  CaseClauseL :: AppendCtorField x inp :-> out -> CaseClauseL inp out ('CaseClauseParam ctor x)

-- | Provides "case" arrow which works on different wrappers for clauses.
class CaseArrow name body clause | clause -> name, clause -> body where
  -- | Lift an instruction to case clause.
  --
  -- You should write out constructor name corresponding to the clause
  -- explicitly. Prefix constructor name with "c" letter, otherwise
  -- your label will not be recognized by Haskell parser.
  -- Passing constructor name can be circumvented but doing so is not recomended
  -- as mentioning contructor name improves readability and allows avoiding
  -- some mistakes.
  (/->) :: Label name -> body -> clause
  infixr 0 /->

instance ( name ~ ("c" `AppendSymbol` ctor)
         , body ~ (AppendCtorField x inp :-> out)
         ) => CaseArrow name body
                (CaseClauseL inp out ('CaseClauseParam ctor x)) where
  (/->) _ = CaseClauseL

-- | Pattern match on the given sum type.
--
-- You have to provide a 'Rec' containing case branches.
-- To construct a case branch use '/->' operator.
case_
  :: forall dt out inp.
     ( InstrCaseC dt
     , RMap (CaseClauses dt)
     )
  => Rec (CaseClauseL inp out) (CaseClauses dt) -> dt & inp :-> out
case_ = LorentzInstr . instrCase @dt . rmap coerceCaseClause
  where
    coerceCaseClause
      :: forall clauses.
         CaseClauseL inp out clauses -> CaseClause (ToTs inp) (ToTs out) clauses
    coerceCaseClause (CaseClauseL (LorentzInstr cc)) =
      CaseClause $ case Proxy @clauses of
        (_ :: Proxy ('CaseClauseParam ctor cc)) ->
          case appendCtorFieldAxiom @cc @inp of Dict -> cc

-- | Like 'case_', accepts a tuple of clauses, which may be more convenient.
--
-- If user is experiencing problems with wierd errors about tuples while using
-- this function, he should take look at `Util.TypeTuple.Instances` and ensure
-- that his tuple isn't bigger than generated instances, if so, he should probably
-- extend number of generated instances.
caseT
  :: forall dt out inp clauses.
     CaseTC dt out inp clauses
  => IsoRecTuple clauses -> dt & inp :-> out
caseT = case_ @dt . recFromTuple

type CaseTC dt out inp clauses =
  ( InstrCaseC dt
  , RMap (CaseClauses dt)
  , RecFromTuple clauses
  , clauses ~ Rec (CaseClauseL inp out) (CaseClauses dt)
  )

-- | Unwrap a constructor with the given name. Useful for sum types.
unwrapUnsafe_
  :: forall dt name st.
     InstrUnwrapC dt name
  => Label name -> dt & st :-> (CtorOnlyField name dt ': st)
unwrapUnsafe_ =
  case appendCtorFieldAxiom @(GetCtorField dt name) @st of
    Dict -> I . instrUnwrapUnsafe @dt