{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE RecordWildCards #-}

module Juno.Consensus.Handle.AppendEntriesResponse
  (handle
  ,handleAlotOfAers
  ,updateCommitProofMap)
where

import Control.Lens hiding (Index)
import Control.Parallel.Strategies
import Control.Monad.Reader
import Control.Monad.State (get)
import Control.Monad.Writer.Strict
import Data.Maybe
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Set (Set)
import qualified Data.Set as Set

import Juno.Consensus.Commit (doCommit)
import Juno.Consensus.Handle.Types

import Juno.Runtime.Timer (resetElectionTimerLeader)
import Juno.Util.Util (debug, updateLNextIndex)
import qualified Juno.Types as JT

data AEResponseEnv = AEResponseEnv {
-- Old Constructors
    _nodeRole             :: Role
  , _term             :: Term
  , _commitProof      :: Map NodeID AppendEntriesResponse
  }
makeLenses ''AEResponseEnv

data AEResponseOut = AEResponseOut
  { _stateMergeCommitProof :: Map NodeID AppendEntriesResponse
  , _leaderState :: LeaderState }

data LeaderState =
  DoNothing |
  NotLeader |
  StatelessSendAE
    { _sendAENodeID :: NodeID } |
  Unconvinced -- sends AE after
    { _sendAENodeID :: NodeID
    , _deleteConvinced :: NodeID } |
  ConvincedAndUnsuccessful -- sends AE after
    { _sendAENodeID :: NodeID
    , _setLaggingLogIndex :: LogIndex } |
  ConvincedAndSuccessful -- does not send AE after
    { _incrementNextIndexNode :: NodeID
    , _incrementNextIndexLogIndex :: LogIndex
    , _insertConvinced :: NodeID}


data Convinced = Convinced | NotConvinced
data AESuccess = Success | Failure
data RequestTermStatus = OldRequestTerm | CurrentRequestTerm | NewerRequestTerm

handleAEResponse :: (MonadWriter [String] m, MonadReader AEResponseEnv m) => AppendEntriesResponse -> m AEResponseOut
handleAEResponse aer@AppendEntriesResponse{..} = do
    --tell ["got an appendEntriesResponse RPC"]
    mcp <- updateCommitProofMap aer <$> view commitProof
    role' <- view nodeRole
    currentTerm' <- view term
    if role' == Leader then
      return $ case (isConvinced, isSuccessful, whereIsTheRequest currentTerm') of
        (NotConvinced, _, OldRequestTerm) -> AEResponseOut mcp $ Unconvinced _aerNodeId _aerNodeId
        (NotConvinced, _, CurrentRequestTerm) -> AEResponseOut mcp $ Unconvinced _aerNodeId _aerNodeId
        (Convinced, Failure, CurrentRequestTerm) -> AEResponseOut mcp $ ConvincedAndUnsuccessful _aerNodeId _aerIndex
        (Convinced, Success, CurrentRequestTerm) -> AEResponseOut mcp $ ConvincedAndSuccessful _aerNodeId _aerIndex _aerNodeId
        -- The next two case are underspecified currently and they should not occur as
        -- they imply that a follow is ahead of us but the current code sends an AER anyway
        (NotConvinced, _, _) -> AEResponseOut mcp $ StatelessSendAE _aerNodeId
        (_, Failure, _) -> AEResponseOut mcp $ StatelessSendAE _aerNodeId
        -- This is just a fall through case in the current code
        -- do nothing as the follower is convinced and successful but out of date? Shouldn't this trigger a replay AE?
        (Convinced, Success, OldRequestTerm) -> AEResponseOut mcp DoNothing
        -- We are a leader, in a term that hasn't happened yet?
        (_, _, NewerRequestTerm) -> AEResponseOut mcp DoNothing
    else
      return $ AEResponseOut mcp NotLeader
  where
    isConvinced = if _aerConvinced then Convinced else NotConvinced
    isSuccessful = if _aerSuccess then Success else Failure
    whereIsTheRequest ct | _aerTerm == ct = CurrentRequestTerm
                         | _aerTerm < ct = OldRequestTerm
                         | otherwise = NewerRequestTerm

updateCommitProofMap :: AppendEntriesResponse -> Map NodeID AppendEntriesResponse -> Map NodeID AppendEntriesResponse
updateCommitProofMap aerNew m = Map.alter go nid m
  where
    nid :: NodeID
    nid = _aerNodeId aerNew
    go = \case
      Nothing   -> Just aerNew
      Just aerOld -> if _aerIndex aerNew > _aerIndex aerOld
                     then Just aerNew
                     else Just aerOld

handle :: Monad m => AppendEntriesResponse -> JT.Raft m ()
handle ae = do
  s <- get
  let ape = AEResponseEnv
              (JT._nodeRole s)
              (JT._term s)
              (JT._commitProof s)
  (AEResponseOut{..}, l) <- runReaderT (runWriterT (handleAEResponse ae)) ape
  mapM_ debug l
  JT.commitProof .= _stateMergeCommitProof
  doCommit
  case _leaderState of
    NotLeader -> return ()
    DoNothing -> resetElectionTimerLeader
    StatelessSendAE{..} ->
      resetElectionTimerLeader
    Unconvinced{..} -> do
      JT.lConvinced %= Set.delete _deleteConvinced
      resetElectionTimerLeader
    ConvincedAndSuccessful{..} -> do
      updateLNextIndex $ Map.insert _incrementNextIndexNode $ _incrementNextIndexLogIndex + 1
      JT.lConvinced %= Set.insert _insertConvinced
      resetElectionTimerLeader
    ConvincedAndUnsuccessful{..} -> do
      updateLNextIndex $ Map.insert _sendAENodeID _setLaggingLogIndex
      resetElectionTimerLeader

handleAlotOfAers :: Monad m => AlotOfAERs -> JT.Raft m ()
handleAlotOfAers (AlotOfAERs m) = do
  ks <- KeySet <$> view (JT.cfg . JT.publicKeys) <*> view (JT.cfg . JT.clientPublicKeys)
  let res = (processSetAer ks <$> Map.elems m) `using` parList rseq
  aers <- catMaybes <$> mapM (\(a,l) -> mapM_ debug l >> return a) res
  mapM_ handle aers

processSetAer :: KeySet -> Set AppendEntriesResponse -> (Maybe AppendEntriesResponse, [String])
processSetAer ks s = go [] (Set.toDescList s)
  where
    go fails [] = (Nothing, fails)
    go fails (aer:rest)
      | _aerWasVerified aer = (Just aer, fails)
      | otherwise = case aerReVerify ks aer of
                      Left f -> go (f:fails) rest
                      Right () -> (Just $ aer {_aerWasVerified = True}, fails)


-- | Verify if needed on `ReceivedMsg` provenance
aerReVerify :: KeySet -> AppendEntriesResponse -> Either String ()
aerReVerify  _ AppendEntriesResponse{ _aerWasVerified = True } = Right ()
aerReVerify  _ AppendEntriesResponse{ _aerProvenance = NewMsg } = Right ()
aerReVerify ks AppendEntriesResponse{
                  _aerWasVerified = False,
                  _aerProvenance = ReceivedMsg{..}
                } = verifySignedRPC ks $ SignedRPC _pDig _pOrig
