{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE TemplateHaskell     #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}

-- | Waiting for things to happen
module Plutus.Trace.Effects.Waiting(
    Waiting(..)
    , waitUntilSlot
    , waitUntilTime
    , nextSlot
    , waitNSlots
    , waitNMilliSeconds
    , handleWaiting
    ) where

import Cardano.Node.Emulator.Internal.Node.TimeSlot qualified as TimeSlot
import Control.Monad.Freer (Eff, Member, type (~>))
import Control.Monad.Freer.Coroutine (Yield)
import Control.Monad.Freer.TH (makeEffect)
import Ledger.Slot (Slot)
import Ledger.Time (DiffMilliSeconds, POSIXTime, fromMilliSeconds)
import Numeric.Natural (Natural)
import Plutus.Trace.Emulator.Types (EmulatorMessage (NewSlot))
import Plutus.Trace.Scheduler (EmSystemCall, Priority (Sleeping), sleep)

data Waiting r where
    WaitUntilSlot :: Slot -> Waiting Slot
    GetSlotConfig :: Waiting TimeSlot.SlotConfig

makeEffect ''Waiting

-- | Wait until the slot where the given time falls into and return latest time
-- we know has passed.
waitUntilTime :: Member Waiting effs => POSIXTime -> Eff effs POSIXTime
waitUntilTime :: POSIXTime -> Eff effs POSIXTime
waitUntilTime POSIXTime
time = do
    SlotConfig
slotConfig <- Eff effs SlotConfig
forall (effs :: [* -> *]).
Member Waiting effs =>
Eff effs SlotConfig
getSlotConfig
    Slot
slot <- Slot -> Eff effs Slot
forall (effs :: [* -> *]).
Member Waiting effs =>
Slot -> Eff effs Slot
waitUntilSlot (SlotConfig -> POSIXTime -> Slot
TimeSlot.posixTimeToEnclosingSlot SlotConfig
slotConfig POSIXTime
time)
    POSIXTime -> Eff effs POSIXTime
forall (m :: * -> *) a. Monad m => a -> m a
return (POSIXTime -> Eff effs POSIXTime)
-> POSIXTime -> Eff effs POSIXTime
forall a b. (a -> b) -> a -> b
$ SlotConfig -> Slot -> POSIXTime
TimeSlot.slotToEndPOSIXTime SlotConfig
slotConfig Slot
slot

-- | Wait until the beginning of the next slot, returning
--   the new slot number.
nextSlot :: Member Waiting effs => Eff effs Slot
nextSlot :: Eff effs Slot
nextSlot = Slot -> Eff effs Slot
forall (effs :: [* -> *]).
Member Waiting effs =>
Slot -> Eff effs Slot
waitUntilSlot Slot
0

-- | Wait for a number of slots
waitNSlots ::
    forall effs.
    ( Member Waiting effs )
    => Natural
    -> Eff effs Slot
waitNSlots :: Natural -> Eff effs Slot
waitNSlots Natural
n
    | Natural
n Natural -> Natural -> Bool
forall a. Ord a => a -> a -> Bool
> Natural
1 = Eff effs Slot
forall (effs :: [* -> *]). Member Waiting effs => Eff effs Slot
nextSlot Eff effs Slot -> Eff effs Slot -> Eff effs Slot
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Natural -> Eff effs Slot
forall (effs :: [* -> *]).
Member Waiting effs =>
Natural -> Eff effs Slot
waitNSlots (Natural
n Natural -> Natural -> Natural
forall a. Num a => a -> a -> a
- Natural
1)
    | Bool
otherwise = Eff effs Slot
forall (effs :: [* -> *]). Member Waiting effs => Eff effs Slot
nextSlot

-- | Convert the given 'n' milliseconds to a number of slots to wait.
--
-- Note: Currently, if n < length of a slot, then 'waitNMilliSeconds' has no
-- effect.
waitNMilliSeconds ::
    forall effs.
    ( Member Waiting effs )
    => DiffMilliSeconds
    -> Eff effs Slot
waitNMilliSeconds :: DiffMilliSeconds -> Eff effs Slot
waitNMilliSeconds DiffMilliSeconds
n = do
    SlotConfig
slotConfig <- Eff effs SlotConfig
forall (effs :: [* -> *]).
Member Waiting effs =>
Eff effs SlotConfig
getSlotConfig
    Natural -> Eff effs Slot
forall (effs :: [* -> *]).
Member Waiting effs =>
Natural -> Eff effs Slot
waitNSlots (Slot -> Natural
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Slot -> Natural) -> Slot -> Natural
forall a b. (a -> b) -> a -> b
$ SlotConfig -> POSIXTime -> Slot
TimeSlot.posixTimeToEnclosingSlot SlotConfig
slotConfig (POSIXTime -> Slot) -> POSIXTime -> Slot
forall a b. (a -> b) -> a -> b
$ DiffMilliSeconds -> POSIXTime
fromMilliSeconds DiffMilliSeconds
n)

handleWaiting ::
    forall effs effs2 a.
    ( Member (Yield (EmSystemCall effs2 EmulatorMessage a) (Maybe EmulatorMessage)) effs
    )
    => TimeSlot.SlotConfig
    -> Waiting
    ~> Eff effs
handleWaiting :: SlotConfig -> Waiting ~> Eff effs
handleWaiting SlotConfig
slotConfig = \case
    Waiting x
GetSlotConfig -> SlotConfig -> Eff effs SlotConfig
forall (f :: * -> *) a. Applicative f => a -> f a
pure SlotConfig
slotConfig
    WaitUntilSlot Slot
s -> Eff effs x
Eff effs Slot
go where
        go :: Eff effs Slot
go = Priority -> Eff effs (Maybe EmulatorMessage)
forall (effs :: [* -> *]) systemEvent (effs2 :: [* -> *]) a.
Member
  (Yield (EmSystemCall effs systemEvent a) (Maybe systemEvent))
  effs2 =>
Priority -> Eff effs2 (Maybe systemEvent)
sleep @effs2 @_ @_ @a Priority
Sleeping Eff effs (Maybe EmulatorMessage)
-> (Maybe EmulatorMessage -> Eff effs Slot) -> Eff effs Slot
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case { Just (NewSlot [Block]
_ Slot
sl) | Slot
sl Slot -> Slot -> Bool
forall a. Ord a => a -> a -> Bool
>= Slot
s -> Slot -> Eff effs Slot
forall (f :: * -> *) a. Applicative f => a -> f a
pure Slot
sl; Maybe EmulatorMessage
_ -> Eff effs Slot
go }