{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}

module Test.Ouroboros.Network.PeerSelection.Utils where

import Control.Monad.Class.MonadTime.SI
import Data.Set (Set)
import Data.Set qualified as Set
import System.Random (mkStdGen)

import Ouroboros.Network.PeerSelection.Governor (AssociationMode (..),
           PeerSelectionTargets (..), TracePeerSelection (..))
import Ouroboros.Network.PeerSelection.Governor qualified as Governor
import Ouroboros.Network.PeerSelection.PublicRootPeers qualified as PublicRootPeers

import Test.Ouroboros.Network.Data.Signal (Events, Signal)
import Test.Ouroboros.Network.Data.Signal qualified as Signal

import Test.Ouroboros.Network.PeerSelection.Cardano.MockEnvironment hiding
           (targets, tests)
import Test.Ouroboros.Network.PeerSelection.Instances


takeFirstNHours :: DiffTime -> [(Time, a)] -> [(Time, a)]
takeFirstNHours :: forall a. DiffTime -> [(Time, a)] -> [(Time, a)]
takeFirstNHours DiffTime
h = ((Time, a) -> Bool) -> [(Time, a)] -> [(Time, a)]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (\(Time
t,a
_) -> Time
t Time -> Time -> Bool
forall a. Ord a => a -> a -> Bool
< DiffTime -> Time
Time (DiffTime
60DiffTime -> DiffTime -> DiffTime
forall a. Num a => a -> a -> a
*DiffTime
60DiffTime -> DiffTime -> DiffTime
forall a. Num a => a -> a -> a
*DiffTime
h))


selectEnvEvents :: Events (TestTraceEvent extraState extraFlags extraPeers extraCounters) -> Events TraceMockEnv
selectEnvEvents :: forall extraState extraFlags extraPeers extraCounters.
Events
  (TestTraceEvent extraState extraFlags extraPeers extraCounters)
-> Events TraceMockEnv
selectEnvEvents = (TestTraceEvent extraState extraFlags extraPeers extraCounters
 -> Maybe TraceMockEnv)
-> Events
     (TestTraceEvent extraState extraFlags extraPeers extraCounters)
-> Events TraceMockEnv
forall a b. (a -> Maybe b) -> Events a -> Events b
Signal.selectEvents
                    (\case MockEnvEvent TraceMockEnv
e -> TraceMockEnv -> Maybe TraceMockEnv
forall a. a -> Maybe a
Just (TraceMockEnv -> Maybe TraceMockEnv)
-> TraceMockEnv -> Maybe TraceMockEnv
forall a b. (a -> b) -> a -> b
$! TraceMockEnv
e
                           TestTraceEvent extraState extraFlags extraPeers extraCounters
_              -> Maybe TraceMockEnv
forall a. Maybe a
Nothing)

selectGovEvents :: Events (TestTraceEvent extraState extraFlags extraPeers extracounters)
                -> Events (TracePeerSelection extraState extraFlags extraPeers PeerAddr)
selectGovEvents :: forall extraState extraFlags extraPeers extracounters.
Events
  (TestTraceEvent extraState extraFlags extraPeers extracounters)
-> Events
     (TracePeerSelection extraState extraFlags extraPeers PeerAddr)
selectGovEvents = (TestTraceEvent extraState extraFlags extraPeers extracounters
 -> Maybe
      (TracePeerSelection extraState extraFlags extraPeers PeerAddr))
-> Events
     (TestTraceEvent extraState extraFlags extraPeers extracounters)
-> Events
     (TracePeerSelection extraState extraFlags extraPeers PeerAddr)
forall a b. (a -> Maybe b) -> Events a -> Events b
Signal.selectEvents
                    (\case GovernorEvent TracePeerSelection extraState extraFlags extraPeers PeerAddr
e -> TracePeerSelection extraState extraFlags extraPeers PeerAddr
-> Maybe
     (TracePeerSelection extraState extraFlags extraPeers PeerAddr)
forall a. a -> Maybe a
Just (TracePeerSelection extraState extraFlags extraPeers PeerAddr
 -> Maybe
      (TracePeerSelection extraState extraFlags extraPeers PeerAddr))
-> TracePeerSelection extraState extraFlags extraPeers PeerAddr
-> Maybe
     (TracePeerSelection extraState extraFlags extraPeers PeerAddr)
forall a b. (a -> b) -> a -> b
$! TracePeerSelection extraState extraFlags extraPeers PeerAddr
e
                           TestTraceEvent extraState extraFlags extraPeers extracounters
_               -> Maybe
  (TracePeerSelection extraState extraFlags extraPeers PeerAddr)
forall a. Maybe a
Nothing)

