{-# LANGUAGE FlexibleContexts #-}

module Haskoin.Store.Manager
  ( StoreConfig (..),
    Store (..),
    withStore,
  )
where

import Control.Monad (forever, unless, when)
import Control.Monad.Logger (MonadLoggerIO)
import Control.Monad.Reader (ReaderT (ReaderT), runReaderT)
import Data.Serialize (decode)
import Data.Time.Clock (NominalDiffTime)
import Data.Word (Word32)
import Haskoin
  ( BlockHash (..),
    Inv (..),
    InvType (..),
    InvVector (..),
    Message (..),
    MessageCommand (..),
    Network,
    NetworkAddress (..),
    NotFound (..),
    Pong (..),
    Reject (..),
    TxHash (..),
    VarString (..),
    sockToHostAddress,
  )
import Haskoin.Node
  ( Chain,
    ChainEvent (..),
    HostPort,
    Node (..),
    NodeConfig (..),
    NodeEvent (..),
    PeerEvent (..),
    PeerManager,
    WithConnection,
    withNode,
  )
import Haskoin.Store.BlockStore
  ( BlockStore,
    BlockStoreConfig (..),
    blockStoreBlockSTM,
    blockStoreHeadSTM,
    blockStoreNotFoundSTM,
    blockStorePeerConnectSTM,
    blockStorePeerDisconnectSTM,
    blockStoreTxHashSTM,
    blockStoreTxSTM,
    withBlockStore,
  )
import Haskoin.Store.Cache
  ( CacheConfig (..),
    CacheWriter,
    cacheNewBlock,
    cacheNewTx,
    cacheWriter,
    connectRedis,
    newCacheMetrics,
  )
import Haskoin.Store.Common
  ( StoreEvent (..),
    createDataMetrics,
  )
import Haskoin.Store.Database.Reader
  ( DatabaseReader (..),
    DatabaseReaderT,
    withDatabaseReader,
  )
import NQE
  ( Inbox,
    Process (..),
    Publisher,
    publishSTM,
    receive,
    withProcess,
    withPublisher,
    withSubscription,
  )
import Network.Socket (SockAddr (..))
import qualified System.Metrics as Metrics (Store)
import UnliftIO
  ( MonadIO,
    MonadUnliftIO,
    STM,
    atomically,
    link,
    withAsync,
  )
import UnliftIO.Concurrent (threadDelay)

-- | Store mailboxes.
data Store = Store
  { storeManager :: !PeerManager,
    storeChain :: !Chain,
    storeBlock :: !BlockStore,
    storeDB :: !DatabaseReader,
    storeCache :: !(Maybe CacheConfig),
    storePublisher :: !(Publisher StoreEvent),
    storeNetwork :: !Network
  }

-- | Configuration for a 'Store'.
data StoreConfig = StoreConfig
  { -- | max peers to connect to
    storeConfMaxPeers :: !Int,
    -- | static set of peers to connect to
    storeConfInitPeers :: ![HostPort],
    -- | discover new peers
    storeConfDiscover :: !Bool,
    -- | RocksDB database path
    storeConfDB :: !FilePath,
    -- | network constants
    storeConfNetwork :: !Network,
    -- | Redis cache configuration
    storeConfCache :: !(Maybe String),
    -- | gap on extended public key with no transactions
    storeConfInitialGap :: !Word32,
    -- | gap for extended public keys
    storeConfGap :: !Word32,
    -- | cache xpubs with more than this many used addresses
    storeConfCacheMin :: !Int,
    -- | maximum number of keys in Redis cache
    storeConfMaxKeys :: !Integer,
    -- | do not index new mempool transactions
    storeConfNoMempool :: !Bool,
    -- | wipe mempool when starting
    storeConfWipeMempool :: !Bool,
    -- | sync mempool from peers
    storeConfSyncMempool :: !Bool,
    -- | disconnect peer if message not received for this many seconds
    storeConfPeerTimeout :: !NominalDiffTime,
    -- | disconnect peer if it has been connected this long
    storeConfPeerMaxLife :: !NominalDiffTime,
    -- | connect to peers using the function 'withConnection'
    storeConfConnect :: !(SockAddr -> WithConnection),
    -- | delay in microseconds to retry getting cache lock
    storeConfCacheRetryDelay :: !Int,
    -- | stats store
    storeConfStats :: !(Maybe Metrics.Store)
  }

withStore ::
  (MonadLoggerIO m, MonadUnliftIO m) =>
  StoreConfig ->
  (Store -> m a) ->
  m a
withStore cfg action =
  connectDB cfg $
    ReaderT $ \db ->
      withPublisher $ \pub ->
        withPublisher $ \node_pub ->
          withSubscription node_pub $ \node_sub ->
            withNode (nodeCfg cfg db node_pub) $ \node ->
              withCache cfg (nodeChain node) db pub $ \mcache ->
                withBlockStore (blockStoreCfg cfg node pub db) $ \b ->
                  withAsync (nodeForwarder b pub node_sub) $ \a1 ->
                    link a1
                      >> action
                        Store
                          { storeManager = nodeManager node,
                            storeChain = nodeChain node,
                            storeBlock = b,
                            storeDB = db,
                            storeCache = mcache,
                            storePublisher = pub,
                            storeNetwork = storeConfNetwork cfg
                          }

connectDB :: MonadUnliftIO m => StoreConfig -> DatabaseReaderT m a -> m a
connectDB cfg f = do
  stats <- mapM createDataMetrics (storeConfStats cfg)
  withDatabaseReader
    (storeConfNetwork cfg)
    (storeConfInitialGap cfg)
    (storeConfGap cfg)
    (storeConfDB cfg)
    stats
    f

blockStoreCfg ::
  StoreConfig ->
  Node ->
  Publisher StoreEvent ->
  DatabaseReader ->
  BlockStoreConfig
blockStoreCfg cfg node pub db =
  BlockStoreConfig
    { blockConfChain = nodeChain node,
      blockConfManager = nodeManager node,
      blockConfListener = pub,
      blockConfDB = db,
      blockConfNet = storeConfNetwork cfg,
      blockConfNoMempool = storeConfNoMempool cfg,
      blockConfWipeMempool = storeConfWipeMempool cfg,
      blockConfSyncMempool = storeConfSyncMempool cfg,
      blockConfPeerTimeout = storeConfPeerTimeout cfg,
      blockConfStats = storeConfStats cfg
    }

nodeCfg ::
  StoreConfig ->
  DatabaseReader ->
  Publisher NodeEvent ->
  NodeConfig
nodeCfg cfg db pub =
  NodeConfig
    { nodeConfMaxPeers = storeConfMaxPeers cfg,
      nodeConfDB = databaseHandle db,
      nodeConfColumnFamily = Nothing,
      nodeConfPeers = storeConfInitPeers cfg,
      nodeConfDiscover = storeConfDiscover cfg,
      nodeConfEvents = pub,
      nodeConfNetAddr =
        NetworkAddress
          0
          (sockToHostAddress (SockAddrInet 0 0)),
      nodeConfNet = storeConfNetwork cfg,
      nodeConfTimeout = storeConfPeerTimeout cfg,
      nodeConfPeerMaxLife = storeConfPeerMaxLife cfg,
      nodeConfConnect = storeConfConnect cfg
    }

withCache ::
  (MonadUnliftIO m, MonadLoggerIO m) =>
  StoreConfig ->
  Chain ->
  DatabaseReader ->
  Publisher StoreEvent ->
  (Maybe CacheConfig -> m a) ->
  m a
withCache cfg chain db pub action =
  case storeConfCache cfg of
    Nothing ->
      action Nothing
    Just redisurl ->
      mapM newCacheMetrics (storeConfStats cfg) >>= \metrics ->
        connectRedis redisurl >>= \conn ->
          withSubscription pub $ \evts ->
            let conf = c conn metrics
             in withProcess (f conf) $ \p ->
                  cacheWriterProcesses evts (getProcessMailbox p) $ do
                    action (Just conf)
  where
    f conf cwinbox = runReaderT (cacheWriter conf cwinbox) db
    c conn metrics =
      CacheConfig
        { cacheConn = conn,
          cacheMin = storeConfCacheMin cfg,
          cacheChain = chain,
          cacheMax = storeConfMaxKeys cfg,
          cacheRetryDelay = storeConfCacheRetryDelay cfg,
          cacheMetrics = metrics
        }

cacheWriterProcesses ::
  MonadUnliftIO m =>
  Inbox StoreEvent ->
  CacheWriter ->
  m a ->
  m a
cacheWriterProcesses evts cwm action =
  withAsync events $ \a1 -> link a1 >> action
  where
    events = cacheWriterEvents evts cwm

cacheWriterEvents :: MonadIO m => Inbox StoreEvent -> CacheWriter -> m ()
cacheWriterEvents evts cwm =
  forever $
    receive evts >>= \e ->
      e `cacheWriterDispatch` cwm

cacheWriterDispatch :: MonadIO m => StoreEvent -> CacheWriter -> m ()
cacheWriterDispatch (StoreBestBlock _) = cacheNewBlock
cacheWriterDispatch (StoreMempoolNew t) = cacheNewTx t
cacheWriterDispatch (StoreMempoolDelete t) = cacheNewTx t
cacheWriterDispatch _ = const (return ())

nodeForwarder ::
  MonadIO m =>
  BlockStore ->
  Publisher StoreEvent ->
  Inbox NodeEvent ->
  m ()
nodeForwarder b pub sub =
  forever $ receive sub >>= atomically . storeDispatch b pub

-- | Dispatcher of node events.
storeDispatch ::
  BlockStore ->
  Publisher StoreEvent ->
  NodeEvent ->
  STM ()
storeDispatch b pub (PeerEvent (PeerConnected p)) = do
  publishSTM (StorePeerConnected p) pub
  blockStorePeerConnectSTM p b
storeDispatch b pub (PeerEvent (PeerDisconnected p)) = do
  publishSTM (StorePeerDisconnected p) pub
  blockStorePeerDisconnectSTM p b
storeDispatch b _ (ChainEvent (ChainBestBlock bn)) =
  blockStoreHeadSTM bn b
storeDispatch _ _ (ChainEvent _) =
  return ()
storeDispatch _ pub (PeerMessage p (MPong (Pong n))) =
  publishSTM (StorePeerPong p n) pub
storeDispatch b _ (PeerMessage p (MBlock block)) =
  blockStoreBlockSTM p block b
storeDispatch b _ (PeerMessage p (MTx tx)) =
  blockStoreTxSTM p tx b
storeDispatch b _ (PeerMessage p (MNotFound (NotFound is))) = do
  let blocks =
        [ BlockHash h
          | InvVector t h <- is,
            t == InvBlock || t == InvWitnessBlock
        ]
  unless (null blocks) $ blockStoreNotFoundSTM p blocks b
storeDispatch b pub (PeerMessage p (MInv (Inv is))) = do
  let txs = [TxHash h | InvVector t h <- is, t == InvTx || t == InvWitnessTx]
  publishSTM (StoreTxAnnounce p txs) pub
  unless (null txs) $ blockStoreTxHashSTM p txs b
storeDispatch _ pub (PeerMessage p (MReject r)) =
  when (rejectMessage r == MCTx) $
    case decode (rejectData r) of
      Left _ -> return ()
      Right th ->
        let reject =
              StoreTxReject
                p
                th
                (rejectCode r)
                (getVarString (rejectReason r))
         in publishSTM reject pub
storeDispatch _ _ _ =
  return ()