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

module Test.Cardano.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.Instances (PeerAddr (..))

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


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 extraTrace)
                -> Events TraceMockEnv
selectEnvEvents :: forall extraState extraFlags extraPeers extraCounters extraTrace.
Events
  (TestTraceEvent
     extraState extraFlags extraPeers extraCounters extraTrace)
-> Events TraceMockEnv
selectEnvEvents = (TestTraceEvent
   extraState extraFlags extraPeers extraCounters extraTrace
 -> Maybe TraceMockEnv)
-> Events
     (TestTraceEvent
        extraState extraFlags extraPeers extraCounters extraTrace)
-> 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 extraTrace
_              -> Maybe TraceMockEnv
forall a. Maybe a
Nothing)

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

selectGovCounters :: Events (TestTraceEvent extraState extraFlags extraPeers
                                            extraCounters extraTrace)
                  -> Events (Governor.PeerSelectionCounters extraCounters)
selectGovCounters :: forall extraState extraFlags extraPeers extraCounters extraTrace.
Events
  (TestTraceEvent
     extraState extraFlags extraPeers extraCounters extraTrace)
-> Events (PeerSelectionCounters extraCounters)
selectGovCounters = (TestTraceEvent
   extraState extraFlags extraPeers extraCounters extraTrace
 -> Maybe (PeerSelectionCounters extraCounters))
-> Events
     (TestTraceEvent
        extraState extraFlags extraPeers extraCounters extraTrace)
-> 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 extraTrace
_                  -> Maybe (PeerSelectionCounters extraCounters)
forall a. Maybe a
Nothing)

selectGovAssociationMode :: Events (TestTraceEvent extraState extraFlags extraPeers
                                                   extraCounters extraTrace)
                         -> Events AssociationMode
selectGovAssociationMode :: forall extraState extraFlags extraPeers extraCounters extraTrace.
Events
  (TestTraceEvent
     extraState extraFlags extraPeers extraCounters extraTrace)
-> Events AssociationMode
selectGovAssociationMode = (TestTraceEvent
   extraState extraFlags extraPeers extraCounters extraTrace
 -> Maybe AssociationMode)
-> Events
     (TestTraceEvent
        extraState extraFlags extraPeers extraCounters extraTrace)
-> 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 extraTrace
_                         -> 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 extraTrace)
               -> Signal a
selectGovState :: forall a extraState extraFlags extraPeers extraCounters extraTrace.
Eq a =>
(forall peerconn.
 PeerSelectionState
   extraState extraFlags extraPeers PeerAddr peerconn
 -> a)
-> extraState
-> extraPeers
-> Events
     (TestTraceEvent
        extraState extraFlags extraPeers extraCounters extraTrace)
-> 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 extraTrace)
    -> Signal a)
-> Events
     (TestTraceEvent
        extraState extraFlags extraPeers extraCounters extraTrace)
-> 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 extraTrace)
    -> Events a)
-> Events
     (TestTraceEvent
        extraState extraFlags extraPeers extraCounters extraTrace)
-> Signal a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TestTraceEvent
   extraState extraFlags extraPeers extraCounters extraTrace
 -> Maybe a)
-> Events
     (TestTraceEvent
        extraState extraFlags extraPeers extraCounters extraTrace)
-> 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 extraTrace
_                                                  -> Maybe a
forall a. Maybe a
Nothing)

selectEnvTargets :: Eq a
                 => (PeerSelectionTargets -> a)
                 -> Events (TestTraceEvent extraState extraFlags extraPeers
                                           extraCounters extraTrace)
                 -> Signal a
selectEnvTargets :: forall a extraState extraFlags extraPeers extraCounters extraTrace.
Eq a =>
(PeerSelectionTargets -> a)
-> Events
     (TestTraceEvent
        extraState extraFlags extraPeers extraCounters extraTrace)
-> 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 extraTrace)
    -> Signal a)
-> Events
     (TestTraceEvent
        extraState extraFlags extraPeers extraCounters extraTrace)
-> 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 extraTrace)
    -> Signal PeerSelectionTargets)
-> Events
     (TestTraceEvent
        extraState extraFlags extraPeers extraCounters extraTrace)
-> 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 extraTrace)
    -> Events PeerSelectionTargets)
-> Events
     (TestTraceEvent
        extraState extraFlags extraPeers extraCounters extraTrace)
-> 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 extraTrace)
    -> Events TraceMockEnv)
-> Events
     (TestTraceEvent
        extraState extraFlags extraPeers extraCounters extraTrace)
-> Events PeerSelectionTargets
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Events
  (TestTraceEvent
     extraState extraFlags extraPeers extraCounters extraTrace)
-> Events TraceMockEnv
forall extraState extraFlags extraPeers extraCounters extraTrace.
Events
  (TestTraceEvent
     extraState extraFlags extraPeers extraCounters extraTrace)
-> 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