module Utils.Failover
  ( Failover, failoverLabels
  , mkFailover, mkFailoverLabeled
  , FailoverMode(..)
  , FailoverException(..)
  , withFailover, withFailoverReTest
  , FailoverMetrics, registerFailoverMetrics
  , mapFailover
  ) where

import ClassyPrelude hiding (try, Vector, finally, onException)

import Data.List.PointedList (PointedList)
import qualified Data.List.PointedList as P

import qualified Data.Set as Set
import qualified Data.Map as Map
import Data.Monoid (Endo(..))

import Numeric.Natural

import System.Clock

import Control.Lens hiding (failover)
import Utils.Lens.TH

import Control.Monad.Catch
import Control.Monad.Trans.Except (ExceptT, runExceptT)
import Control.Monad.State.Strict (runState)
import Control.Monad.Logger

import Control.Concurrent.STM.TVar (stateTVar)

import Data.Fixed

import Data.Time.Clock.POSIX
import Prometheus

import qualified Data.Foldable as F

import Data.Unique

import Utils (foldMapM)


data FailoverItem a = FailoverItem
  { failoverValue      :: a
  , failoverLabel      :: Text
  , failoverLastTest   :: Maybe TimeSpec
  , failoverReferences :: Set Unique
  }
makeLenses_ ''FailoverItem

newtype Failover a = Failover { failover :: TVar (PointedList (FailoverItem a)) }
  deriving (Eq)

data FailoverMode
  = FailoverUnlimited
  | FailoverLimited Natural
  | FailoverNone
  deriving (Eq, Ord, Read, Show, Generic, Typeable)

data FailoverException
  = FailoverNoItemsAvailable
  deriving (Eq, Ord, Read, Show, Generic, Typeable)
  deriving anyclass (Exception)

failoverLabels :: MonadIO m => Failover a -> m [Text]
failoverLabels Failover{..} = do
  items <- readTVarIO failover
  return $ items ^.. folded . _failoverLabel


mkFailover :: MonadIO m
           => PointedList a
           -> m (Failover a)
mkFailover opts = fmap Failover . liftIO $ newTVarIO opts'
  where opts' = flip (iover $ indexing traverse) opts $ \i failoverValue -> FailoverItem{ failoverLabel = tshow i, failoverLastTest = Nothing, failoverReferences = Set.empty, .. }

mkFailoverLabeled :: MonadIO m
                  => PointedList (Text, a)
                  -> m (Failover a)
mkFailoverLabeled opts = fmap Failover . liftIO $ newTVarIO opts'
  where opts' = flip map opts $ \(failoverLabel, failoverValue) -> FailoverItem{ failoverLastTest = Nothing, failoverReferences = Set.empty, .. }


mapFailover :: ( MonadIO m, Monoid b)
            => (a -> m b)
            -> Failover a
            -> m b
mapFailover f Failover{..} = do
  as <- toListOf (folded . _failoverValue) <$> readTVarIO failover
  foldMapM f as

withFailoverReference :: (MonadIO m, MonadMask m)
                      => Failover a
                      -> (Unique -> m b)
                      -> m b
withFailoverReference Failover{..} cont = do
  ref <- liftIO newUnique
  finally (cont ref) $
    atomically . modifyTVar failover $ traverse . _failoverReferences %~ Set.delete ref


withFailover :: ( MonadIO m, MonadMask m
                , Exception e
                , MonadLogger m
                )
             => Failover a
             -> FailoverMode
             -> (b -> ExceptT e m c)
             -> (a -> m b)
             -> m c
withFailover = withFailover' $ const P.focus

withFailoverReTest :: forall m a b c e.
                 ( MonadIO m, MonadMask m
                 , Exception e
                 , MonadLogger m
                 )
              => Failover a
              -> (Nano -> Bool)
              -> FailoverMode
              -> (b -> ExceptT e m c)
              -> (a -> m b)
              -> m c
withFailoverReTest f doTest = withFailover' testTargets f
  where
    testTargets :: TimeSpec -> Traversal' (PointedList (FailoverItem a)) (FailoverItem a)
    testTargets now = failing (P.prefix . traverse . filtered shouldTest) P.focus
      where shouldTest FailoverItem{failoverLastTest}
              = maybe True (\lT -> doTest . MkFixed . toNanoSecs $ now - lT) failoverLastTest

