module Cardano.Wallet.Primitive.Delegation.UTxO
    ( stakeKeyCoinDistr
    ) where

import Prelude

import Cardano.Wallet.Primitive.Types.Address
    ( Address (..) )
import Cardano.Wallet.Primitive.Types.Coin
    ( Coin )
import Cardano.Wallet.Primitive.Types.RewardAccount
    ( RewardAccount (..) )
import Cardano.Wallet.Primitive.Types.Tx
    ( TxOut (..) )
import Cardano.Wallet.Primitive.Types.UTxO
    ( UTxO (..) )
import Data.Map
    ( Map )

import qualified Cardano.Wallet.Primitive.Types.TokenBundle as TokenBundle
import qualified Data.Map as Map

-- | Calculate how much `Coin` exists on each `Maybe RewardAccount` in the
-- `UTxO` given a way to extract `Maybe RewardAccount` from an `Address`.
--
-- This is intended to be used with `rewardAccountFromAddress`, which exists
-- elsewhere because of the cardano-wallet-core / cardano-wallet split.
stakeKeyCoinDistr
    :: (Address -> Maybe RewardAccount)
    -> UTxO
    -> Map (Maybe RewardAccount) Coin
stakeKeyCoinDistr :: (Address -> Maybe RewardAccount)
-> UTxO -> Map (Maybe RewardAccount) Coin
stakeKeyCoinDistr Address -> Maybe RewardAccount
stakeRef =
    (Coin -> Coin -> Coin)
-> [(Maybe RewardAccount, Coin)] -> Map (Maybe RewardAccount) Coin
forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith Coin -> Coin -> Coin
forall a. Semigroup a => a -> a -> a
(<>) ([(Maybe RewardAccount, Coin)] -> Map (Maybe RewardAccount) Coin)
-> (UTxO -> [(Maybe RewardAccount, Coin)])
-> UTxO
-> Map (Maybe RewardAccount) Coin
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TxOut -> (Maybe RewardAccount, Coin))
-> [TxOut] -> [(Maybe RewardAccount, Coin)]
forall a b. (a -> b) -> [a] -> [b]
map TxOut -> (Maybe RewardAccount, Coin)
classifyOut ([TxOut] -> [(Maybe RewardAccount, Coin)])
-> (UTxO -> [TxOut]) -> UTxO -> [(Maybe RewardAccount, Coin)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map TxIn TxOut -> [TxOut]
forall k a. Map k a -> [a]
Map.elems (Map TxIn TxOut -> [TxOut])
-> (UTxO -> Map TxIn TxOut) -> UTxO -> [TxOut]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UTxO -> Map TxIn TxOut
unUTxO
  where
    classifyOut :: TxOut -> (Maybe RewardAccount, Coin)
    classifyOut :: TxOut -> (Maybe RewardAccount, Coin)
classifyOut (TxOut Address
addr TokenBundle
b) = (Address -> Maybe RewardAccount
stakeRef Address
addr, TokenBundle -> Coin
TokenBundle.getCoin TokenBundle
b)