-- | This module exposes the internal state in which our direct simulation is
-- run, and functions to update and query it.
module Cooked.MockChain.MockChainState
  ( MockChainState (..),
    mcstParamsL,
    mcstLedgerStateL,
    mcstOutputsL,
    mcstConstitutionL,
    mcstToUtxoState,
    addOutput,
    removeOutput,
  )
where

import Cardano.Node.Emulator.Internal.Node qualified as Emulator
import Cooked.MockChain.UtxoState
import Cooked.Skeleton
import Data.Default
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Ledger.Orphans ()
import Optics.Core
import Optics.TH
import Plutus.Script.Utils.Scripts qualified as Script
import PlutusLedgerApi.V3 qualified as Api

-- | The state used to run the simulation in 'Cooked.MockChain.Direct'
data MockChainState = MockChainState
  { MockChainState -> Params
mcstParams :: Emulator.Params,
    MockChainState -> EmulatedLedgerState
mcstLedgerState :: Emulator.EmulatedLedgerState,
    -- | Associates to each 'Api.TxOutRef' the 'TxSkelOut' that produced it,
    -- alongside a boolean to state whether this UTxO is still present in the
    -- index ('True') or has already been consumed ('False').
    MockChainState -> Map TxOutRef (TxSkelOut, Bool)
mcstOutputs :: Map Api.TxOutRef (TxSkelOut, Bool),
    -- | The constitution script to be used with proposals
    MockChainState -> Maybe (Versioned Script)
mcstConstitution :: Maybe (Script.Versioned Script.Script)
  }
  deriving (Int -> MockChainState -> ShowS
[MockChainState] -> ShowS
MockChainState -> String
(Int -> MockChainState -> ShowS)
-> (MockChainState -> String)
-> ([MockChainState] -> ShowS)
-> Show MockChainState
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> MockChainState -> ShowS
showsPrec :: Int -> MockChainState -> ShowS
$cshow :: MockChainState -> String
show :: MockChainState -> String
$cshowList :: [MockChainState] -> ShowS
showList :: [MockChainState] -> ShowS
Show)

-- | A lens to set or get the parameters of the 'MockChainState'
makeLensesFor [("mcstParams", "mcstParamsL")] ''MockChainState

-- | A lens to set or get the ledger state of the 'MockChainState'
makeLensesFor [("mcstLedgerState", "mcstLedgerStateL")] ''MockChainState

-- | A lens to set or get the outputs of the 'MockChainState'
makeLensesFor [("mcstOutputs", "mcstOutputsL")] ''MockChainState

-- | A lens to set or get the constitution script of the 'MockChainState'
makeLensesFor [("mcstConstitution", "mcstConstitutionL")] ''MockChainState

instance Default MockChainState where
  def :: MockChainState
def = Params
-> EmulatedLedgerState
-> Map TxOutRef (TxSkelOut, Bool)
-> Maybe (Versioned Script)
-> MockChainState
MockChainState Params
forall a. Default a => a
def (Params -> EmulatedLedgerState
Emulator.initialState Params
forall a. Default a => a
def) Map TxOutRef (TxSkelOut, Bool)
forall k a. Map k a
Map.empty Maybe (Versioned Script)
forall a. Maybe a
Nothing

-- | Builds a 'UtxoState' from a 'MockChainState'
mcstToUtxoState :: MockChainState -> UtxoState
mcstToUtxoState :: MockChainState -> UtxoState
mcstToUtxoState =
  (UtxoState -> (TxOutRef, (TxSkelOut, Bool)) -> UtxoState)
-> UtxoState -> [(TxOutRef, (TxSkelOut, Bool))] -> UtxoState
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl UtxoState -> (TxOutRef, (TxSkelOut, Bool)) -> UtxoState
extractPayload UtxoState
forall a. Monoid a => a
mempty ([(TxOutRef, (TxSkelOut, Bool))] -> UtxoState)
-> (MockChainState -> [(TxOutRef, (TxSkelOut, Bool))])
-> MockChainState
-> UtxoState
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map TxOutRef (TxSkelOut, Bool) -> [(TxOutRef, (TxSkelOut, Bool))]
forall k a. Map k a -> [(k, a)]
Map.toList (Map TxOutRef (TxSkelOut, Bool) -> [(TxOutRef, (TxSkelOut, Bool))])
-> (MockChainState -> Map TxOutRef (TxSkelOut, Bool))
-> MockChainState
-> [(TxOutRef, (TxSkelOut, Bool))]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MockChainState -> Map TxOutRef (TxSkelOut, Bool)
mcstOutputs
  where
    extractPayload :: UtxoState -> (Api.TxOutRef, (TxSkelOut, Bool)) -> UtxoState
    extractPayload :: UtxoState -> (TxOutRef, (TxSkelOut, Bool)) -> UtxoState
extractPayload UtxoState
utxoState (TxOutRef
txOutRef, (TxSkelOut
txSkelOut, Bool
bool)) =
      let newAddress :: Address
newAddress = Optic' A_Getter NoIx TxSkelOut Address -> TxSkelOut -> Address
forall k (is :: IxList) s a.
Is k A_Getter =>
Optic' k is s a -> s -> a
view Optic' A_Getter NoIx TxSkelOut Address
txSkelOutAddressG TxSkelOut
txSkelOut
          newPayloadSet :: UtxoPayloadSet
