module Lorentz.Instr
  ( nop
  , drop
  , dropN
  , dup
  , swap
  , digPeano
  , dig
  , dug
  , push
  , some
  , none
  , unit
  , ifNone
  , pair
  , car
  , cdr
  , left
  , right
  , ifLeft
  , nil
  , cons
  , size
  , emptySet
  , emptyMap
  , emptyBigMap
  , map
  , iter
  , mem
  , get
  , update
  , failingWhenPresent
  , updateNew
  , if_
  , ifCons
  , loop
  , loopLeft
  , lambda
  , exec
  , execute
  , apply
  , dip
  , ConstraintDIPNLorentz
  , dipNPeano
  , dipN
  , failWith
  , cast
  , pack
  , unpack
  , concat
  , concat'
  , slice, isNat, add, sub, rsub, mul, ediv, abs
  , neg
  , lsl
  , lsr
  , or
  , and
  , xor
  , not
  , compare
  , eq0
  , neq0
  , lt0
  , gt0
  , le0
  , ge0
  , int
  , toTAddress_
  , self
  , selfCalling
  , contract
  , contractCalling
  , contractCallingUnsafe
  , runFutureContract
  , epAddressToContract
  , transferTokens
  , setDelegate
  , createContract
  , implicitAccount
  , now
  , amount
  , balance
  , checkSignature
  , sha256
  , sha512
  , blake2B
  , hashKey
  , stepsToQuota
  , source
  , sender
  , address
  , chainId
  , framed
  , LorentzFunctor (..)
  , nonZero
  ) where

import Prelude hiding
  (EQ, GT, LT, abs, and, compare, concat, drop, get, map, not, or, some, swap, xor)

import Data.Constraint (Dict(..), (\\))
import qualified Data.Kind as Kind
import Data.Singletons (SingI, sing)
import qualified GHC.TypeNats as GHC (Nat)