withFailover' :: forall m a b c e.
                 ( MonadIO m, MonadMask m
                 , Exception e
                 , MonadLogger m
                 )
              => (TimeSpec -> Traversal' (PointedList (FailoverItem a)) (FailoverItem a))
              -> Failover a
              -> FailoverMode
              -> (b -> ExceptT e m c)
              -> (a -> m b)
              -> m c
withFailover' testTarget' f@Failover{..} mode detAcceptable act = withFailoverReference f $ \alreadyTested ->
  let loop pErr mode' = withFailoverReference f $ \currentlyTesting -> do
        now <- liftIO $ getTime Monotonic

        FailoverItem{failoverLabel, failoverValue} <- maybe (maybe (throwM FailoverNoItemsAvailable) throwM pErr) return <=< atomically $
          stateTVar failover . runState $ do
            let testTarget :: Traversal' (PointedList (FailoverItem a)) (FailoverItem a)
                testTarget = taking 1 $ testTarget' now . filtered (views _failoverReferences $ Set.notMember alreadyTested)
            testTarget . _failoverReferences %= Set.insert currentlyTesting
            testTarget . _failoverLastTest %= Just . maybe now (max now)
            tgt <- preuse testTarget
            testTarget . _failoverReferences %= Set.insert alreadyTested
            return tgt

        let
          recordFailure = do
            $logErrorS "withFailover'" $ tshow (hashUnique alreadyTested) <> " recording failure for item " <> failoverLabel
            atomically . modifyTVar failover $ \failover' -> if
              | views (P.focus . _failoverReferences) (Set.member currentlyTesting) failover'
                -> fromMaybe (goFirst failover') $ P.next failover'
              | otherwise
                -> failover'
            where goFirst l = maybe l goFirst $ P.previous l

        $logDebugS "withFailover'" $ tshow (hashUnique alreadyTested) <> " using item " <> failoverLabel
        res' <- handleAll (\err -> $logErrorS "withFailover'" (tshow (hashUnique alreadyTested) <> " exception during act or detAcceptable: " <> tshow err) >> recordFailure >> throwM err) $
          runExceptT . detAcceptable =<< act failoverValue
        case res' of
          Right _  -> $logDebugS "withFailover'" $ tshow (hashUnique alreadyTested) <> " used item " <> failoverLabel <> " successfully"
          Left err -> $logErrorS "withFailover'" $ tshow (hashUnique alreadyTested) <> " used item " <> failoverLabel <> ", " <> tshow mode' <> ": " <> tshow err

        let
          recordSuccess = do
            $logDebugS "withFailover'" $ tshow (hashUnique alreadyTested) <> " recording success for item " <> failoverLabel
            didRecord <- atomically . stateTVar failover $ \failover' -> if
              | (newPrefix, newFocus : newSuffix) <- break (views _failoverReferences $ Set.member currentlyTesting) $ failover' ^. P.prefix
                -> ( True
                   , P.singleton newFocus
                       & P.prefix .~ newPrefix
                       & P.suffix .~ newSuffix ++ (view P.focus failover' : view P.suffix failover')
                   )
              | otherwise
                -> ( False, failover' )
            when didRecord $
              $logInfoS "withFailover'" $ tshow (hashUnique alreadyTested) <> " recorded success for item " <> failoverLabel
          doRetry err = do
            recordFailure
            case mode' of
              FailoverUnlimited -> loop (Just err) FailoverUnlimited
              FailoverLimited n | n > 0 -> loop (Just err) . FailoverLimited $ pred n
              _other -> throwM err

        case (res', mode') of
          (Left  err  , FailoverUnlimited)
            -> doRetry err
          (Left  err  , FailoverLimited n) | n > 0
            -> doRetry err
          (Left  err  , _)
            -> recordFailure >> throwM err
          (Right res'', _)
            -> res'' <$ recordSuccess
   in loop Nothing mode


data FailoverMetrics = forall a. FailoverMetrics
  { failoverMetricsFailover :: Failover a
  , failoverMetricsLabelsSeen :: TVar (Map Text (Maybe TimeSpec))
  }

failoverMetrics :: Text -> Failover a -> Metric FailoverMetrics
failoverMetrics fLbl failoverMetricsFailover@Failover{..} = Metric $ do
  failoverMetricsLabelsSeen <- atomically $ do
    foState <- readTVar failover
    newTVar . Map.fromList $ [ (failoverLabel, failoverLastTest) | FailoverItem{..} <- F.toList foState ]
  return (FailoverMetrics{..}, collectFailoverMetrics failoverMetricsLabelsSeen)
  where
    collectFailoverMetrics labelsSeen = do
      now <- liftIO $ getTime Monotonic
      cTime <- liftIO getPOSIXTime

      atomically $ do
        foState <- readTVar failover
        labelsSeen' <- stateTVar labelsSeen $ \labelsSeen' -> labelsSeen'
          & ala Endo foldMap [ Map.insert failoverLabel failoverLastTest | FailoverItem{..} <- F.toList foState ]
          & (id &&& id)
        return
          [ SampleGroup failoverActiveInfo GaugeType
            [ Sample "uni2work_failover_active_bool" [("failover", fLbl), ("label", lbl)] . bool "0" "1" $ views (P.focus . _failoverLabel) (== lbl) foState
            | (lbl, _) <- Map.toList labelsSeen'
            ]
          , SampleGroup failoverLastTestInfo GaugeType
            [ Sample "uni2work_failover_last_test_time" [("failover", fLbl), ("label", lbl)] . encodeUtf8 . tshow $ realToFrac cTime - (MkFixed . toNanoSecs $ now - lastTest :: Nano)
            | (lbl, Just lastTest) <- Map.toList labelsSeen'
            ]
          ]

    failoverActiveInfo = Info "uni2work_failover_active_bool"
                              "Currently active item in the failover set"
    failoverLastTestInfo = Info "uni2work_failover_last_test_time"
                                "Time of last attempt to use failover item"

registerFailoverMetrics :: MonadIO m => Text -> Failover a -> m ()
registerFailoverMetrics fLbl = void . register . failoverMetrics fLbl