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

-- | 'Expr' data type and its generalizations

module Indigo.Internal.Expr.Types
  ( -- * The Expr data type
    Expr (..)

  -- * Generalizations of Expr
  , IsExpr
  , ToExpr
  , ExprType
  , (:~>)
  , toExpr

  -- * Arithmetic Expr
  , IsArithExpr
  , IsUnaryArithExpr

  -- * Polymorphic Expr
  , IsConcatExpr
  , IsConcatListExpr
  , IsDivExpr
  , IsModExpr
  , IsGetExpr
  , IsMemExpr
  , IsSizeExpr
  , IsSliceExpr
  , IsUpdExpr

  , ObjectManipulation (..)
  , ObjectExpr
  , NamedFieldExpr (..)
  ) where

import Prelude hiding (Const (..))
import qualified Data.Kind as Kind
import Data.Vinyl.Core (RMap(..), RecordToList(..))
import Data.Vinyl.Functor (Const (..))
import Data.Typeable (typeRep)
import Fmt (Buildable (..), pretty)

import Indigo.Lorentz
import Indigo.Internal.Expr.TH (deriveExprBuildable)
import Indigo.Internal.Field
import Indigo.Internal.Object (ComplexObjectC, FieldTypes, IndigoObjectF (..))
import Indigo.Internal.Var (Var (..))
import qualified Michelson.Typed.Arith as M
import Michelson.Typed (GetFieldType, CtorOnlyField, InstrUnwrapC, InstrWrapOneC)

----------------------------------------------------------------------------
-- The Expr data type
----------------------------------------------------------------------------