selectGovCounters :: Events (TestTraceEvent extraState extraFlags extraPeers extraCounters)
                  -> Events (Governor.PeerSelectionCounters extraCounters)
selectGovCounters :: forall extraState extraFlags extraPeers extraCounters.
Events
  (TestTraceEvent extraState extraFlags extraPeers extraCounters)
-> Events (PeerSelectionCounters extraCounters)
selectGovCounters = (TestTraceEvent extraState extraFlags extraPeers extraCounters
 -> Maybe (PeerSelectionCounters extraCounters))
-> Events
     (TestTraceEvent extraState extraFlags extraPeers extraCounters)
-> Events (PeerSelectionCounters extraCounters)
forall a b. (a -> Maybe b) -> Events a -> Events b
Signal.selectEvents
                      (\case GovernorCounters PeerSelectionCounters extraCounters
e -> PeerSelectionCounters extraCounters
-> Maybe (PeerSelectionCounters extraCounters)
forall a. a -> Maybe a
Just (PeerSelectionCounters extraCounters
 -> Maybe (PeerSelectionCounters extraCounters))
-> PeerSelectionCounters extraCounters
-> Maybe (PeerSelectionCounters extraCounters)
forall a b. (a -> b) -> a -> b
$! PeerSelectionCounters extraCounters
e
                             TestTraceEvent extraState extraFlags extraPeers extraCounters
_                  -> Maybe (PeerSelectionCounters extraCounters)
forall a. Maybe a
Nothing)

selectGovAssociationMode :: Events (TestTraceEvent extraState extraFlags extraPeers extraCounters)
                         -> Events AssociationMode
selectGovAssociationMode :: forall extraState extraFlags extraPeers extraCounters.
Events
  (TestTraceEvent extraState extraFlags extraPeers extraCounters)
-> Events AssociationMode
selectGovAssociationMode = (TestTraceEvent extraState extraFlags extraPeers extraCounters
 -> Maybe AssociationMode)
-> Events
     (TestTraceEvent extraState extraFlags extraPeers extraCounters)
-> Events AssociationMode
forall a b. (a -> Maybe b) -> Events a -> Events b
Signal.selectEvents
                             (\case GovernorAssociationMode AssociationMode
e -> AssociationMode -> Maybe AssociationMode
forall a. a -> Maybe a
Just (AssociationMode -> Maybe AssociationMode)
-> AssociationMode -> Maybe AssociationMode
forall a b. (a -> b) -> a -> b
$! AssociationMode
e
                                    TestTraceEvent extraState extraFlags extraPeers extraCounters
_                         -> Maybe AssociationMode
forall a. Maybe a
Nothing)

selectGovState :: Eq a
               => (forall peerconn. Governor.PeerSelectionState extraState extraFlags extraPeers PeerAddr peerconn -> a)
               -> extraState
               -> extraPeers
               -> Events (TestTraceEvent extraState extraFlags extraPeers extraCounters)
               -> Signal a
selectGovState :: forall a extraState extraFlags extraPeers extraCounters.
Eq a =>
(forall peerconn.
 PeerSelectionState
   extraState extraFlags extraPeers PeerAddr peerconn
 -> a)
-> extraState
-> extraPeers
-> Events
     (TestTraceEvent extraState extraFlags extraPeers extraCounters)
-> Signal a
selectGovState forall peerconn.
PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
-> a
f extraState
es extraPeers
ep =
    Signal a -> Signal a
forall a. Eq a => Signal a -> Signal a
Signal.nub
  -- TODO: #3182 Rng seed should come from quickcheck.
  --       and `NumberOfBigLedgerPeers`
  (Signal a -> Signal a)
-> (Events
      (TestTraceEvent extraState extraFlags extraPeers extraCounters)
    -> Signal a)
-> Events
     (TestTraceEvent extraState extraFlags extraPeers extraCounters)