newPayloadSet =
            [UtxoPayload] -> UtxoPayloadSet
UtxoPayloadSet
              [ TxOutRef
-> Value -> UtxoPayloadDatum -> Maybe ScriptHash -> UtxoPayload
UtxoPayload
                  TxOutRef
txOutRef
                  (Optic' A_Lens NoIx TxSkelOut Value -> TxSkelOut -> Value
forall k (is :: IxList) s a.
Is k A_Getter =>
Optic' k is s a -> s -> a
view Optic' A_Lens NoIx TxSkelOut Value
txSkelOutValueL TxSkelOut
txSkelOut)
                  ( case TxSkelOut
txSkelOut TxSkelOut
-> Optic' A_Lens NoIx TxSkelOut TxSkelOutDatum -> TxSkelOutDatum
forall k s (is :: IxList) a.
Is k A_Getter =>
s -> Optic' k is s a -> a
^. Optic' A_Lens NoIx TxSkelOut TxSkelOutDatum
txSkelOutDatumL of
                      TxSkelOutDatum
NoTxSkelOutDatum -> UtxoPayloadDatum
NoUtxoPayloadDatum
                      SomeTxSkelOutDatum dat
content DatumKind
kind -> dat -> Bool -> UtxoPayloadDatum
forall dat. DatumConstrs dat => dat -> Bool -> UtxoPayloadDatum
SomeUtxoPayloadDatum dat
content (DatumKind
kind DatumKind -> DatumKind -> Bool
forall a. Eq a => a -> a -> Bool
/= DatumKind
Inline)
                  )
                  (Optic' An_AffineFold NoIx TxSkelOut ScriptHash
-> TxSkelOut -> Maybe ScriptHash
forall k (is :: IxList) s a.
Is k An_AffineFold =>
Optic' k is s a -> s -> Maybe a
preview (Lens' TxSkelOut TxSkelOutReferenceScript
txSkelOutReferenceScriptL Lens' TxSkelOut TxSkelOutReferenceScript
-> Optic
     An_AffineFold
     NoIx
     TxSkelOutReferenceScript
     TxSkelOutReferenceScript
     ScriptHash
     ScriptHash
-> Optic' An_AffineFold NoIx TxSkelOut ScriptHash
forall k l m (is :: IxList) (js :: IxList) (ks :: IxList) s t u v a
       b.
(JoinKinds k l m, AppendIndices is js ks) =>
Optic k is s t u v -> Optic l js u v a b -> Optic m ks s t a b
% Optic
  An_AffineFold
  NoIx
  TxSkelOutReferenceScript
  TxSkelOutReferenceScript
  ScriptHash
  ScriptHash
txSkelOutReferenceScriptHashAF) TxSkelOut
txSkelOut)
              ]
       in if Bool
bool
            then UtxoState
utxoState {availableUtxos = Map.insertWith (<>) newAddress newPayloadSet (availableUtxos utxoState)}
            else UtxoState
utxoState {consumedUtxos = Map.insertWith (<>) newAddress newPayloadSet (consumedUtxos utxoState)}

-- | Stores an output in a 'MockChainState'
addOutput :: Api.TxOutRef -> TxSkelOut -> MockChainState -> MockChainState
addOutput :: TxOutRef -> TxSkelOut -> MockChainState -> MockChainState
addOutput TxOutRef
oRef TxSkelOut
txSkelOut = Lens' MockChainState (Map TxOutRef (TxSkelOut, Bool))
-> (Map TxOutRef (TxSkelOut, Bool)
    -> Map TxOutRef (TxSkelOut, Bool))
-> MockChainState
-> MockChainState
forall k (is :: IxList) s t a b.
Is k A_Setter =>
Optic k is s t a b -> (a -> b) -> s -> t
over Lens' MockChainState (Map TxOutRef (TxSkelOut, Bool))
mcstOutputsL (TxOutRef
-> (TxSkelOut, Bool)
-> Map TxOutRef (TxSkelOut, Bool)
-> Map TxOutRef (TxSkelOut, Bool)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert TxOutRef
oRef (TxSkelOut
txSkelOut, Bool
True))

-- | Removes an output from the 'MockChainState'
removeOutput :: Api.TxOutRef -> MockChainState -> MockChainState
removeOutput :: TxOutRef -> MockChainState -> MockChainState
removeOutput TxOutRef
oRef = Lens' MockChainState (Map TxOutRef (TxSkelOut, Bool))
-> (Map TxOutRef (TxSkelOut, Bool)
    -> Map TxOutRef (TxSkelOut, Bool))
-> MockChainState
-> MockChainState
forall k (is :: IxList) s t a b.
Is k A_Setter =>
Optic k is s t a b -> (a -> b) -> s -> t
over Lens' MockChainState (Map TxOutRef (TxSkelOut, Bool))
mcstOutputsL (((TxSkelOut, Bool) -> Maybe (TxSkelOut, Bool))
-> TxOutRef
-> Map TxOutRef (TxSkelOut, Bool)
-> Map TxOutRef (TxSkelOut, Bool)
forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
Map.update (\(TxSkelOut
output, Bool
_) -> (TxSkelOut, Bool) -> Maybe (TxSkelOut, Bool)
forall a. a -> Maybe a
Just (TxSkelOut
output, Bool
False)) TxOutRef
oRef)