data Expr a where
  C :: NiceConstant a => a -> Expr a

  V :: KnownValue a => Var a -> Expr a

  ObjMan :: ObjectManipulation a -> Expr a

  Cast :: KnownValue a => Expr a -> Expr a

  Size :: SizeOpHs c => Expr c -> Expr Natural

  Update
    :: (UpdOpHs c, KnownValue c)
    => Expr c -> Expr (UpdOpKeyHs c) -> Expr (UpdOpParamsHs c) -> Expr c

  Add
    :: (ArithOpHs M.Add n m, KnownValue (ArithResHs M.Add n m))
    => Expr n -> Expr m -> Expr (ArithResHs M.Add n m)

  Sub
    :: (ArithOpHs M.Sub n m, KnownValue (ArithResHs M.Sub n m))
    => Expr n -> Expr m -> Expr (ArithResHs M.Sub n m)

  Mul
    :: (ArithOpHs M.Mul n m, KnownValue (ArithResHs M.Mul n m))
    => Expr n -> Expr m -> Expr (ArithResHs M.Mul n m)

  Div
    :: (EDivOpHs n m, KnownValue (EDivOpResHs n m))
    => Expr n -> Expr m -> Expr (EDivOpResHs n m)

  Mod
    :: (EDivOpHs n m, KnownValue (EModOpResHs n m))
    => Expr n -> Expr m -> Expr (EModOpResHs n m)

  Abs
    :: (UnaryArithOpHs M.Abs n, KnownValue (UnaryArithResHs M.Abs n))
    => Expr n -> Expr (UnaryArithResHs M.Abs n)

  Neg
    :: (UnaryArithOpHs M.Neg n, KnownValue (UnaryArithResHs M.Neg n))
    => Expr n -> Expr (UnaryArithResHs M.Neg n)


  Lsl
    :: (ArithOpHs M.Lsl n m, KnownValue (ArithResHs M.Lsl n m))
    => Expr n -> Expr m -> Expr (ArithResHs M.Lsl n m)

  Lsr
    :: (ArithOpHs M.Lsr n m, KnownValue (ArithResHs M.Lsr n m))
    => Expr n -> Expr m -> Expr (ArithResHs M.Lsr n m)


  Eq' :: NiceComparable n => Expr n -> Expr n -> Expr Bool

  Neq :: NiceComparable n => Expr n -> Expr n -> Expr Bool

  Le :: NiceComparable n => Expr n -> Expr n -> Expr Bool

  Lt :: NiceComparable n => Expr n -> Expr n -> Expr Bool

  Ge :: NiceComparable n => Expr n -> Expr n -> Expr Bool

  Gt :: NiceComparable n => Expr n -> Expr n -> Expr Bool

  Or
    :: (ArithOpHs M.Or n m, KnownValue (ArithResHs M.Or n m))
    => Expr n -> Expr m -> Expr (ArithResHs M.Or n m)

  Xor
    :: (ArithOpHs M.Xor n m, KnownValue (ArithResHs M.Xor n m))
    => Expr n -> Expr m -> Expr (ArithResHs M.Xor n m)

  And
    :: (ArithOpHs M.And n m, KnownValue (ArithResHs M.And n m))
    => Expr n -> Expr m -> Expr (ArithResHs M.And n m)

  Not
    :: (UnaryArithOpHs M.Not n, KnownValue (UnaryArithResHs M.Not n))
    => Expr n -> Expr (UnaryArithResHs M.Not n)

  Int' :: Expr Natural -> Expr Integer

  IsNat :: Expr Integer -> Expr (Maybe Natural)

  Coerce
    :: (Castable_ a b, KnownValue b)
    => Expr a -> Expr b

  ForcedCoerce
    :: (MichelsonCoercible a b, KnownValue b)
    => Expr a -> Expr b

  Fst :: KnownValue n => Expr (n, m) -> Expr n
  Snd :: KnownValue m => Expr (n, m) -> Expr m

  Pair :: KnownValue (n, m) => Expr n -> Expr m -> Expr (n, m)

  Some :: KnownValue (Maybe t) => Expr t -> Expr (Maybe t)
  None :: KnownValue t => Expr (Maybe t)

  Right' :: (KnownValue y, KnownValue (Either y x)) => Expr x -> Expr (Either y x)
  Left' :: (KnownValue x, KnownValue (Either y x)) => Expr y -> Expr (Either y x)

  Mem :: MemOpHs c => Expr (MemOpKeyHs c) -> Expr c -> Expr Bool

  StGet
    :: ( StoreHasSubmap store name key value
       , KnownValue value
       )
    => Label name -> Expr key -> Expr store -> Expr (Maybe value)

  StInsertNew
    :: ( StoreHasSubmap store name key value
       , KnownValue store
       , IsError err
       , Buildable err
       )
    => Label name -> err
    -> Expr key -> Expr value -> Expr store -> Expr store

  StInsert
    :: (StoreHasSubmap store name key value, KnownValue store)
    => Label name
    -> Expr key -> Expr value -> Expr store -> Expr store

  StMem
    :: ( StoreHasSubmap store name key val
       , KnownValue val
       )
    => Label name -> Expr key -> Expr store -> Expr Bool

  StUpdate
    :: (StoreHasSubmap store name key val, KnownValue store)
    => Label name -> Expr key -> Expr (Maybe val) -> Expr store -> Expr store

  StDelete
    :: (StoreHasSubmap store name key val, KnownValue store, KnownValue val)
    => Label name -> Expr key -> Expr store -> Expr store

  Wrap
    :: ( InstrWrapOneC dt name
       , KnownValue dt
       )
    => Label name
    -> Expr (CtorOnlyField name dt)
    -> Expr dt
  Unwrap
    :: ( InstrUnwrapC dt name
       , KnownValue (CtorOnlyField name dt)
       )
    => Label name
    -> Expr dt
    -> Expr (CtorOnlyField name dt)

  Construct
    :: ( InstrConstructC dt
       , RMap (ConstructorFieldTypes dt)
       , RecordToList (ConstructorFieldTypes dt)
       , KnownValue dt
       )
    => Proxy dt -> Rec Expr (ConstructorFieldTypes dt) -> Expr dt

  -- TODO remove Construct and rename this one
  ConstructWithoutNamed
    :: ComplexObjectC dt
    => Proxy dt -> Rec Expr (FieldTypes dt) -> Expr dt

  Name
    :: KnownValue (name :! t)
    => Label name -> Expr t -> Expr (name :! t)
  UnName
    :: KnownValue t
    => Label name -> Expr (name :! t) -> Expr t

  EmptySet
    :: (NiceComparable key, KnownValue (Set key))
    => Expr (Set key)

  Get
    :: ( GetOpHs c
       , KnownValue (Maybe (GetOpValHs c))
       , KnownValue (GetOpValHs c)
       )
    => Expr (GetOpKeyHs c) -> Expr c -> Expr (Maybe (GetOpValHs c))

  EmptyMap
    :: (KnownValue value, NiceComparable key, KnownValue (Map key value))
    => Expr (Map key value)

  EmptyBigMap
    :: (KnownValue value, NiceComparable key, KnownValue (BigMap key value))
    => Expr (BigMap key value)

  Pack
    :: NicePackedValue a
    => Expr a -> Expr (Packed a)

  Unpack
    :: NiceUnpackedValue a
    => Expr (Packed a) -> Expr (Maybe a)

  PackRaw
    :: NicePackedValue a
    => Expr a -> Expr ByteString

  UnpackRaw
    :: NiceUnpackedValue a
    => Expr ByteString -> Expr (Maybe a)

  Cons :: KnownValue (List a) => Expr a -> Expr (List a) -> Expr (List a)

  Nil :: KnownValue a => Expr (List a)

  Concat
    :: (ConcatOpHs c, KnownValue c)
    => Expr c -> Expr c -> Expr c

  Concat'
    :: (ConcatOpHs c, KnownValue c)
    => Expr (List c) -> Expr c

  Slice
    :: (SliceOpHs c, KnownValue c)
    => Expr Natural -> Expr Natural -> Expr c -> Expr (Maybe c)

  Contract
    :: ( NiceParameterFull p
       , NoExplicitDefaultEntrypoint p
       , ToTAddress p addr
       , ToT addr ~ ToT Address
       )
    => Expr addr -> Expr (Maybe (ContractRef p))

  Self
    :: ( NiceParameterFull p
       , NoExplicitDefaultEntrypoint p
       )
    => Expr (ContractRef p)

  ContractAddress
    :: Expr (ContractRef p) -> Expr Address

  ContractCallingUnsafe
    :: NiceParameter arg
    => EpName -> Expr Address -> Expr (Maybe (ContractRef arg))

  RunFutureContract
    :: NiceParameter p
    => Expr (FutureContract p) -> Expr (Maybe (ContractRef p))

  ImplicitAccount :: Expr KeyHash -> Expr (ContractRef ())

  ConvertEpAddressToContract
    :: NiceParameter p => Expr EpAddress -> Expr (Maybe (ContractRef p))

  MakeView
    :: KnownValue (View a r)
    => Expr a -> Expr (ContractRef r) -> Expr (View a r)

  MakeVoid
    :: KnownValue (Void_ a b)
    => Expr a -> Expr (Lambda b b) -> Expr (Void_ a b)

  CheckSignature
    :: BytesLike bs => Expr PublicKey -> Expr (TSignature bs) -> Expr bs -> Expr Bool

  Sha256 :: BytesLike bs => Expr bs -> Expr (Hash Sha256 bs)
  Sha512 :: BytesLike bs => Expr bs -> Expr (Hash Sha512 bs)
  Blake2b :: BytesLike bs => Expr bs -> Expr (Hash Blake2b bs)
  HashKey :: Expr PublicKey -> Expr KeyHash

  ChainId :: Expr ChainId

  Now :: Expr Timestamp
  Amount :: Expr Mutez
  Balance :: Expr Mutez
  Sender :: Expr Address

  Exec
    :: KnownValue b
    => Expr a -> Expr (Lambda a b) -> Expr b

  NonZero
    :: (NonZero n, KnownValue (Maybe n))
    => Expr n -> Expr (Maybe n)