-> Signal a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Events a -> Signal a
forall a. a -> Events a -> Signal a
Signal.fromChangeEvents (PeerSelectionState
  extraState extraFlags extraPeers PeerAddr (ZonkAny 0)
-> a
forall peerconn.
PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
-> a
f (PeerSelectionState
   extraState extraFlags extraPeers PeerAddr (ZonkAny 0)
 -> a)
-> PeerSelectionState
     extraState extraFlags extraPeers PeerAddr (ZonkAny 0)
-> a
forall a b. (a -> b) -> a -> b
$! StdGen
-> extraState
-> extraPeers
-> PeerSelectionState
     extraState extraFlags extraPeers PeerAddr (ZonkAny 0)
forall extraState extraPeers extraFlags peeraddr peerconn.
StdGen
-> extraState
-> extraPeers
-> PeerSelectionState
     extraState extraFlags extraPeers peeraddr peerconn
Governor.emptyPeerSelectionState (Int -> StdGen
mkStdGen Int
42) extraState
es extraPeers
ep)
  (Events a -> Signal a)
-> (Events
      (TestTraceEvent extraState extraFlags extraPeers extraCounters)
    -> Events a)
-> Events
     (TestTraceEvent extraState extraFlags extraPeers extraCounters)
-> Signal a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TestTraceEvent extraState extraFlags extraPeers extraCounters
 -> Maybe a)
-> Events
     (TestTraceEvent extraState extraFlags extraPeers extraCounters)
-> Events a
forall a b. (a -> Maybe b) -> Events a -> Events b
Signal.selectEvents
      (\case GovernorDebug (Governor.TraceGovernorState Time
_ Maybe DiffTime
_ PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
st) -> a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$! PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
-> a
forall peerconn.
PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
-> a
f PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
st
             TestTraceEvent extraState extraFlags extraPeers extraCounters
_                                                  -> Maybe a
forall a. Maybe a
Nothing)

selectEnvTargets :: Eq a
                 => (PeerSelectionTargets -> a)
                 -> Events (TestTraceEvent extraState extraFlags extraPeers extraCounters)
                 -> Signal a
selectEnvTargets :: forall a extraState extraFlags extraPeers extraCounters.
Eq a =>
(PeerSelectionTargets -> a)
-> Events
     (TestTraceEvent extraState extraFlags extraPeers extraCounters)
-> Signal a
selectEnvTargets PeerSelectionTargets -> a
f =
    Signal a -> Signal a
forall a. Eq a => Signal a -> Signal a
Signal.nub
  (Signal a -> Signal a)
-> (Events
      (TestTraceEvent extraState extraFlags extraPeers extraCounters)
    -> Signal a)
-> Events
     (TestTraceEvent extraState extraFlags extraPeers extraCounters)
-> Signal a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (PeerSelectionTargets -> a)
-> Signal PeerSelectionTargets -> Signal a
forall a b. (a -> b) -> Signal a -> Signal b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap PeerSelectionTargets -> a
f
  (Signal PeerSelectionTargets -> Signal a)
-> (Events
      (TestTraceEvent extraState extraFlags extraPeers extraCounters)
    -> Signal PeerSelectionTargets)
-> Events
     (TestTraceEvent extraState extraFlags extraPeers extraCounters)
-> Signal a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PeerSelectionTargets
-> Events PeerSelectionTargets -> Signal PeerSelectionTargets
forall a. a -> Events a -> Signal a
Signal.fromChangeEvents PeerSelectionTargets
Governor.nullPeerSelectionTargets
  (Events PeerSelectionTargets -> Signal PeerSelectionTargets)
-> (Events
      (TestTraceEvent extraState extraFlags extraPeers extraCounters)
    -> Events PeerSelectionTargets)
-> Events
     (TestTraceEvent extraState extraFlags extraPeers extraCounters)
-> Signal PeerSelectionTargets
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TraceMockEnv -> Maybe PeerSelectionTargets)
-> Events TraceMockEnv -> Events PeerSelectionTargets
forall a b. (a -> Maybe b) -> Events a -> Events b
Signal.selectEvents
      (\case TraceEnvSetTargets PeerSelectionTargets
targets -> PeerSelectionTargets -> Maybe PeerSelectionTargets
forall a. a -> Maybe a
Just (PeerSelectionTargets -> Maybe PeerSelectionTargets)
-> PeerSelectionTargets -> Maybe PeerSelectionTargets
forall a b. (a -> b) -> a -> b
$! PeerSelectionTargets
targets
             TraceMockEnv
_                          -> Maybe PeerSelectionTargets
forall a. Maybe a
Nothing)
  (Events TraceMockEnv -> Events PeerSelectionTargets)
