-- |
--   Copyright   :  (c) Sam Truzjan 2013-2014
--   License     :  BSD3
--   Maintainer  :  pxqr.sta@gmail.com
--   Stability   :  experimental
--   Portability :  portable
--
--   This module defines internal state of a node instance. You can
--   have multiple nodes per application but usually you don't have
--   to. Normally, you don't need to import this module, use
--   "Network.BitTorrent.DHT" instead.
--
{-# LANGUAGE CPP                        #-}
{-# LANGUAGE RecordWildCards            #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE TemplateHaskell            #-}
module Network.BitTorrent.DHT.Session
       ( -- * Options
         -- | Use @optFooBar def@ to get default 'Alpha' or 'K'.
         Alpha
       , K
       , Options (..)

         -- * Session
       , Node
       , options
       , tentativeNodeId
       , myNodeIdAccordingTo
       , routingInfo
       , routableAddress

         -- ** Initialization
       , LogFun
       , NodeHandler
       , newNode
       , closeNode

         -- * DHT
         -- | Use @asks options@ to get options passed to 'startNode'
         -- or @asks thisNodeId@ to get id of locally running node.
       , DHT
       , runDHT

         -- ** Tokens
       , grantToken
       , checkToken

         -- ** Routing table
       , getTable
       , getClosest
       , getSwarms

         -- ** Peer storage
       , insertPeer
       , getPeerList
       , insertTopic
       , deleteTopic

         -- ** Messaging
       , queryParallel
       ) where

import Prelude hiding (ioError)

import Control.Concurrent.STM
import Control.Concurrent.Async.Lifted
import Control.Exception.Lifted hiding (Handler)
import Control.Monad.Base
import Control.Monad.Logger
import Control.Monad.Reader
import Control.Monad.Trans.Control
import Control.Monad.Trans.Resource
import Data.ByteString
import Data.Conduit.Lazy
import Data.Default
import Data.Fixed
import Data.Hashable
import Data.List as L
import Data.Maybe
import Data.Set as S
import Data.Time
import Network (PortNumber)
import System.Random (randomIO)

import Data.Torrent as Torrent
import Network.KRPC as KRPC hiding (Options, def)
import qualified Network.KRPC as KRPC (def)
import Network.BitTorrent.Address
import Network.BitTorrent.DHT.ContactInfo (PeerStore)
import qualified Network.BitTorrent.DHT.ContactInfo as P
import Network.BitTorrent.DHT.Message
import Network.BitTorrent.DHT.Routing as R
import Network.BitTorrent.DHT.Token   as T

{-----------------------------------------------------------------------
-- Options
-----------------------------------------------------------------------}

-- | Node lookups can proceed asynchronously.
type Alpha = Int

-- NOTE: libtorrent uses 5, azureus uses 10
-- | The quantity of simultaneous lookups is typically three.
defaultAlpha :: Alpha
defaultAlpha = 3

-- TODO add replication loop

-- TODO do not insert infohash -> peeraddr if infohash is too far from
-- this node id
{-
data Order
  = NearFirst
  | FarFirst
  | Random

data Traversal
  = Greedy     -- ^ aggressive short-circuit traversal
  | Exhaustive -- ^
-}

-- | Original Kamelia DHT uses term /publish/ for data replication
-- process. BitTorrent DHT uses term /announce/ since the purpose of
-- the DHT is peer discovery. Later in documentation, we use terms
-- /publish/ and /announce/ interchangible.
data Options = Options
  { -- | The degree of parallelism in 'find_node' queries. More
    -- parallism lead to faster bootstrapping and lookup operations,
    -- but also increase resource usage.
    --
    --   Normally this parameter should not exceed 'optK'.
    optAlpha       :: {-# UNPACK #-} !Alpha

    -- | /K/ parameter - number of nodes to return in 'find_node'
    -- responses.
  , optK           :: {-# UNPACK #-} !K

    -- | Number of buckets to maintain. This parameter depends on
    -- amount of nodes in the DHT network.
  , optBucketCount :: {-# UNPACK #-} !BucketCount

    -- | RPC timeout.
  , optTimeout     :: !NominalDiffTime

    -- | /R/ parameter - how many target nodes the 'announce' query
    -- should affect.
    --
    --   A large replica set compensates for inconsistent routing and
    --   reduces the need to frequently republish data for
    --   persistence. This comes at an increased cost for
    --   'Network.BitTorrent.DHT.insert' in terms of time, nodes
    --   contacted, and storage.
  , optReplication :: {-# UNPACK #-} !NodeCount

    -- | How often this node should republish (or reannounce) its
    -- data.
    --
    -- Large replica set ('optReplication') should require
    -- smaller reannounce intervals ('optReannounce').
  , optReannounce  :: !NominalDiffTime

    -- | The time it takes for data to expire in the
    --   network. Publisher of the data should republish (or
    --   reannounce) data to keep it in the network.
    --
    --   The /data expired timeout/ should be more than 'optReannounce'
    --   interval.
  , optDataExpired :: !NominalDiffTime
  } deriving (Show, Eq)

-- | Optimal options for bittorrent client. For short-lifetime
-- utilities you most likely need to tune 'optAlpha' and
-- 'optBucketCount'.
instance Default Options where
  def = Options
    { optAlpha       = defaultAlpha
    , optK           = defaultK

      -- see Fig.2 from "BitTorrent Mainline DHT Measurement" paper.
    , optBucketCount = defaultBucketCount

      -- see Fig.4 from "Profiling a Million User DHT" paper.
    , optTimeout     = 5  -- seconds
    , optReplication = 20 -- nodes
    , optReannounce  = 15 * 60
    , optDataExpired = 60 * 60
    }

seconds :: NominalDiffTime -> Int
seconds dt = fromEnum (realToFrac dt :: Uni)

{-----------------------------------------------------------------------
-- Tokens policy
-----------------------------------------------------------------------}

data SessionTokens = SessionTokens
  { tokenMap    :: !TokenMap
  , lastUpdate  :: !UTCTime
  , maxInterval :: !NominalDiffTime
  }

nullSessionTokens :: IO SessionTokens
nullSessionTokens = SessionTokens
  <$> (tokens <$> liftIO randomIO)
  <*> liftIO getCurrentTime
  <*> pure defaultUpdateInterval

-- TODO invalidate *twice* if needed
invalidateTokens :: UTCTime -> SessionTokens -> SessionTokens
invalidateTokens curTime ts @ SessionTokens {..}
  | curTime `diffUTCTime` lastUpdate > maxInterval = SessionTokens
    { tokenMap    = update tokenMap
    , lastUpdate  = curTime
    , maxInterval = maxInterval
    }
  |                  otherwise                     = ts

{-----------------------------------------------------------------------
-- Session
-----------------------------------------------------------------------}

-- | A set of torrents this peer intends to share.
type AnnounceSet = Set (InfoHash, PortNumber)

-- | Logger function.
type LogFun = Loc -> LogSource -> LogLevel -> LogStr -> IO ()

-- | DHT session keep track state of /this/ node.
data Node ip = Node
  { -- | Session configuration;
    options       :: !Options

    -- | Pseudo-unique self-assigned session identifier. This value is
    -- constant during DHT session and (optionally) between sessions.
  , tentativeNodeId    :: !NodeId

  , resources     :: !InternalState
  , manager       :: !(Manager (DHT       ip     )) -- ^ RPC manager;
  , routingInfo   :: !(TVar    (Maybe (R.Info ip))) -- ^ search table;
  , contactInfo   :: !(TVar    (PeerStore ip     )) -- ^ published by other nodes;
  , announceInfo  :: !(TVar     AnnounceSet       ) -- ^ to publish by this node;
  , sessionTokens :: !(TVar     SessionTokens     ) -- ^ query session IDs.
  , loggerFun     :: !LogFun
  }

-- | DHT keep track current session and proper resource allocation for
-- safe multithreading.
newtype DHT ip a = DHT { unDHT :: ReaderT (Node ip) IO a }
  deriving ( Functor, Applicative, Monad, MonadIO
           , MonadBase IO, MonadReader (Node ip), MonadThrow
           )

#if MIN_VERSION_monad_control(1,0,0)
newtype DHTStM ip a = StM {
      unSt :: StM (ReaderT (Node ip) IO) a
    }
#endif

instance MonadBaseControl IO (DHT ip) where
#if MIN_VERSION_monad_control(1,0,0)
  type StM (DHT ip) a = DHTStM ip a
#else
  newtype StM (DHT ip) a = StM {
      unSt :: StM (ReaderT (Node ip) IO) a
    }
#endif
  liftBaseWith cc = DHT $ liftBaseWith $ \ cc' ->
      cc $ \ (DHT m) -> StM <$> cc' m
  {-# INLINE liftBaseWith #-}

  restoreM = DHT . restoreM . unSt
  {-# INLINE restoreM #-}

-- | Check is it is possible to run 'queryNode' or handle pending
-- query from remote node.
instance MonadActive (DHT ip) where
  monadActive = getManager >>= liftIO . isActive
  {-# INLINE monadActive #-}

-- | All allocated resources will be closed at 'closeNode'.
instance MonadResource (DHT ip) where
  liftResourceT m = do
    s <- asks resources
    liftIO $ runInternalState m s

instance MonadKRPC (DHT ip) (DHT ip) where
  getManager = asks manager

instance MonadLogger (DHT ip) where
  monadLoggerLog loc src lvl msg = do
    logger <- asks loggerFun
    liftIO $ logger loc src lvl (toLogStr msg)

type NodeHandler ip = Handler (DHT ip)

-- | Run DHT session. You /must/ properly close session using
-- 'closeNode' function, otherwise socket or other scarce resources may
-- leak.
newNode :: Address ip
          => [NodeHandler ip] -- ^ handlers to run on accepted queries;
          -> Options          -- ^ various dht options;
          -> NodeAddr ip      -- ^ node address to bind;
          -> LogFun           -- ^
          -> Maybe NodeId     -- ^ use this NodeId, if not given a new one is generated.
          -> IO (Node ip)     -- ^ a new DHT node running at given address.
newNode hs opts naddr logger mbid = do
    s <- createInternalState
    runInternalState initNode s
      `onException` closeInternalState s
  where
    rpcOpts  = KRPC.def { optQueryTimeout = seconds (optTimeout opts) }
    nodeAddr = toSockAddr naddr
    initNode = do
      s      <- getInternalState
      (_, m) <- allocate (newManager rpcOpts nodeAddr hs) closeManager
      liftIO $ do
        myId   <- maybe genNodeId return mbid
        node   <- Node opts myId s m
                <$> atomically (newTVar Nothing)
                <*> newTVarIO def
                <*> newTVarIO S.empty
                <*> (newTVarIO =<< nullSessionTokens)
                <*> pure logger
        runReaderT (unDHT KRPC.listen) node
        return node

-- | Some resources like listener thread may live for
-- some short period of time right after this DHT session closed.
closeNode :: Node ip -> IO ()
closeNode Node {..} = closeInternalState resources

-- | Run DHT operation on the given session.
runDHT :: Node ip -> DHT ip a -> IO a
runDHT node action = runReaderT (unDHT action) node
{-# INLINE runDHT #-}

{-----------------------------------------------------------------------
--  Routing
-----------------------------------------------------------------------}

-- /pick a random ID/ in the range of the bucket and perform a
-- find_nodes search on it.

{-----------------------------------------------------------------------
-- Tokens
-----------------------------------------------------------------------}

tryUpdateSecret :: DHT ip ()
tryUpdateSecret = do
  curTime <- liftIO getCurrentTime
  toks    <- asks sessionTokens
  liftIO $ atomically $ modifyTVar' toks (invalidateTokens curTime)

grantToken :: Hashable a => NodeAddr a -> DHT ip Token
grantToken addr = do
  tryUpdateSecret
  toks <- asks sessionTokens >>= liftIO . readTVarIO
  return $ T.lookup addr $ tokenMap toks

-- | Throws 'HandlerError' if the token is invalid or already
-- expired. See 'TokenMap' for details.
checkToken :: Hashable a => NodeAddr a -> Token -> DHT ip Bool
checkToken addr questionableToken = do
  tryUpdateSecret
  toks <- asks sessionTokens >>= liftIO . readTVarIO
  return $ T.member addr questionableToken (tokenMap toks)


{-----------------------------------------------------------------------
-- Routing table
-----------------------------------------------------------------------}

-- | This nodes externally routable address reported by remote peers.
routableAddress :: DHT ip (Maybe SockAddr)
routableAddress = do
    info <- asks routingInfo >>= liftIO . atomically . readTVar
    return $ myAddress <$> info

-- | The current NodeId that the given remote node should know us by.
myNodeIdAccordingTo :: NodeAddr ip -> DHT ip NodeId
myNodeIdAccordingTo _ = do
  info <- asks routingInfo >>= liftIO . atomically . readTVar
  fallback <- asks tentativeNodeId
  return $ maybe fallback myNodeId info

-- | Get current routing table. Normally you don't need to use this
-- function, but it can be usefull for debugging and profiling purposes.
getTable :: Eq ip => DHT ip (Table ip)
getTable = do
  Node { tentativeNodeId = myId
       , routingInfo = var
       , options = opts } <- ask
  let nil = nullTable myId (optBucketCount opts)
  liftIO (maybe nil R.myBuckets <$> atomically (readTVar var))

getSwarms :: Ord ip => DHT ip [ (InfoHash, Int, Maybe ByteString) ]
getSwarms = do
    store <- asks contactInfo >>= liftIO . atomically . readTVar
    return $ P.knownSwarms store

-- | Find a set of closest nodes from routing table of this node. (in
-- no particular order)
--
--   This operation used for 'find_nodes' query.
--
getClosest :: Eq ip => TableKey k => k -> DHT ip [NodeInfo ip]
getClosest node = do
  k <- asks (optK . options)
  kclosest k node <$> getTable

{-----------------------------------------------------------------------
-- Peer storage
-----------------------------------------------------------------------}

refreshContacts :: DHT ip ()
refreshContacts =
    -- TODO limit dht peer store in size (probably by removing oldest peers)
    return ()


-- | Insert peer to peer store. Used to handle announce requests.
insertPeer :: Ord ip => InfoHash -> Maybe ByteString -> PeerAddr ip -> DHT ip ()
insertPeer ih name addr = do
  refreshContacts
  var <- asks contactInfo
  liftIO $ atomically $ modifyTVar' var (P.insertPeer ih name addr)

-- | Get peer set for specific swarm.
lookupPeers :: Ord ip => InfoHash -> DHT ip [PeerAddr ip]
lookupPeers ih = do
  refreshContacts
  var <- asks contactInfo
  liftIO $ P.lookup ih <$> readTVarIO var

-- | Prepare result for 'get_peers' query.
--
--   This operation use 'getClosest' as failback so it may block.
--
getPeerList :: Ord ip => InfoHash -> DHT ip (PeerList ip)
getPeerList ih = do
  ps <- lookupPeers ih
  if L.null ps
    then Left <$> getClosest ih
    else return (Right ps)

insertTopic :: InfoHash -> PortNumber -> DHT ip ()
insertTopic ih p = do
  var <- asks announceInfo
  liftIO $ atomically $ modifyTVar' var (S.insert (ih, p))

deleteTopic :: InfoHash -> PortNumber -> DHT ip ()
deleteTopic ih p = do
  var <- asks announceInfo
  liftIO $ atomically $ modifyTVar' var (S.delete (ih, p))

{-----------------------------------------------------------------------
-- Messaging
-----------------------------------------------------------------------}

-- | Failed queries are ignored.
queryParallel :: [DHT ip a] -> DHT ip [a]
queryParallel queries = do
    -- TODO: use alpha
    -- alpha <- asks (optAlpha . options)
    cleanup <$> mapConcurrently try queries
  where
    cleanup :: [Either QueryFailure a] -> [a]
    cleanup = mapMaybe (either (const Nothing) Just)