{-# LANGUAGE StrictData #-}
{-# LANGUAGE ViewPatterns #-}

module Data.Raft.Node.Internal where

import Prelude as P
import qualified Data.Map.Strict as Map
import qualified Data.List as List
import qualified Data.Sequence as Seq
import qualified Data.Set as Set
import System.Random.MWC

import Data.Raft.Index
import Data.RaftPrelude
import Data.Raft.Node.Type
import Data.Raft.NodeIdentifier
import Data.Raft.Entry
import Data.Raft.Messages
import Data.Raft.Ballot
import Data.Raft.Term
import Data.Raft.Log

newNode :: (Default st) => Self id -> AverageSeconds -> IO (Node id st a)
newNode = newNodeWithPeers Set.empty

newNodeWithPeers :: (Default st) => Peers id -> Self id -> AverageSeconds -> IO (Node id st a)
newNodeWithPeers peers self asec = do
  variance <- withSystemRandom . asGenST $ \gen -> uniformR (0.9, 1.1) gen
  let timeout = (P.fromIntegral asec) * (variance :: Double)
  time <- getCurrentTime
  return $ Node self peers def def (P.realToFrac timeout) time def def

demoteToFollower :: Standing id -> Standing id
demoteToFollower f@Follower{} = f
demoteToFollower (Candidate _) = Follower def
demoteToFollower (Leader _) = Follower def

hasMajority :: Ballots id -> Bool
hasMajority ballots
  | foldl' countGranted 0 ballots > (Map.size ballots `div` 2) = True
  | otherwise = False
  where
    countGranted i (Just Granted) = i + 1
    countGranted i _              = i

defaultBallots :: (Ord id) => Peers id -> Ballots id
defaultBallots peers = Map.fromList $ zip (Set.toList peers) (repeat def)

defaultSubjects :: Node id st a -> Subjects id
defaultSubjects = undefined

-- Handle AppendEntries

appendEntries :: EntryAppention id a -> CurrentTime -> Node id st a -> (EntryReceipt id, Node id st a)
appendEntries append ctime node
  | (append ^. leaderTerm) < (node ^. term) = (receiptFail node, node)
  | (append ^. leaderTerm) == (node ^. term) && (node ^. isCandidate) = (receiptFail node, node)
  | otherwise =
    append (append ^. previousEntryReference) (append ^. eappendEntries) (node ^. log)
    >>= return . (commit (append ^. leaderCommit))
    >>= \newLog -> return $ node & log .~ newLog

receiptFail :: Node id st a -> EntryReceipt id
receiptFail node = EntryReceipt (node ^. nodeSelf) (node ^. term) False

-- Handle Apply

applyCommitted :: (Applicable st a) => Node id st a -> Node id st a
applyCommitted node = node & state .~ newSt
                           & log .~ newLog
  where
    (pruned, newLog) = prune (node ^. log)
    newSt = foldl' (\st' (Entry _ _ a) -> apply st' a) (node ^. state) pruned

-- Handle RequestVote

requestVote :: VoteRequest id -> CurrentTime -> Node id st a -> (Ballot, Node id st a)
requestVote req ctime node =
  if (req ^. previousEntryReference . index) < (node ^. log . lastEntryReference . index)
  then (Rejected, node)
  else let cterm = req ^. term
           nterm = node ^. term
       in case node ^. nodeStanding of
            Follower Undecided -> grantIf $ cterm >= nterm
            _                  -> grantIf $ cterm > nterm
  where
    grantIf True = (Granted, node & nodeStanding .~ Follower (Supporting (req ^. candidate))
                                  & nodeLastTick .~ ctime)
    grantIf False = (Rejected, node)

acceptBallot :: BallotEnvelope id -> CurrentTime -> Node id st a -> (NodeAction id a, Node id st a)
acceptBallot env ctime node@((^. nodeStanding) -> Candidate ballots)
  | (env ^. term) > (node ^. term) = (NoAction, node & standing .~ Follower def)
  | (env ^. ballot) == Granted =
    let newNode = node & standing . _Candidate %~ Map.insert (env ^. peer) Granted
    in if newNode ^. candidateHasMajorityVotes
       then (sendHeartbeats node, newNode & standing .~ Leader (defaultSubjects newNode))
       else (NoAction, newNode)
  | otherwise = (NoAction, node)
acceptBallot _ _ node = (NoAction, node)

-- Handle Tick

tick :: (Ord id) => CurrentTime -> Node id st a -> (NodeAction id a, Node id st a)
tick ctime node@((^. nodeStanding) -> Follower supporting) =
  if addUTCTime (node ^. nodeTimeout) (node ^. nodeLastTick) < ctime
  then let peers = node ^. nodePeers
       in (requestVoteFromAll peers, node & nodeStanding .~ Candidate (defaultBallots peers)
                                          & nodeTerm %~ succ)
  else (NoAction, node & nodeLastTick .~ ctime)
tick ctime node@((^. nodeStanding) -> Candidate ballots) =
  if addUTCTime (node ^. nodeTimeout) (node ^. nodeLastTick) < ctime
  then let peers = node ^. nodePeers
       in (requestVoteFromAll peers, node & nodeStanding .~ Candidate (defaultBallots peers)
                                          & nodeTerm %~ succ)
  else (NoAction, node & nodeLastTick .~ ctime)
tick ctime node@((^. nodeStanding) -> Leader subjects) =
  if addUTCTime (node ^. nodeTimeout) (node ^. nodeLastTick) < oneThirdTimeout
  then (sendHeartbeats node, node & nodeLastTick .~ ctime)
  else (NoAction, node)
  where
    oneThirdTimeout = (node ^. nodeTimeout) P.* 0.33

requestVoteFromAll :: Node id st a -> NodeAction id a
requestVoteFromAll node =
  RequestVotesFrom $ List.zip (Set.toList (node ^. nodePeers)) (repeat req)
  where
    req = VoteRequest (node ^. nodeSelf) (node ^. nodeTerm) (lastEntryReference (node ^. nodeLog))

sendHeartbeats :: Node id st a -> NodeAction id a
sendHeartbeats node =
  SendAppendEntries $ List.zip (node ^. nodePeers) (List.repeat req)
  where
    req = AppendEntriesRequest
      (node ^. nodeTerm)
      (node ^. nodeSelf)
      (node ^. nodeLog . commit)
      (node ^. log . lastEntryReference)
      Seq.empty

-- Client: Append Data

clientAppend :: (Ord id) => ClientAppendRequest id a -> Node id st a -> ([NodeAction id a], Node id st a)
clientAppend (ClientAppendRequest _ content) (Node standing st log)
  | standingIsLeader standing =
    let prev@(EntryReference index term) = lastEntryReference log
        newEntries = (Seq.singleton (Entry (succ index) term content))
        (Just log') = append prev newEntries log
        peerReq = AppendEntriesRequest
                  (standingTerm standing)
                  (standingSelf standing)
                  (logCommit log')
                  (lastEntryReference log')
                  newEntries
        hbeats = List.zip (Set.toList (Set.delete (standingSelf standing) (standingPeers standing))) (List.repeat peerReq)
        clientRes = SendClientResponse (ClientAppendAcknowledge)
    in ([SendAppendEntries hbeats, clientRes], Node standing st log')
  | otherwise = ([SendClientResponse (ClientAppendFailure (standingSelf standing))], Node standing st log)