-> (Events
      (TestTraceEvent extraState extraFlags extraPeers extraCounters)
    -> Events TraceMockEnv)
-> Events
     (TestTraceEvent extraState extraFlags extraPeers extraCounters)
-> Events PeerSelectionTargets
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Events
  (TestTraceEvent extraState extraFlags extraPeers extraCounters)
-> Events TraceMockEnv
forall extraState extraFlags extraPeers extraCounters.
Events
  (TestTraceEvent extraState extraFlags extraPeers extraCounters)
-> Events TraceMockEnv
selectEnvEvents


-- | filter big ledger peers
--
takeBigLedgerPeers
    :: (Governor.PeerSelectionState extraState extraFlags extraPeers PeerAddr peerconn -> Set PeerAddr)
    ->  Governor.PeerSelectionState extraState extraFlags extraPeers PeerAddr peerconn -> Set PeerAddr
takeBigLedgerPeers :: forall extraState extraFlags extraPeers peerconn.
(PeerSelectionState
   extraState extraFlags extraPeers PeerAddr peerconn
 -> Set PeerAddr)
-> PeerSelectionState
     extraState extraFlags extraPeers PeerAddr peerconn
-> Set PeerAddr
takeBigLedgerPeers PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
-> Set PeerAddr
f =
  \PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
st -> PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
-> Set PeerAddr
f PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
st Set PeerAddr -> Set PeerAddr -> Set PeerAddr
forall a. Ord a => Set a -> Set a -> Set a
`Set.intersection` (PublicRootPeers extraPeers PeerAddr -> Set PeerAddr
forall extraPeers peeraddr.
PublicRootPeers extraPeers peeraddr -> Set peeraddr
PublicRootPeers.getBigLedgerPeers (PublicRootPeers extraPeers PeerAddr -> Set PeerAddr)
-> (PeerSelectionState
      extraState extraFlags extraPeers PeerAddr peerconn
    -> PublicRootPeers extraPeers PeerAddr)
-> PeerSelectionState
     extraState extraFlags extraPeers PeerAddr peerconn
-> Set PeerAddr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
-> PublicRootPeers extraPeers PeerAddr
forall extraState extraFlags extraPeers peeraddr peerconn.
PeerSelectionState
  extraState extraFlags extraPeers peeraddr peerconn
-> PublicRootPeers extraPeers peeraddr
Governor.publicRootPeers) PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
st

-- | filter out big ledger peers
--
dropBigLedgerPeers
    :: (Governor.PeerSelectionState extraState extraFlags extraPeers PeerAddr peerconn -> Set PeerAddr)
    ->  Governor.PeerSelectionState extraState extraFlags extraPeers PeerAddr peerconn -> Set PeerAddr
dropBigLedgerPeers :: forall extraState extraFlags extraPeers peerconn.
(PeerSelectionState
   extraState extraFlags extraPeers PeerAddr peerconn
 -> Set PeerAddr)
-> PeerSelectionState
     extraState extraFlags extraPeers PeerAddr peerconn
-> Set PeerAddr
dropBigLedgerPeers PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
-> Set PeerAddr
f =
  \PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
st -> PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
-> Set PeerAddr
f PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
st Set PeerAddr -> Set PeerAddr -> Set PeerAddr
forall a. Ord a => Set a -> Set a -> Set a
Set.\\ (PublicRootPeers extraPeers PeerAddr -> Set PeerAddr
forall extraPeers peeraddr.
PublicRootPeers extraPeers peeraddr -> Set peeraddr
PublicRootPeers.getBigLedgerPeers (PublicRootPeers extraPeers PeerAddr -> Set PeerAddr)
-> (PeerSelectionState
      extraState extraFlags extraPeers PeerAddr peerconn
    -> PublicRootPeers extraPeers PeerAddr)
-> PeerSelectionState
     extraState extraFlags extraPeers PeerAddr peerconn
-> Set PeerAddr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
-> PublicRootPeers extraPeers PeerAddr
forall extraState extraFlags extraPeers peeraddr peerconn.
PeerSelectionState
  extraState extraFlags extraPeers peeraddr peerconn
-> PublicRootPeers extraPeers peeraddr
Governor.publicRootPeers) PeerSelectionState
  extraState extraFlags extraPeers PeerAddr peerconn
st