----------------------------------------------------------------------------
-- Object manipulation
----------------------------------------------------------------------------

-- | Datatype describing access to an inner fields of object, like
-- @object !. field1 !. field2 ~. (field3, value3) ~. (field4, value4)@
data ObjectManipulation a where
  Object :: Expr a -> ObjectManipulation a

  ToField
    :: HasField dt fname ftype
    => ObjectManipulation dt
    -> Label fname
    -> ObjectManipulation ftype

  -- NB. @SetField (Object expr) field1
  --       (ObjMan $ SetField (ToField (Object expr) field1) field2 targetExpr)@
  -- is a bad representation, which will cause generation of not optimal code
  -- (like expr would be materialized object),
  -- so it would be nice to enforce only
  -- @SetField (Object expr) (field1 . field2) targetExpr@ representation.
  SetField
    :: HasField dt fname ftype
    => ObjectManipulation dt
    -> Label fname
    -> Expr ftype
    -> ObjectManipulation dt

instance Buildable (ObjectManipulation a) where
  build (Object e) = build e
  build (ToField oman lab) = pretty oman <> " #! " <> pretty lab
  build (SetField oman lab fEx) = pretty oman <> " !! (" <> pretty lab <> ", " <> pretty fEx <> ")"

-- | Auxiliary datatype where each field refers to
-- an expression the field equals to. It's not recursive one.
data NamedFieldExpr a name where
  NamedFieldExpr
    :: { unNamedFieldExpr :: Expr (GetFieldType a name) }
    -> NamedFieldExpr a name

