content: {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE GeneralisedNewtypeDeriving #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE TemplateHaskell #-}
  {-# LANGUAGE ViewPatterns #-}
  {-# OPTIONS_GHC -Wno-deprecations #-}
  
  -- | This module simulates running multiple copies of consensus together in a
  --  deterministic fashion, without consideration for real time.  The goal is to
  --  have a faster and more reproducible way of testing/profiling/benchmarking
  --  performance-related issues.
  --
  --  Note that it is expected that you will edit this file depending on what
  --  you wish to test.
  module Main where
  
  import Control.Monad
  import Control.Monad.IO.Class
  import Control.Monad.Reader.Class
  import Control.Monad.State.Class
  import Control.Monad.Trans.Reader (ReaderT (..))
  import Control.Monad.Trans.State.Strict (StateT (..), execStateT)
  import qualified Data.PQueue.Min as MinPQ
  import qualified Data.Sequence as Seq
  import Data.Time.Clock
  import Data.Time.Clock.POSIX
  import qualified Data.Vector as Vec
  import Data.Word
  import Lens.Micro.Platform
  import System.IO
  import System.Mem
  import System.Random
  
  import Concordium.Afgjort.Finalize.Types
  import Concordium.GlobalState
  import qualified Concordium.GlobalState.AccountMap.LMDB as LMDBAccountMap
  import Concordium.GlobalState.BakerInfo
  import Concordium.GlobalState.Block
  import qualified Concordium.GlobalState.BlockPointer as BS
  import Concordium.GlobalState.Finalization
  import Concordium.GlobalState.Parameters
  import qualified Concordium.GlobalState.TreeState as TS
  import Concordium.Types
  import Concordium.Types.AnonymityRevokers
  import Concordium.Types.IdentityProviders
  import Concordium.Types.Transactions
  
  import Concordium.Afgjort.Finalize
  import Concordium.Birk.Bake
  import Concordium.Kontrol
  import Concordium.Logger
  import Concordium.Skov.MonadImplementations
  import Concordium.Startup
  import Concordium.TimeMonad
  import Concordium.TimerMonad
  
  import qualified Concordium.Crypto.DummyData as Dummy
  import Concordium.GlobalState.DummyData (dummyKeyCollection)
  import qualified Concordium.GlobalState.DummyData as Dummy
  import qualified Concordium.Types.DummyData as Dummy
  
  import qualified Data.ByteString.Lazy as LBS
  import qualified Data.Serialize as S
  import System.Directory
  
  -- | Protocol version
  type PV = 'P5
  
  -- | Construct the global state configuration.
  --  Can be customised if changing the configuration.
  makeGlobalStateConfig :: RuntimeParameters -> FilePath -> FilePath -> FilePath -> IO GlobalStateConfig
  makeGlobalStateConfig rt treeStateDir blockStateFile accMapDirectory = do
      accMap <- LMDBAccountMap.openDatabase accMapDirectory
      return $ GlobalStateConfig rt treeStateDir blockStateFile accMap
  
  {-
  type TreeConfig = PairGSConfig MemoryTreeMemoryBlockConfig DiskTreeDiskBlockConfig
  
  -- |Construct the global state configuration.
  -- Can be customised if changing the configuration.
  makeGlobalStateConfig :: RuntimeParameters -> FilePath -> FilePath -> IO TreeConfig
  makeGlobalStateConfig rp treeStateDir blockStateFile =
     return $ PairGSConfig (MTMBConfig rp, GlobalStateConfig rp treeStateDir blockStateFile)
  -}
  
  -- | A timer is represented as an integer identifier.
  --  Timers are issued with increasing identifiers.
  newtype DummyTimer = DummyTimer Integer
      deriving (Num, Eq, Ord)
  
  -- | Configuration to use for bakers.
  --  Can be customised for different global state configurations (disk/memory/paired)
  --  or to enable/disable finalization buffering.
  type BakerConfig = SkovConfig PV (ActiveFinalization DummyTimer) NoHandler
  
  -- | The identity providers to use.
  dummyIdentityProviders :: IdentityProviders
  dummyIdentityProviders = emptyIdentityProviders
  
  dummyArs :: AnonymityRevokers
  dummyArs = emptyAnonymityRevokers
  
  -- | Monad that provides a deterministic implementation of 'TimeMonad' -- i.e. that is
  --  not dependent on real time.
  newtype DeterministicTime m a = DeterministicTime {runDeterministic' :: ReaderT UTCTime m a}
      deriving (Functor, Applicative, Monad, MonadIO)
  
  instance (Monad m) => TimeMonad (DeterministicTime m) where
      currentTime = DeterministicTime ask
  
  -- | Run the 'DeterministicTime' action with the given time.
  runDeterministic :: DeterministicTime m a -> UTCTime -> m a
  runDeterministic (DeterministicTime a) t = runReaderT a t
  
  -- | The base monad (with logging and time).
  type LogBase = LoggerT (DeterministicTime IO)
  
  -- | How 'SkovHandlers' should be instantiated for our setting.
  type MyHandlers = SkovHandlers PV DummyTimer BakerConfig (StateT SimState LogBase)
  
  -- | The monad for bakers to run in.
  type BakerM = SkovT PV MyHandlers BakerConfig (StateT SimState LogBase)
  
  -- | Events that trigger actions by bakers.
  data Event
      = -- | Attempt to bake a block in the given slot; generates a new event for the next slot
        EBake !Slot
      | -- | Receive a block
        EBlock !BakedBlock
      | -- | Receive a transaction
        ETransaction !BlockItem
      | -- | Receive a finalization message
        EFinalization !FinalizationPseudoMessage
      | -- | Receive a finalization record
        EFinalizationRecord !FinalizationRecord
      | -- | Trigger a timer event
        ETimer !DummyTimer !(BakerM ())
  
  instance Show Event where
      show (EBake sl) = "Bake in slot " ++ show sl
      show (EBlock _) = "Receive block"
      show (ETransaction _) = "Receive transaction"
      show (EFinalization _) = "Receive finalization message"
      show (EFinalizationRecord _) = "Receive finalization record"
      show (ETimer _ _) = "Timer event"
  
  -- | Both baker-specific and generic events.
  data GEvent
      = -- | An event for a particular baker
        BakerEvent !Int !Event
      | -- | Spawn the next transaction to send to bakers
        TransactionEvent [(Integer, BlockItem)]
  
  -- | An event with a time at which it should occur.
  --  The time is used for determining priority.
  data PEvent = PEvent !Integer !GEvent
  
  instance Eq PEvent where
      (PEvent i1 _) == (PEvent i2 _) = i1 == i2
  instance Ord PEvent where
      compare (PEvent i1 _) (PEvent i2 _) = compare i1 i2
  
  -- | The state of a particular baker.
  data BakerState = BakerState
      { _bsIdentity :: !BakerIdentity,
        _bsInfo :: !FullBakerInfo,
        _bsContext :: !(SkovContext BakerConfig),
        _bsState :: !(SkovState BakerConfig)
      }
  
  -- | Typeclass of a datastructure that collects events.
  --  The instance of this structure determines how events
  --  are ordered.
  class Events e where
      -- | Add an event to the collection.
      addEvent :: PEvent -> e -> e
  
      -- | Filter the collection.
      filterEvents :: (PEvent -> Bool) -> e -> e
  
      -- | Make a collection from a list (in a default manner).
      makeEvents :: [PEvent] -> e
  
      -- | Extract the next event from the collection.
      nextEvent :: e -> Maybe (PEvent, e)
  
  instance Events (MinPQ.MinQueue PEvent) where
      addEvent = MinPQ.insert
      filterEvents = MinPQ.filter
      makeEvents = MinPQ.fromList
      nextEvent = MinPQ.minView
  
  -- | An instance of 'Events' that selects events randomly
  --  (without regard for time).  The randomness is determined
  --  by the 'StdGen', which is used to pick the next element
  --  from the sequence.
  data RandomisedEvents = RandomisedEvents
      { _reEvents :: !(Seq.Seq PEvent),
        _reGen :: !StdGen
      }
  
  -- | State of the simulation.
  data SimState = SimState
      { _ssBakers :: !(Vec.Vector BakerState),
        _ssEvents :: !(MinPQ.MinQueue PEvent),
        --  _ssEvents :: RandomisedEvents,
        _ssNextTimer :: !DummyTimer,
        _ssCurrentTime :: !UTCTime
      }
  
  makeLenses ''RandomisedEvents
  makeLenses ''BakerState
  makeLenses ''SimState
  
  -- | Pick an element from a sequence, returning the element
  --  and the sequence with that element removed.
  selectFromSeq :: (RandomGen g) => g -> Seq.Seq a -> (a, Seq.Seq a, g)
  selectFromSeq g s =
      let (n, g') = randomR (0, length s - 1) g
      in  (Seq.index s n, Seq.deleteAt n s, g')
  
  -- | Seed to use for randomness in 'RandomisedEvents'.
  --  This can be varied to generate different runs (when 'RandomisedEvents'
  --  is used.)
  randomisedEventsSeed :: Int
  randomisedEventsSeed = 0
  
  instance Events RandomisedEvents where
      addEvent e = reEvents %~ (Seq.|> e)
      filterEvents f = reEvents %~ Seq.filter f
      makeEvents l = RandomisedEvents (Seq.fromList l) (mkStdGen randomisedEventsSeed)
      nextEvent RandomisedEvents{..} = case _reEvents of
          Seq.Empty -> Nothing
          _ ->
              let (v, s', g') = selectFromSeq _reGen _reEvents
              in  Just (v, RandomisedEvents s' g')
  
  -- | Maximal baker ID.
  maxBakerId :: (Integral a) => a
  maxBakerId = 0 -- 9
  
  -- | List of all baker IDs.
  allBakers :: (Integral a) => [a]
  allBakers = [0 .. maxBakerId]
  
  transactions :: StdGen -> [(Integer, BlockItem)]
  transactions gen = trs (1 :: Nonce) (randoms gen :: [Word8])
    where
      trs (Nonce n) (amnt : amnts) = (toInteger n `div` 100, Dummy.makeTransferTransaction (Dummy.mateuszKP, Dummy.mateuszAccount) Dummy.mateuszAccount (fromIntegral amnt) (Nonce n)) : trs (Nonce (n + 1)) amnts
      trs _ _ = error "Ran out of transaction data"
  
  -- | Transactions from the extra accounts.
  extraAccountTransactions :: Int -> [(Integer, BlockItem)]
  extraAccountTransactions numAccts = trs 1
    where
      maxAcc = numAccts
      trs (Nonce n) =
          [ ( toInteger maxAcc * toInteger (n - 1) + toInteger i,
              Dummy.makeTransferTransaction
                  (Dummy.alesKP, Dummy.accountAddressFrom i)
                  Dummy.mateuszAccount
                  1
                  (Nonce n) -- \$ fromInteger $ toInteger maxAcc * toInteger (n - 1) + toInteger i)
            )
          | i <- [1 .. maxAcc]
          ]
              ++ trs (Nonce (n + 1))
  
  -- | Genesis accounts. For convenience, these all use the same keys.
  extraAccounts :: Int -> [GenesisAccount]
  extraAccounts numAccts = [Dummy.createCustomAccount 1000000 Dummy.alesKP (Dummy.accountAddressFrom i) | i <- [1 .. numAccts]]
  
  -- | Number of execution steps between blocks.
  --  Note: this can be used to control the number of transactions per block, in preference to
  --  having multiple transactions in a single execution step.
  ticksPerSlot :: (Num a) => a
  ticksPerSlot = 100
  
  -- | The initial state of the simulation.
  initialState :: Int -> IO SimState
  initialState numAccts = do
      -- This timestamp is only used for naming the database files.
      now <- currentTimestamp
      -- Change the following line to write the genesis to a file, if desired.
      when False $
          LBS.writeFile ("data/genesis-" ++ show now ++ ".dat") $
              S.runPutLazy (putPVGenesisData (PVGenesisData genData))
      _ssBakers <- Vec.fromList <$> mapM (mkBakerState now) (zip [0 ..] bakers)
      return SimState{..}
    where
      chainParams =
          Dummy.dummyChainParameters
              { _cpConsensusParameters = ConsensusParametersV0{_cpElectionDifficulty = makeElectionDifficulty 50000},
                _cpExchangeRates = makeExchangeRates 1 1,
                _cpFoundationAccount = maxBakerId + 1
              }
      -- The genesis parameters could be changed.
      -- The slot duration is set to 'ticksPerSlot' seconds, since the deterministic time
      -- advances 1 second per tick and baking is set to occur once every 'ticksPerSlot' ticks.
      (genData, bakers, _) =
          makeGenesisDataV0
              0 -- Start at time 0, to match time
              (maxBakerId + 1) -- Number of bakers
              (ticksPerSlot * 1000) -- Slot time is 100 seconds, for baking blocks every 100 ticks
              defaultFinalizationParameters
              Dummy.dummyCryptographicParameters
              dummyIdentityProviders
              dummyArs
              ([Dummy.createCustomAccount 1000000000000 Dummy.mateuszKP Dummy.mateuszAccount] ++ extraAccounts numAccts)
              (Energy maxBound)
              dummyKeyCollection
              chainParams
      mkBakerState :: Timestamp -> (BakerId, (BakerIdentity, FullBakerInfo)) -> IO BakerState
      mkBakerState now (bakerId, (_bsIdentity, _bsInfo)) = do
          createDirectoryIfMissing True "data"
          gsconfig <-
              makeGlobalStateConfig
                  defaultRuntimeParameters{rpAccountsCacheSize = 5000}
                  ("data/treestate-" ++ show now ++ "-" ++ show bakerId)
                  ("data/blockstate-" ++ show now ++ "-" ++ show bakerId ++ ".dat")
                  ("data/accountmap-" ++ show now ++ "-" ++ show bakerId ++ ".dat")
          let
              finconfig = ActiveFinalization (FinalizationInstance (bakerSignKey _bsIdentity) (bakerElectionKey _bsIdentity) (bakerAggregationKey _bsIdentity))
              hconfig = NoHandler
              config = SkovConfig gsconfig finconfig hconfig
          (_bsContext, _bsState) <- runLoggerT (initialiseSkov genData config) (logFor (fromIntegral bakerId))
          return BakerState{..}
      _ssEvents = makeEvents $ (PEvent 0 (TransactionEvent (extraAccountTransactions numAccts))) : [PEvent 0 (BakerEvent i (EBake 0)) | i <- allBakers]
      _ssNextTimer = 0
      _ssCurrentTime = posixSecondsToUTCTime 0
  
  -- | Log an event for a particular baker.
  logFor :: (MonadIO m) => Int -> LogMethod m
  logFor _ _ _ _ = return ()
  
  -- logFor i src lvl msg = liftIO $ do
  --     putStrLn $ "[" ++ show i ++ ":" ++ show src ++ ":" ++ show lvl ++ "] " ++ show msg
  --     hFlush stdout
  
  -- | Run a baker action in the state monad.
  runBaker :: Integer -> Int -> BakerM a -> StateT SimState IO a
  runBaker curTime bid a = do
      bakerState <- (Vec.! bid) <$> use ssBakers
      (r, bsState') <- runBase (runSkovT a handlers (_bsContext bakerState) (_bsState bakerState))
      ssBakers . ix bid . bsState .= bsState'
      return r
    where
      runBase z = do
          s <- get
          (r, s') <- liftIO $ runDeterministic (runLoggerT (runStateT z s) (logFor bid)) curTimeUTC
          put s'
          return r
      curTimeUTC = posixSecondsToUTCTime (fromIntegral curTime)
      handlers = SkovHandlers{..}
      shBroadcastFinalizationMessage = broadcastEvent curTime . EFinalization
      shOnTimeout timeout action = do
          tmr <- ssNextTimer <<%= (1 +)
          let t = case timeout of
                  DelayFor delay -> curTime + round delay
                  DelayUntil z -> curTime + max 1 (round (diffUTCTime z curTimeUTC))
          ssEvents %= addEvent (PEvent t (BakerEvent bid (ETimer tmr (void action))))
          return tmr
      shCancelTimer tmr = ssEvents %= filterEvents f
        where
          f (PEvent _ (BakerEvent _ (ETimer tmr' _))) = tmr' /= tmr
          f _ = True
      shPendingLive = return ()
  
  -- | Send an event to all bakers.
  broadcastEvent ::
      (MonadState SimState m) =>
      Integer ->
      Event ->
      m ()
  broadcastEvent curTime ev = ssEvents %= \e -> foldr addEvent e [PEvent curTime (BakerEvent i ev) | i <- allBakers]
  
  -- | Display an event for a baker.
  displayBakerEvent :: (MonadIO m) => Int -> Event -> m ()
  displayBakerEvent i ev = liftIO $ putStrLn $ show i ++ "> " ++ show ev
  
  bpBlock :: TS.BlockPointerType BakerM -> Block PV
  -- bpBlock (PairBlockData (l, _)) = BS._bpBlock l
  bpBlock = BS._bpBlock
  
  -- | Run a step of the consensus. This takes the next event and
  --  executes that.
  stepConsensus :: StateT SimState IO ()
  stepConsensus =
      (nextEvent <$> use ssEvents) >>= \case
          Nothing -> return ()
          Just (nextEv, evs') -> do
              ssEvents .= evs'
              case nextEv of
                  (PEvent _ (TransactionEvent trs)) -> case trs of
                      [] -> return ()
                      ((te, ev) : r) -> do
                          ssEvents %= \e -> foldr addEvent e [PEvent te (BakerEvent bid (ETransaction ev)) | bid <- allBakers]
                          case r of
                              [] -> return ()
                              ((t', _) : _) -> ssEvents %= addEvent (PEvent t' (TransactionEvent r))
                  (PEvent t (BakerEvent i ev)) -> {- displayBakerEvent i ev >> -} case ev of
                      EBake sl -> do
                          bIdentity <- (^. bsIdentity) . (Vec.! i) <$> use ssBakers
                          let doBake =
                                  bakeForSlot bIdentity sl
                          mb <- runBaker t i doBake
                          forM_ (bpBlock <$> mb) $ \case
                              GenesisBlock{} -> return ()
                              NormalBlock b -> broadcastEvent t (EBlock b)
                          ssEvents %= addEvent (PEvent (t + ticksPerSlot) (BakerEvent i (EBake (sl + 1))))
                      EBlock bb -> do
                          let pb = makePendingBlock bb (posixSecondsToUTCTime (fromIntegral t))
                          _ <- runBaker t i (receiveExecuteBlock pb)
                          return ()
                      ETransaction tr -> do
                          _ <- runBaker t i (receiveTransaction tr)
                          return ()
                      EFinalization fm -> do
                          _ <- runBaker t i (finalizationReceiveMessage fm)
                          return ()
                      EFinalizationRecord fr -> do
                          _ <- runBaker t i (finalizationReceiveRecord False fr)
                          return ()
                      ETimer _ a -> runBaker t i a
  
  -- | Main runner. Simply runs consensus for a certain number of steps.
  main :: IO ()
  main = do
      putStr "Number of accounts: "
      hFlush stdout
      numAccts <- readLn
      putStrLn "Initialising"
      b numAccts (1000000 :: Int)
    where
      loop 0 _ = return ()
      loop n s = do
          s' <- execStateT stepConsensus s
          loop (n - 1) s'
      b numAccts steps = do
          !s0 <- initialState numAccts
          putStrLn "Initialisation complete; press Enter to start"
          performGC
          _ <- getLine
          putStrLn "Starting"
          loop steps s0
  