import Lorentz.Arith
import Lorentz.Base
import Lorentz.Constraints
import Lorentz.EntryPoints
import Lorentz.Polymorphic
import Lorentz.Run (compileLorentzContract)
import Lorentz.Value
import Lorentz.Zip
import Michelson.Typed
  (pattern CAR, pattern CDR, ConstraintDIG, ConstraintDIG', ConstraintDIPN, ConstraintDIPN',
  ConstraintDUG, ConstraintDUG', pattern DefEpName, EntryPointCallT(..), Instr(..), RemFail(..),
  SomeEntryPointCallT(..), ToTs, Value'(..), sepcName, starNotes)
import Michelson.Typed.Arith
import Michelson.Typed.Haskell.Value
import Util.Peano
import Util.Type

nop :: s :-> s
nop = I Nop

drop :: a & s :-> s
drop = I DROP

-- | Drop top @n@ elements from the stack.
dropN ::
  forall (n :: GHC.Nat) (s :: [Kind.Type]).
  -- Note: if we introduce `nPeano ~ ToPeano n` variable,
  -- GHC will complain that this constraint is redundant.
  ( SingI (ToPeano n), KnownPeano (ToPeano n)
  , RequireLongerOrSameLength (ToTs s) (ToPeano n)
  -- ↓ Kinda obvious, but not to GHC.
  , Drop (ToPeano n) (ToTs s) ~ ToTs (Drop (ToPeano n) s)
  ) => s :-> Drop (ToPeano n) s
dropN = I (DROPN $ sing @(ToPeano n))
  where
    _example :: '[ Integer, Integer, Integer ] :-> '[]
    _example = dropN @3

dup  :: a & s :-> a & a & s
dup = I DUP

swap :: a & b & s :-> b & a & s
swap = I SWAP

-- See a comment about `ConstraintDIPNLorentz'.
type ConstraintDIGLorentz (n :: Peano) (inp :: [Kind.Type]) (out :: [Kind.Type])
  (a :: Kind.Type) =
  ( ConstraintDIG n (ToTs inp) (ToTs out) (ToT a)
  , ConstraintDIG' Kind.Type n inp out a
  )

type ConstraintDUGLorentz (n :: Peano) (inp :: [Kind.Type]) (out :: [Kind.Type])
  (a :: Kind.Type) =
  ( ConstraintDUG n (ToTs inp) (ToTs out) (ToT a)
  , ConstraintDUG' Kind.Type n inp out a
  )

-- | Version of `dig` which uses Peano number.
-- It is inteded for internal usage in Lorentz.
digPeano ::
  forall (n :: Peano) inp out a.
  ( ConstraintDIGLorentz n inp out a
  ) => inp :-> out
digPeano = I (DIG $ sing @n)

dig ::
  forall (n :: GHC.Nat) inp out a.
  ( ConstraintDIGLorentz (ToPeano n) inp out a
  ) => inp :-> out
dig = digPeano @(ToPeano n)
  where
    _example ::
      '[ Integer, Integer, Integer, Bool ] :->
      '[ Bool, Integer, Integer, Integer ]
    _example = dig @3

dug ::
  forall (n :: GHC.Nat) inp out a.
  ( ConstraintDUGLorentz (ToPeano n) inp out a
  ) => inp :-> out
dug = I (DUG $ sing @(ToPeano n))
  where
    _example ::
      '[ Bool, Integer, Integer, Integer ] :->
      '[ Integer, Integer, Integer, Bool ]
    _example = dug @3

push :: forall t s . NiceConstant t => t -> (s :-> t & s)
push a = I $ PUSH (toVal a) \\ niceConstantEvi @t

some :: a & s :-> Maybe a & s
some = I SOME

none :: forall a s . KnownValue a => s :-> (Maybe a & s)
none = I NONE

unit :: s :-> () & s
unit = I UNIT

ifNone
  :: (s :-> s') -> (a & s :-> s') -> (Maybe a & s :-> s')
ifNone = iGenericIf IF_NONE

pair :: a & b & s :-> (a, b) & s
pair = I PAIR

car :: (a, b) & s :-> a & s
car = I CAR

cdr :: (a, b) & s :-> b & s
cdr = I CDR

left :: forall a b s. KnownValue b => a & s :-> Either a b & s
left = I LEFT

right :: forall a b s. KnownValue a => b & s :-> Either a b & s
right = I RIGHT

ifLeft
  :: (a & s :-> s') -> (b & s :-> s') -> (Either a b & s :-> s')
ifLeft = iGenericIf IF_LEFT

nil :: KnownValue p => s :-> List p & s
nil = I NIL

cons :: a & List a & s :-> List a & s
cons = I CONS

ifCons
  :: (a & List a & s :-> s') -> (s :-> s') -> (List a & s :-> s')
ifCons = iGenericIf IF_CONS

size :: SizeOpHs c => c & s :-> Natural & s
size = I SIZE

emptySet :: (KnownCValue e) => s :-> Set e & s
emptySet = I EMPTY_SET

emptyMap :: (KnownCValue k, KnownValue v)
         => s :-> Map k v & s
emptyMap = I EMPTY_MAP

emptyBigMap :: (KnownCValue k, KnownValue v)
            => s :-> BigMap k v & s
emptyBigMap = I EMPTY_BIG_MAP

map
  :: (MapOpHs c, IsoMapOpRes c b, HasCallStack)
  => (MapOpInpHs c & s :-> b & s) -> (c & s :-> MapOpResHs c b & s)
map (iNonFailingCode -> action) = I (MAP action)

iter
  :: (IterOpHs c, HasCallStack)
  => (IterOpElHs c & s :-> s) -> (c & s :-> s)
iter (iNonFailingCode -> action) = I (ITER action)

mem :: MemOpHs c => MemOpKeyHs c & c & s :-> Bool & s
mem = I MEM

get :: GetOpHs c => GetOpKeyHs c & c & s :-> Maybe (GetOpValHs c) & s
get = I GET

update :: UpdOpHs c => UpdOpKeyHs c & UpdOpParamsHs c & c & s :-> c & s
update = I UPDATE

if_ :: (s :-> s') -> (s :-> s') -> (Bool & s :-> s')
if_ = iGenericIf IF

loop :: (s :-> Bool & s) -> (Bool & s :-> s)
loop (iAnyCode -> b) = I (LOOP b)

loopLeft
  :: (a & s :-> Either a b & s) -> (Either a b & s :-> b & s)
loopLeft (iAnyCode -> b) = I (LOOP_LEFT b)

lambda
  :: (ZipInstrs [i, o], KnownValue (ZippedStack i), KnownValue (ZippedStack o))
  => (i :-> o) -> (s :-> (i :-> o) & s)
lambda instr = case zippingStack instr of
  I l -> I (LAMBDA . VLam $ RfNormal l)
  FI l -> I (LAMBDA . VLam $ RfAlwaysFails l)

exec :: a & Lambda a b & s :-> b & s
exec = I EXEC

-- | Similar to 'exec' but works for lambdas with arbitrary size of input
-- and output.
--
-- Note that this instruction has its arguments flipped, lambda goes first.
-- This seems to be the only reasonable way to achieve good inference.
execute
  :: forall i o s.
      (Each [KnownList, ZipInstr] [i, o])
  => ((i :-> o) : (i ++ s)) :-> (o ++ s)
execute = framed @s $
  dip (zipInstr @i) # swap # I EXEC # unzipInstr @o
  where
    _example
      :: ([Integer, Natural] :-> [(), ()]) : Integer : Natural : s
      :-> () : () : s
    _example = execute

apply
  :: forall a b c s.
     (NiceConstant a)
  => a & Lambda (a, b) c & s :-> Lambda b c & s
apply = I $ APPLY \\ niceConstantEvi @a

dip :: forall a s s'. HasCallStack => (s :-> s') -> (a & s :-> a & s')
dip (iNonFailingCode -> a) = I (DIP a)

-- Helper constraint we need for 'dipN'.
-- The first constraint here is sufficient to make 'dipN' compile.
-- However, it is not enough for good type inference. If we use only the first
-- constraint, '_example' below will not compile because GHC will not be able
-- to deduce type of the input stack for 'unit'.
-- It can only deduce that 'ToTs s0' is empty (where 's0' is input stack), but
-- 'ToTs' is not injective, hence 's0' is ambiguous.
-- So we need both and we merge them into one to avoid a warning about
-- a redundant constraint.
type ConstraintDIPNLorentz (n :: Peano) (inp :: [Kind.Type]) (out :: [Kind.Type])
  (s :: [Kind.Type]) (s' :: [Kind.Type]) =
  ( ConstraintDIPN n (ToTs inp) (ToTs out) (ToTs s) (ToTs s')
  , ConstraintDIPN' Kind.Type n inp out s s'
  )

-- | Version of `dipN` which uses Peano number.
-- It is inteded for internal usage in Lorentz.
dipNPeano ::
  forall (n :: Peano) (inp :: [Kind.Type]) (out :: [Kind.Type]) (s :: [Kind.Type]) (s' :: [Kind.Type]).
  ( ConstraintDIPNLorentz n inp out s s'
  ) => s :-> s' -> inp :-> out
dipNPeano (iNonFailingCode -> a) = I (DIPN (sing @n) a)

dipN ::
  forall (n :: GHC.Nat) (inp :: [Kind.Type]) (out :: [Kind.Type]) (s :: [Kind.Type]) (s' :: [Kind.Type]).
  ( ConstraintDIPNLorentz (ToPeano n) inp out s s'
  ) => s :-> s' -> inp :-> out
dipN = dipNPeano @(ToPeano n)
  where
    _example :: '[ Integer, Integer, Integer ] :-> '[ Integer, Integer, Integer, () ]
    _example = dipN @3 unit

failWith :: (KnownValue a) => a & s :-> t
failWith = FI FAILWITH

cast :: KnownValue a => (a & s :-> a & s)
cast = I CAST

pack
  :: forall a s. (NicePackedValue a)
  => a & s :-> ByteString & s
pack = I $ PACK \\ nicePackedValueEvi @a

unpack
  :: forall a s. (NiceUnpackedValue a)
  => ByteString & s :-> Maybe a & s
unpack = I $ UNPACK \\ niceUnpackedValueEvi @a

concat :: ConcatOpHs c => c & c & s :-> c & s
concat = I CONCAT

concat' :: ConcatOpHs c => List c & s :-> c & s
concat' = I CONCAT'

slice :: SliceOpHs c => Natural & Natural & c & s :-> Maybe c & s
slice = I SLICE

isNat :: Integer & s :-> Maybe Natural & s
isNat = I ISNAT

add
  :: ArithOpHs Add n m
  => n & m & s :-> ArithResHs Add n m & s
add = I ADD

sub
  :: ArithOpHs Sub n m
  => n & m & s :-> ArithResHs Sub n m & s
sub = I SUB

rsub
  :: ArithOpHs Sub n m
  => m & n & s :-> ArithResHs Sub n m & s
rsub = swap # sub

mul
  :: ArithOpHs Mul n m
  => n & m & s :-> ArithResHs Mul n m & s
mul = I MUL

ediv :: EDivOpHs n m
     => n & m & s
     :-> Maybe ((EDivOpResHs n m, EModOpResHs n m)) & s
ediv = I EDIV

abs :: UnaryArithOpHs Abs n => n & s :-> UnaryArithResHs Abs n & s
abs = I ABS

neg :: UnaryArithOpHs Neg n => n & s :-> UnaryArithResHs Neg n & s
neg = I NEG


lsl
  :: ArithOpHs Lsl n m
  => n & m & s :-> ArithResHs Lsl n m & s
lsl = I LSL

lsr
  :: ArithOpHs Lsr n m
  => n & m & s :-> ArithResHs Lsr n m & s
lsr = I LSR

or
  :: ArithOpHs Or n m
  => n & m & s :-> ArithResHs Or n m & s
or = I OR

and
  :: ArithOpHs And n m
  => n & m & s :-> ArithResHs And n m & s
and = I AND

xor
  :: (ArithOpHs Xor n m)
  => n & m & s :-> ArithResHs Xor n m & s
xor = I XOR

not :: UnaryArithOpHs Not n => n & s :-> UnaryArithResHs Not n & s
not = I NOT

compare :: NiceComparable n
        => n & n & s :-> Integer & s
compare = I COMPARE

eq0 :: UnaryArithOpHs Eq' n => n & s :-> UnaryArithResHs Eq' n & s
eq0 = I EQ

neq0 :: UnaryArithOpHs Neq n => n & s :-> UnaryArithResHs Neq n & s
neq0 = I NEQ

lt0 :: UnaryArithOpHs Lt n => n & s :-> UnaryArithResHs Lt n & s
lt0 = I LT

gt0 :: UnaryArithOpHs Gt n => n & s :-> UnaryArithResHs Gt n & s
gt0 = I GT

le0 :: UnaryArithOpHs Le n => n & s :-> UnaryArithResHs Le n & s
le0 = I LE

ge0 :: UnaryArithOpHs Ge n => n & s :-> UnaryArithResHs Ge n & s
ge0 = I GE

int :: Natural & s :-> Integer & s
int = I INT

-- | Cast something appropriate to 'TAddress'.
-- TODO [TM-280]: try to move somewhere
toTAddress_
  :: (ToTAddress_ cp addr)
  => addr : s :-> TAddress cp : s
toTAddress_ = I Nop

-- | Something coercible to 'TAddress cp'.
type ToTAddress_ cp addr = (ToTAddress cp addr, ToT addr ~ ToT Address)

-- | Get a reference to the current contract.
--
-- Note that, similar to 'CONTRACT' instruction, in Michelson 'SELF' instruction
-- can accept an entrypoint as field annotation, and without annotation specified
-- it creates a @contract@ value which calls the default entrypoint.
--
-- This particular function carries the behaviour of @SELF@ before introduction
-- of lightweight entrypoints feature.
-- Thus the contract must __not__ have explicit "default" entrypoint for this to
-- work.
--
-- If you are going to call a specific entrypoint of the contract, see 'selfCalling'.
self
  :: forall p s.
      (NiceParameterFull p, ForbidExplicitDefaultEntryPoint p)
  => s :-> ContractRef p & s
self = I (SELF $ sepcCallRootChecked @p) \\ niceParameterEvi @p

-- | Make a reference to the current contract, maybe a specific entrypoint.
--
-- Note that, since information about parameter of the current contract is not
-- carried around, in this function you need to specify parameter type @p@
-- explicitly.
selfCalling
  :: forall p mname s.
     (NiceParameterFull p)
  => EntryPointRef mname
  -> s :-> ContractRef (GetEntryPointArgCustom p mname) & s
selfCalling epRef = I $
  withDict (niceParameterEvi @p) $
  case parameterEntryPointCallCustom @p epRef of
    epc@EntryPointCall{} -> SELF (SomeEpc epc)

-- | Get a reference to a contract by its address.
--
-- This instruction carries the behaviour of @CONTRACT@ before introduction
-- of lightweight entrypoints feature.
-- The contract must __not__ have explicit "default" entrypoint for this to work.
--
-- If you are going to call a specific entrypoint of the contract, see 'contractCalling'.
contract
  :: forall p addr s.
      ( NiceParameterFull p, ForbidExplicitDefaultEntryPoint p
      , ToTAddress_ p addr
      )
  => addr & s :-> Maybe (ContractRef p) & s
contract = I (CONTRACT starNotes epName) \\ niceParameterEvi @p
  where
    epName = sepcName (sepcCallRootChecked @p)

-- | Make a reference to a contract, maybe a specific entrypoint.
--
-- When calling this function, make sure that parameter type is known.
-- It's recommended that you supply 'TAddress' with a concrete parameter as the
-- stack argument.
contractCalling
  :: forall cp epRef epArg addr s.
     (HasEntryPointArg cp epRef epArg, ToTAddress_ cp addr)
  => epRef
  -> addr & s :-> Maybe (ContractRef epArg) & s
contractCalling epRef = I $
  case useHasEntryPointArg @cp @epRef @epArg epRef of
    (Dict, epName) -> CONTRACT starNotes epName

-- | Specialized version of 'contractCalling' for the case when you do
-- not have compile-time evidence of appropriate 'HasEntryPointArg'.
-- For instance, if you have untyped 'EpName' you can not have this
-- evidence (the value is only available in runtime).
-- If you have typed 'EntryPointRef', use 'eprName' to construct 'EpName'.
contractCallingUnsafe
  :: forall arg s.
     (NiceParameter arg)
  => EpName
  -> Address & s :-> Maybe (ContractRef arg) & s
contractCallingUnsafe epName = contractCalling (TrustEpName epName)

-- | Version of 'contract' instruction which may accept address with already
-- specified entrypoint name.
--
-- Also you cannot specify entrypoint name here because this could result in
-- conflict.
runFutureContract
  :: forall p s. (NiceParameter p)
  => FutureContract p & s :-> Maybe (ContractRef p) & s
runFutureContract =
  I Nop # epAddressToContract

-- | Similar to 'runFutureContract', works with 'EpAddress'.
--
-- Validity of such operation cannot be ensured at compile time.
epAddressToContract
  :: forall p s. (NiceParameter p)
  => EpAddress & s :-> Maybe (ContractRef p) & s
epAddressToContract =
  I (CONTRACT starNotes DefEpName) \\ niceParameterEvi @p

transferTokens
  :: forall p s. (NiceParameter p)
  => p & Mutez & ContractRef p & s :-> Operation & s
transferTokens = I $ TRANSFER_TOKENS \\ niceParameterEvi @p

setDelegate :: Maybe KeyHash & s :-> Operation & s
setDelegate = I SET_DELEGATE

createContract
  :: forall p g s. (NiceStorage g, NiceParameterFull p)
  => Contract p g
  -> Maybe KeyHash & Mutez & g & s
  :-> Operation & Address & s
createContract cntrc =
  I $ CREATE_CONTRACT (compileLorentzContract cntrc)
    \\ niceParameterEvi @p
    \\ niceStorageEvi @g

implicitAccount :: KeyHash & s :-> ContractRef () & s
implicitAccount = I IMPLICIT_ACCOUNT

now :: s :-> Timestamp & s
now = I NOW

amount :: s :-> Mutez & s
amount = I AMOUNT

balance :: s :-> Mutez & s
balance = I BALANCE

checkSignature :: PublicKey & Signature & ByteString & s :-> Bool & s
checkSignature = I CHECK_SIGNATURE

sha256 :: ByteString & s :-> ByteString & s
sha256 = I SHA256

sha512 :: ByteString & s :-> ByteString & s
sha512 = I SHA512

blake2B :: ByteString & s :-> ByteString & s
blake2B = I BLAKE2B

hashKey :: PublicKey & s :-> KeyHash & s
hashKey = I HASH_KEY

{-# WARNING stepsToQuota "STEPS_TO_QUOTA instruction is deprecated in Michelson 005" #-}
stepsToQuota :: s :-> Natural & s
stepsToQuota = I STEPS_TO_QUOTA

{-# WARNING source
    "Using `source` is considered a bad practice.\n\
\    Consider using `sender` instead until further investigation" #-}
source :: s :-> Address & s
source = I SOURCE

sender :: s :-> Address & s
sender = I SENDER

address :: ContractRef a & s :-> Address & s
address = I ADDRESS

chainId :: s :-> ChainId & s
chainId = I CHAIN_ID

-- | Execute given instruction on truncated stack.
--
-- This instruction requires you to specify the piece of stack to truncate
-- as type argument.
framed
  :: forall s i o.
      (KnownList i, KnownList o)
  => (i :-> o) -> ((i ++ s) :-> (o ++ s))
framed (iNonFailingCode -> i) =
  I $ FrameInstr (Proxy @(ToTs s)) i
    \\ totsKnownLemma @i
    \\ totsKnownLemma @o
    \\ totsAppendLemma @i @s
    \\ totsAppendLemma @o @s

----------------------------------------------------------------------------
-- Non-canonical instructions
----------------------------------------------------------------------------

-- | Helper instruction.
--
-- Checks whether given key present in the storage and fails if it is.
-- This instruction leaves stack intact.
failingWhenPresent
  :: forall c k s v st e.
     ( MemOpHs c, k ~ MemOpKeyHs c
     , KnownValue e
     , st ~ (k & v & c & s)
     )
  => (forall s0. k : s0 :-> e : s0)
  -> st :-> st
failingWhenPresent mkErr =
  dip (dip dup # swap) # swap # dip dup # swap # mem #
  if_ (mkErr # failWith) nop

-- | Like 'update', but throw an error on attempt to overwrite existing entry.
updateNew
  :: forall c k s e.
     ( UpdOpHs c, MemOpHs c, k ~ UpdOpKeyHs c, k ~ MemOpKeyHs c
     , KnownValue e
     )
  => (forall s0. k : s0 :-> e : s0)
  -> k & UpdOpParamsHs c & c & s :-> c & s
updateNew mkErr = failingWhenPresent mkErr # update

class LorentzFunctor (c :: Kind.Type -> Kind.Type) where
  lmap :: KnownValue b => (a : s :-> b : s) -> (c a : s :-> c b : s)

instance LorentzFunctor Maybe where
  lmap f = ifNone none (f # some)

class NonZero t where
  -- | Retain the value only if it is not zero.
  nonZero :: t : s :-> Maybe t : s

instance NonZero Integer where
  nonZero = dup # eq0 # if_ (drop # none) some

instance NonZero Natural where
  nonZero = dup # int # eq0 # if_ (drop # none) some