type ObjectExpr a = IndigoObjectF (NamedFieldExpr a) a

----------------------------------------------------------------------------
-- Generalizations of Expr
----------------------------------------------------------------------------

type IsExpr op n = (ToExpr op, ExprType op ~ n, KnownValue n)
type (:~>) op n = IsExpr op n

type ExprType a = ExprType' (Decide a) a

toExpr :: forall a . ToExpr a => a -> Expr (ExprType a)
toExpr = toExpr' @(Decide a) @a

class ToExpr' (Decide x) x => ToExpr x
instance ToExpr' (Decide x) x => ToExpr x

-- This type class is needed to cope with overlapping instances.
class ToExpr' decision c where
  type family ExprType' decision c :: Kind.Type
  toExpr' :: c -> Expr (ExprType' decision c)

-- Instance for a var
instance KnownValue (a :: Kind.Type) => ToExpr' 'VarD (Var a) where
  type instance ExprType' 'VarD (Var a) = a
  toExpr' = V

-- Instance for a value
instance NiceConstant a => ToExpr' 'ValD a where
  type instance ExprType' 'ValD a = a
  toExpr' = C

-- Instance for StructManipulation
instance ToExpr' 'ObjManD (ObjectManipulation a) where
  type instance ExprType' 'ObjManD (ObjectManipulation a) = a
  toExpr' = ObjMan

-- Instance for Expr itself
instance ToExpr' 'ExprD (Expr a) where
  type instance ExprType' 'ExprD (Expr a) = a
  toExpr' = id

data Decision = VarD | ValD | ExprD | ObjManD

type family Decide x :: Decision where
  Decide (Var _) = 'VarD
  Decide (Expr _) = 'ExprD
  Decide (ObjectManipulation _) = 'ObjManD
  Decide _ = 'ValD

type IsUnaryArithExpr exN a n =
  ( exN :~> n
  , UnaryArithOpHs a n
  , KnownValue (UnaryArithResHs a n)
  )

type IsArithExpr exN exM a n m =
  ( exN :~> n, exM :~> m
  , ArithOpHs a n m
  , KnownValue (ArithResHs a n m)
  )

type IsDivExpr exN exM n m =
  ( exN :~> n, exM :~> m
  , EDivOpHs n m
  , KnownValue (EDivOpResHs n m)
  )

type IsModExpr exN exM n m =
  ( exN :~> n, exM :~> m
  , EDivOpHs n m
  , KnownValue (EModOpResHs n m)
  )

type IsConcatExpr exN1 exN2 n =
  ( exN1 :~> n
  , exN2 :~> n
  , ConcatOpHs n
  )

type IsConcatListExpr exN n =
  ( exN :~> List n
  , ConcatOpHs n
  , KnownValue n
  )

type IsSliceExpr exN n =
  ( exN :~> n
  , SliceOpHs n
  )

type IsGetExpr exKey exMap map =
  ( exKey :~> GetOpKeyHs map
  , exMap :~> map
  , GetOpHs map
  , KnownValue (GetOpValHs map)
  )

type IsUpdExpr exKey exVal exMap map =
  ( exKey :~> UpdOpKeyHs map
  , exVal :~> UpdOpParamsHs map
  , exMap :~> map
  , UpdOpHs map
  )

type IsMemExpr exKey exN n =
  ( exKey :~> MemOpKeyHs n
  , exN :~> n
  , MemOpHs n
  )

type IsSizeExpr exN n =
  ( exN :~> n
  , SizeOpHs n
  )

--data Expr2 a where
--  C2 :: (NiceConstant a, HasNoOp (ToT a)) => a -> Expr2 a
--
--instance Buildable (Expr2 a) where
--  build (C2 (x :: a)) = pretty (toVal x)

deriveExprBuildable ''Expr