{-# LANGUAGE RecordWildCards #-}

module Juno.Consensus.Commit
  (doCommit
  ,makeCommandResponse
  ,makeCommandResponse')
where

import Data.List
import Data.Maybe (fromMaybe)
import Control.Lens
import Control.Monad
import Data.AffineSpace ((.-.))
import Data.Int (Int64)
import Data.Thyme.Clock (UTCTime, microseconds)

import qualified Data.Sequence as Seq
import qualified Data.Map as Map

import Data.Foldable (toList)

import Juno.Types hiding (valid)
import Juno.Util.Util
import Juno.Runtime.Sender (sendResults)

-- THREAD: SERVER MAIN.
doCommit :: Monad m => Raft m ()
doCommit = do
  commitUpdate <- updateCommitIndex
  when commitUpdate applyLogEntries

applyLogEntries :: Monad m => Raft m ()
applyLogEntries = do
  la <- use lastApplied
  ci <- use commitIndex
  le <- use logEntries
  now <- join $ view (rs.getTimestamp)
  let leToApply = Seq.drop (fromIntegral $ la + 1) . takeEntries (ci + 1) $ le
  results <- mapM (applyCommand now . _leCommand) leToApply
  r <- use nodeRole
  lastApplied .= ci
  logMetric $ MetricAppliedIndex ci
  if not (null results)
    then if r == Leader
         then do
           debug $ "Applied and Responded to " ++ show (length results) ++ " CMD(s)"
           sendResults $! toList results
         else debug $ "Applied " ++ show (length results) ++ " CMD(s)"
    else debug "Applied log entries but did not send results?"

interval :: UTCTime -> UTCTime -> Int64
interval start end = view microseconds $ end .-. start

logApplyLatency :: Monad m => Command -> Raft m ()
logApplyLatency (Command _ _ _ provenance) = case provenance of
  NewMsg -> return ()
  ReceivedMsg _digest _orig mReceivedAt -> case mReceivedAt of
    Just (ReceivedAt arrived) -> do
      now <- join $ view (rs.getTimestamp)
      logMetric $ MetricApplyLatency $ fromIntegral $ interval arrived now
    Nothing -> return ()

applyCommand :: Monad m => UTCTime -> Command -> Raft m (NodeID, CommandResponse)
applyCommand tEnd cmd@Command{..} = do
  apply <- view (rs.applyLogEntry)
  logApplyLatency cmd
  result <- apply cmd
  updateCmdStatusMap cmd result tEnd -- shared with the API and to query state
  replayMap %= Map.insert (_cmdClientId, getCmdSigOrInvariantError "applyCommand" cmd) (Just result)
  (,) _cmdClientId <$> makeCommandResponse tEnd cmd result

updateCmdStatusMap :: Monad m => Command -> CommandResult -> UTCTime -> Raft m ()
updateCmdStatusMap cmd cmdResult tEnd = do
  let rid = _cmdRequestId cmd
  mvarMap <- view (rs.cmdStatusMap)
  updateMapFn <- view (rs.updateCmdMap)
  let lat = case _pTimeStamp $ _cmdProvenance cmd of
        Nothing -> 1 -- don't want a div by zero error downstream and this is for demo purposes
        Just (ReceivedAt tStart) -> interval tStart tEnd
  void $ updateMapFn mvarMap rid (CmdApplied cmdResult lat)

makeCommandResponse :: Monad m => UTCTime -> Command -> CommandResult -> Raft m CommandResponse
makeCommandResponse tEnd cmd result = do
  nid <- view (cfg.nodeId)
  mlid <- use currentLeader
  let lat = case _pTimeStamp $ _cmdProvenance cmd of
              Nothing -> 1 -- don't want a div by zero error downstream and this is for demo purposes
              Just (ReceivedAt tStart) -> interval tStart tEnd
  return $ makeCommandResponse' nid mlid cmd result lat

makeCommandResponse' :: NodeID -> Maybe NodeID -> Command -> CommandResult -> Int64 -> CommandResponse
makeCommandResponse' nid mlid Command{..} result lat = CommandResponse
             result
             (fromMaybe nid mlid)
             nid
             _cmdRequestId
             lat
             NewMsg

logCommitChange :: Monad m => LogIndex -> LogIndex -> Raft m ()
logCommitChange before after
  | after > before = do
      logMetric $ MetricCommitIndex after
      mLastTime <- use lastCommitTime
      now <- join $ view (rs.getTimestamp)
      case mLastTime of
        Nothing -> return ()
        Just lastTime ->
          let duration = interval lastTime now
              (LogIndex numCommits) = after - before
              period = fromIntegral duration / fromIntegral numCommits
          in logMetric $ MetricCommitPeriod period
      lastCommitTime ?= now
  | otherwise = return ()

updateCommitIndex :: Monad m => Raft m Bool
updateCommitIndex = do
  ci <- use commitIndex
  proof <- use commitProof
  qsize <- view quorumSize
  es <- use logEntries

  let maxLogIndex = maxIndex es

  let evidence = reverse $ sortOn _aerIndex $ Map.elems proof

  case checkCommitProof qsize es maxLogIndex evidence of
    Left 0 -> return False
    Left n -> if maxLogIndex > fromIntegral ci
              then do
                    debug $ "Not enough evidence to commit yet, need " ++ show (qsize - n) ++ " more"
                    return False
              else return False
    Right qci -> if qci > ci
                then do
                  commitIndex .= qci
                  logCommitChange ci qci
                  commitProof %= Map.filter (\a -> qci < _aerIndex a)
                  debug $ "Commit index is now: " ++ show qci
                  return True
                else return False

checkCommitProof :: Int -> Log LogEntry -> LogIndex -> [AppendEntriesResponse] -> Either Int LogIndex
checkCommitProof qsize les maxLogIdx evidence = go 0 evidence
  where
    go n        []               = Left n
    go n (ev : evs)
      | _aerIndex ev > maxLogIdx = go n evs
      | Just (_aerHash ev) == (_leHash <$> lookupEntry (_aerIndex ev) les)
                                 = if (n + 1) >= qsize then
                                     Right $ _aerIndex ev
                                   else
                                     go (n + 1) evs
      | otherwise                = go n evs
