{-# LANGUAGE StandaloneDeriving #-}
module SybilLimit where

import Data.List
import Data.IntMap.Strict ( IntMap, (!) )
import qualified Data.IntMap.Strict as IntMap
import Data.Traversable
import Control.Applicative
-- import System.Random
import Stochastic
import Data.Map.Strict ( Map )
import qualified Data.Map as Map

data NodeId

deriving instance Ord NodeId
deriving instance Eq NodeId

data ThisNode = ThisNode
    { selfId :: NodeId
    , friends :: IntMap FriendNode
    , friendById :: Map NodeId Int
    , routeTails :: IntMap (NodeId,NodeId)
    -- ^ Terminal edge by instance number.
    , routeCount :: Int
    }

data FriendNode = FriendNode
    { friendId :: NodeId
    , routesTo :: IntMap Int
    -- ^ Forward random-route hop by instance number.
    , routesFrom :: IntMap Int
    -- ^ Backward random-route hop by instance number.
    , registeredTerminal :: Maybe NodeId
    -- ^ Currently registered terminal edge.
    -- The NodeId indicates the source vertex, friendId is the destination vertex.
    }


data PeerMessage
    = ForwardRRMessage { rrRoute :: Int
                       , rrCount :: Int
                       , rrSuspect :: NodeId }
    | BackwardRRMessage { rrRoute :: Int
                        , rrCount :: Int
                        , rrTail :: (NodeId,NodeId) }



friendNode :: NodeId -> FriendNode
friendNode nid = FriendNode nid IntMap.empty IntMap.empty Nothing

friendCount :: ThisNode -> Int
friendCount me = IntMap.size $ friends me

todo :: a
todo = error "unimplemented"

-- | @ addFriend @ adds a friend and updates the random route instances.
--
-- Arguments:
--
--   [@them@] The 'NodeId' of the friend being added.
--
--   [@me@]   'ThisNode' prior to the friend being added.
--
-- Returns the altered 'ThisNode' and the list of friend numbers
-- for which we should send notifications about the change.
--
addFriend :: NodeId -> ThisNode -> Stochastic ( ThisNode, [Int] )
addFriend them me = fmap (\ks -> (addFriend' ks them me,ks)) mks
 where
    d   = friendCount me
    mks = sequenceA $ replicate (routeCount me) $ random 0 d

    addFriend' :: [Int] -> NodeId -> ThisNode -> ThisNode
    addFriend' ks them me = me { friends = fs, friendById = ids }
     where
        d   = friendCount me
        fs0 = IntMap.insert d (friendNode them) $ friends me
        ids = Map.insert them d (friendById me)
        fs  = foldl' (flip $ uncurry updateRoute) fs0 edges
            where edges = zip [0..] (map mkEdge ks)
                  mkEdge k = if k==d then (d,d) else (k,d)

-- | @ updateRoute @ updates one of the random routes used by SybilLimit.
--
--  [@i@]     The random route instance number.
--
--  [@(k,j)@] A pair of friend numbers.  The routes inbound from friend @k@
--            will be routed out to friend @j@ for the s-direction (forwards)
--            and vice versa for the v-direction (backwards).
--
--  [@fs@]    The current friend map prior to changed route.
--
-- Returns:  The friend map with changes applied.
--
updateRoute :: Int -> (Int,Int) -> IntMap FriendNode -> IntMap FriendNode
updateRoute i (k,j) fs = jfromk $ ktoj fs
 where
    ktoj fs   = IntMap.insert k (setTo   i j (fs!k)) fs
    jfromk fs = IntMap.insert j (setFrom i k (fs!j)) fs

    setTo i x f   = f { routesTo   = IntMap.insert i x (routesTo f)}
    setFrom i x f = f { routesFrom = IntMap.insert i x (routesFrom f)}


forwardMessage me next (srcId,msg) = msg'
      where msg' = ( dest, msg { rrCount = rrCount msg + 1 } )
            srcno = friendById me Map.! srcId
            destno = next (friends me ! srcno) ! rrRoute msg
            dest = friendId (friends me ! destno)

reactToMessage :: Int -> (NodeId,PeerMessage) -> ThisNode -> (ThisNode, [(NodeId,PeerMessage)])
reactToMessage w (srcId, msg) me =
    if rrCount msg == w then terminate else (me, [forwardMessage me next (srcId,msg)])
 where
    (terminate,next) = case msg of
                        ForwardRRMessage {} -> ( terminateForward, routesTo )
                        BackwardRRMessage {} -> ( terminateBackward, routesFrom )

    terminateForward = ( me { friends = friends' }, [pong] )
      where pong =( srcId, BackwardRRMessage { rrCount = 1
                                             , rrRoute = rrRoute msg
                                             , rrTail = (srcId, selfId me)
                                             } )
            friends' = IntMap.adjust adj suspectNo $ friends me
            suspectNo = friendById me Map.! rrSuspect msg
            adj f = f { registeredTerminal = Just $ rrSuspect msg }

    terminateBackward = ( me', [])
      where me' = me { routeTails =
                        IntMap.insert (rrRoute msg) (rrTail msg)
                                     $ routeTails me
                     }