{-# LANGUAGE StandaloneDeriving, TupleSections #-}
module SybilLimit where

import Data.List ( foldl', minimumBy )
import Data.Maybe ( fromMaybe )
import Data.Ord ( comparing )
import Data.Tuple ( swap )
import Data.IntMap.Strict ( IntMap, (!) )
import qualified Data.IntMap.Strict as IntMap
import Data.Traversable ( sequenceA )
import Control.Applicative
-- import System.Random
import Data.Map.Strict ( Map )
import qualified Data.Map as Map

import Stochastic

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.
    , tailCounters :: IntMap Int
    -- ^ verificaiton counter by instance number.
    , routeCount :: Int
    -- ^ The r parameter of SybilLimit.
    , pendingChecks :: Map NodeId PendingSybilCheck
    }

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.
    }


-- | When 'sybPendingTails' is empty, the intersection condition is passed if and
-- only if 'sybVerifiedTails' is not empty.
data PendingSybilCheck = PendingSybilCheck
    { sybPendingTails :: Map (NodeId,NodeId) [Int]
    , sybVerifiedTails :: Map (NodeId,NodeId) [Int]
    }

data PeerMessage
    = ForwardRRMessage { rrRoute :: Int
                       , rrCount :: Int
                       , rrSuspect :: NodeId }
    | BackwardRRMessage { rrRoute :: Int
                        , rrCount :: Int
                        , rrTail :: (NodeId,NodeId) }
    | RegistrationQuery { rrSuspect :: NodeId
                        , rrRegisteredFrom :: NodeId }
    | RegistrationResponse { rrSuspect :: NodeId
                           , rrRegisteredFrom :: NodeId
                           , rrValidRegistration :: Bool }


data SybilCheck
    = SybilCheck { chkSuspect :: NodeId
                 , chkAccepted :: Bool }

data MessageReaction = MessageReaction
    { changedState :: ThisNode
    , outgoingMessages :: [(NodeId, PeerMessage)]
    , sybilChecks :: [SybilCheck]
    }

data SuspectCredentials
    = SuspectCredentials { suspectId :: NodeId
                         , suspectTails :: [(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)}


-- | Required:
--
-- (1) srcId must be a valid friend node.
--
-- (2) msg must be either ForwardRRMessage or BackwardRRMessage
forwardMessage :: ThisNode -> (FriendNode -> IntMap Int) -> (NodeId,PeerMessage) -> (NodeId, PeerMessage)
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 -> MessageReaction
reactToMessage w (srcId, msg@(RegistrationResponse {})) me =
    MessageReaction me' [] sybils
 where
    me0 = me { pendingChecks = vs' }
    vs' :: Map NodeId PendingSybilCheck
    vs' = Map.adjust adj (rrSuspect msg) $ pendingChecks me
    tail = (rrRegisteredFrom msg, srcId)
    adj p = p { sybPendingTails = Map.delete tail $ sybPendingTails p
              , sybVerifiedTails =
                    if rrValidRegistration msg then goodTail else sybVerifiedTails p
              }
        where goodTail = Map.insert tail indexes
                                    $ sybVerifiedTails p
              indexes = fromMaybe [] $ Map.lookup tail $ sybPendingTails p
    (me',sybils) =
            case Map.lookup (rrSuspect msg) $ pendingChecks me0 of
              Nothing -> (me0,[])
              Just icheck -> if Map.null (sybPendingTails icheck)
                                then bcheck icheck
                                else (me0, [])
        where bcheck icheck
                = if Map.null (sybVerifiedTails icheck)
                        then (me0, dishonest)
                        else balanceCheck
                where
                  dishonest = [SybilCheck (rrSuspect msg) False]
                  balanceCheck = second chk $ balanceCondition tail indexes me0
                    where chk = (:[]) . SybilCheck (rrSuspect msg)
                          second f (x,y) = (x,f y)
                          indexes = fromMaybe []
                                    . Map.lookup tail
                                    $ sybVerifiedTails icheck

reactToMessage w (srcId, msg@(RegistrationQuery {})) me =
    MessageReaction me [(srcId,resp)] []
 where
    resp = RegistrationResponse { rrSuspect = rrSuspect msg
                                , rrRegisteredFrom = rrRegisteredFrom msg
                                , rrValidRegistration = check }
    check = fromMaybe False $ do
        key <- Map.lookup (rrRegisteredFrom msg) $ friendById me
        return $ registeredTerminal (friends me ! key) == Just (rrSuspect msg)

reactToMessage w (srcId, msg) me =
    case Map.lookup srcId $ friendById me of
        Nothing -> MessageReaction me [] []
        Just srcNo -> 
            if rrCount msg == w
                then terminate srcNo
                else MessageReaction me [forwardMessage me next (srcId,msg)] []
 where
    (terminate,next) = case msg of
                        ForwardRRMessage {} -> ( terminateForward, routesTo )
                        BackwardRRMessage {} -> ( terminateBackward, routesFrom )

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

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

balanceCondition :: (NodeId,NodeId) -> [Int] -> ThisNode -> (ThisNode, Bool)
balanceCondition (ka,kb) routeNums me = (me',didPass)
 where
    me' = me { tailCounters = counters' }
     where
        counters' = if didPass then incremented else tailCounters me
        incremented = IntMap.adjust (+1) cmin $ tailCounters me

    didPass = fromIntegral (cmin + 1) <= b

    b = h * max (logBase 2 r) a :: Double
     where
        r = fromIntegral $ routeCount me
        a = fromIntegral ( 1 + sum (IntMap.elems $ tailCounters me) ) / r
        h = 4 -- h > 1 is some universal constant that is not too small (they used
              -- h = 4 in their experiments).
    cmin = minimumBy indexCompare routeNums
     where
        indexCompare ca cb = tieBreak $ comparing (tailCounters me !) ca cb
         where
            tieBreak EQ = compare ca cb
            tieBreak x = x

initiateSybilCheck :: SuspectCredentials -> ThisNode -> MessageReaction
initiateSybilCheck cred me = MessageReaction me' msgs []
 where
    me' = me { pendingChecks = Map.insert (suspectId cred) p
                                    $ pendingChecks me }
    msgs = map (\(a,b) -> (b, RegistrationQuery (suspectId cred) a))
              $ Map.keys tmap
    p = PendingSybilCheck { sybPendingTails = tmap
                          , sybVerifiedTails = Map.empty
                          }
    tmap = tmap0 `Map.intersection` smap
        where
            smap = Map.fromList $ map (,()) (suspectTails cred)
            tmap0 = foldl' build
                           Map.empty
                           (map swap $ IntMap.toList $ routeTails me)
            build mp (tl,i) = Map.alter insert tl mp
             where insert Nothing = Just [i]
                   insert (Just is) = Just (i:is)