{-# LANGUAGE BangPatterns #-}
module Network.BitTorrent.DHT.ContactInfo
       ( PeerStore
       , Network.BitTorrent.DHT.ContactInfo.lookup
       , Network.BitTorrent.DHT.ContactInfo.freshPeers
       , Network.BitTorrent.DHT.ContactInfo.insertPeer
       , knownSwarms
       ) where

import Control.Applicative
import Data.Default
import Data.List as L
import Data.Maybe
import Data.HashMap.Strict as HM
import Data.Serialize
import Data.Wrapper.PSQ as PSQ
import Data.Time.Clock.POSIX
import Data.ByteString (ByteString)
import Data.Word
import Network.Socket (SockAddr(..))

import Data.Torrent
import Network.BitTorrent.Address

-- {-
-- import Data.HashMap.Strict as HM
--
-- import Data.Torrent.InfoHash
-- import Network.BitTorrent.Address
--
-- -- increase prefix when table is too large
-- -- decrease prefix when table is too small
-- -- filter outdated peers
--
-- {-----------------------------------------------------------------------
-- --  PeerSet
-- -----------------------------------------------------------------------}
--
-- type PeerSet a = [(PeerAddr a, NodeInfo a, Timestamp)]
--
-- -- compare PSQueue vs Ordered list
--
-- takeNewest :: PeerSet a -> [PeerAddr a]
-- takeNewest = undefined
--
-- dropOld :: Timestamp -> PeerSet a -> PeerSet a
-- dropOld = undefined
--
-- insert :: PeerAddr a -> Timestamp -> PeerSet a -> PeerSet a
-- insert = undefined
--
-- type Mask = Int
-- type Size = Int
-- type Timestamp = Int
--
-- {-----------------------------------------------------------------------
-- --  InfoHashMap
-- -----------------------------------------------------------------------}
--
-- -- compare handwritten prefix tree versus IntMap
--
-- data Tree a
--   = Nil
--   | Tip !InfoHash !(PeerSet a)
--   | Bin !InfoHash !Mask !Size !Timestamp (Tree a) (Tree a)
--
-- insertTree :: InfoHash -> a -> Tree a -> Tree a
-- insertTree = undefined
--
-- type Prio = Int
--
-- --shrink :: ContactInfo ip -> Int
-- shrink  Nil      = Nil
-- shrink (Tip _ _) = undefined
-- shrink (Bin _ _) = undefined
--
-- {-----------------------------------------------------------------------
-- -- InfoHashMap
-- -----------------------------------------------------------------------}
--
-- -- compare new design versus HashMap
--
-- data IntMap k p a
-- type ContactInfo = Map InfoHash Timestamp (Set (PeerAddr IP) Timestamp)
--
-- data ContactInfo ip = PeerStore
--   { maxSize    :: Int
--   , prefixSize :: Int
--   , thisNodeId :: NodeId
--
--   , count      :: Int  -- ^ Cached size of the 'peerSet'
--   , peerSet    :: HashMap InfoHash [PeerAddr ip]
--   }
--
-- size :: ContactInfo ip -> Int
-- size = undefined
--
-- prefixSize :: ContactInfo ip -> Int
-- prefixSize = undefined
--
-- lookup :: InfoHash -> ContactInfo ip -> [PeerAddr ip]
-- lookup = undefined
--
-- insert :: InfoHash -> PeerAddr ip -> ContactInfo ip -> ContactInfo ip
-- insert = undefined
--
-- -- | Limit in size.
-- prune :: NodeId -> Int -> ContactInfo ip -> ContactInfo ip
-- prune pref targetSize  Nil      = Nil
-- prune pref targetSize (Tip _ _) = undefined
--
-- -- | Remove expired entries.
-- splitGT :: Timestamp -> ContactInfo ip -> ContactInfo ip
-- splitGT = undefined
-- -}

-- | Storage used to keep track a set of known peers in client,
-- tracker or DHT sessions.
newtype PeerStore ip = PeerStore (HashMap InfoHash (SwarmData ip))

type Timestamp = POSIXTime

data SwarmData ip = SwarmData
    { peers :: !(PSQ (PeerAddr ip) Timestamp)
    , name :: !(Maybe ByteString)
    }


newtype SerializeAddress a = SerializeAddress { unserializeAddress :: a }

instance Address a => Serialize (SerializeAddress a) where
    get = SerializeAddress <$> do
        c <- get
        case (c::Word8) of
            0x34 -> do ip4 <- get
                       return $ fromJust $ fromAddr (ip4::IPv4)
            0x36 -> do ip6 <- get
                       return $ fromJust $ fromAddr (ip6::IPv6)
            _ -> return $ error "cannot deserialize non-IP SerializeAddress"
    put (SerializeAddress a)
        | Just ip4 <- fromAddr a
                    = put (0x34::Word8) >> put (ip4::IPv4)
        | Just ip6 <- fromAddr a
                    = put (0x36::Word8) >> put (ip6::IPv6)
        | otherwise = return $ error "cannot serialize non-IP SerializeAddress"


instance (Ord ip, Address ip) => Serialize (SwarmData ip) where
    get = flip SwarmData <$> get
                         <*> ( PSQ.fromList . L.map parseAddr <$> get )
     where
        parseAddr addr = (unserializeAddress <$> addr)
                          :-> 0

    put SwarmData{..} = do
        put name
        put $ L.map (\(addr :-> _) -> (SerializeAddress <$> addr))
                -- XXX: should we serialize the timestamp?
            $ PSQ.toList peers


knownSwarms :: PeerStore ip -> [ (InfoHash, Int, Maybe ByteString) ]
knownSwarms (PeerStore m) = L.map (\(ih,SwarmData q n) -> (ih,PSQ.size q,n)) $ HM.toList m

swarmSingleton :: Ord ip => PeerAddr ip -> SwarmData ip
swarmSingleton a = SwarmData
    { peers = PSQ.singleton a 0
    , name = Nothing }

swarmInsert :: Ord ip => SwarmData ip -> SwarmData ip -> SwarmData ip
swarmInsert old new = SwarmData
    { peers = L.foldl' (\q (a :-> t) -> PSQ.insertWith (\p _ -> p) a t q) (peers old) (PSQ.toList $ peers new)
    , name = name new <|> name old -- TODO: decodeUtf8' check
    }

isSwarmOccupied SwarmData{..} = not $ PSQ.null peers

-- | Empty store.
instance Default (PeerStore a) where
  def = PeerStore HM.empty
  {-# INLINE def #-}

-- | Monoid under union operation.
instance Ord a => Monoid (PeerStore a) where
  mempty  = def
  {-# INLINE mempty #-}

  mappend (PeerStore a) (PeerStore b) =
    PeerStore (HM.unionWith swarmInsert a b)
  {-# INLINE mappend #-}

-- | Can be used to store peers between invocations of the client
-- software.
instance (Ord a, Address a) => Serialize (PeerStore a) where
  get = PeerStore . HM.fromList <$> get
  put (PeerStore m) = put (L.filter (isSwarmOccupied . snd) $ HM.toList m)

-- | Returns all peers associated with a given info hash.
lookup :: Ord a => InfoHash -> PeerStore a -> [PeerAddr a]
lookup ih (PeerStore m) = maybe [] (PSQ.keys . peers) $ HM.lookup ih m

batchSize = 64

-- | Used in 'get_peers' DHT queries.
freshPeers :: Ord a => InfoHash -> Timestamp -> PeerStore a -> ([PeerAddr a], PeerStore a)
freshPeers ih tm (PeerStore m) = fromMaybe ([],PeerStore m) $ do
    swarm <- HM.lookup ih m
    let ps0 = take batchSize $ unfoldr (incomp minView) (peers swarm)
        peers' = case reverse ps0 of
                    (_,psq):_ -> psq
                    _         -> peers swarm
        ps = L.map (key . fst) ps0
        m' = HM.insert ih swarm { peers = L.foldl' (\q p -> PSQ.insert p tm q) peers' ps } m
    return $! (ps,PeerStore m')

incomp :: (x -> Maybe (r,x)) -> x -> Maybe ((r,x),x)
incomp !f !x = do
   (result,x') <- f x
   pure $! ( (result,x'), x' )

-- | Used in 'announce_peer' DHT queries.
insertPeer :: Ord a => InfoHash -> Maybe ByteString -> PeerAddr a -> PeerStore a -> PeerStore a
insertPeer !ih !name !a !(PeerStore m) = seq a' $ PeerStore (HM.insertWith swarmInsert ih a' m)
 where
    a' = SwarmData { peers = PSQ.singleton a 0
                   , name = name }