module Utils.Metrics
  ( withHealthReportMetrics
  , registerGHCMetrics
  , observeHTTPRequestLatency
  , registerReadyMetric
  , withJobDuration
  ) where

import Import.NoFoundation hiding (Vector, Info)
import Prometheus
import Prometheus.Metric.GHC

import qualified Data.List as List

import Data.Time.Clock.POSIX

import Network.Wai (Middleware)
import qualified Network.Wai as Wai
import qualified Network.HTTP.Types as HTTP


histogramBuckets :: Rational -- ^ min
                 -> Rational -- ^ max
                 -> [Double]
histogramBuckets bMin bMax = map fromRational . takeWhile (<= bMax) . go bMin $ List.cycle factors
  where
    go n []     = [n]
    go n (f:fs) = n : go (f * n) fs

    factors
      | bMin' `mod` 2 == 0 = [2.5, 2, 2]
      | bMin' `mod` 5 == 0 = [2, 2, 2.5]
      | otherwise          = [2, 2.5, 2]
      where
        bMin' :: Integer
        bMin' = floor . List.head . dropWhile (< 1) $ List.iterate (* 10) bMin
  

{-# NOINLINE healthReportTime #-}
healthReportTime :: Vector Label2 Gauge
healthReportTime = unsafeRegister . vector ("check", "status") $ gauge info
  where
    info = Info "uni2work_health_check_time"
                "POSIXTime of last health check performed by this Uni2work-instance"

{-# NOINLINE healthReportDuration #-}
healthReportDuration :: Vector Label2 Histogram
healthReportDuration = unsafeRegister . vector ("check", "status") $ histogram info buckets
  where
    info = Info "uni2work_health_check_duration_seconds"
                "Duration of last health check performed by this Uni2work-instance"
    buckets = histogramBuckets 5e-6 100e-3

{-# NOINLINE httpRequestLatency #-}
httpRequestLatency :: Vector Label3 Histogram
httpRequestLatency = unsafeRegister . vector ("handler", "method", "status") $ histogram info buckets
  where info = Info "http_request_duration_seconds"
                    "HTTP request latency"
        buckets = histogramBuckets 50e-6 500

data ReadySince = MkReadySince

readyMetric :: POSIXTime -> Metric ReadySince
readyMetric ts = Metric $ return (MkReadySince, collectReadySince)
  where
    collectReadySince = return [SampleGroup info GaugeType [Sample "ready_time" [] sample]]
    info = Info "ready_time"
                "POSIXTime this Uni2work-instance became ready"
    sample = encodeUtf8 $ tshow (realToFrac ts :: Nano)

{-# NOINLINE jobDuration #-}
jobDuration :: Vector Label2 Histogram
jobDuration = unsafeRegister . vector ("kind", "status") $ histogram info buckets
  where info = Info "uni2work_job_duration_seconds"
                    "Duration of time taken to execute a background job"
        buckets = histogramBuckets 5e-6 500

{-# NOINLINE jobQueueDepth #-}
jobQueueDepth :: Vector Label2 Gauge
jobQueueDepth = unsafeRegister . vector ("worker", "priority") $ gauge info
  where info = Info "uni2work_job_queue_size_count"
                    "Current depth of worker queue"


withHealthReportMetrics :: MonadIO m => m HealthReport -> m HealthReport
withHealthReportMetrics act = do
  before <- liftIO getPOSIXTime
  report <- act
  after <- liftIO getPOSIXTime

  let checkVal = toPathPiece $ classifyHealthReport report
      statusVal = toPathPiece $ healthReportStatus report
      timeSample, durationSample :: Double
      timeSample = realToFrac (realToFrac after :: Nano)
      durationSample = realToFrac (realToFrac $ after - before :: Nano)
  liftIO $ withLabel healthReportTime (checkVal, statusVal) $ flip setGauge timeSample
  liftIO $ withLabel healthReportDuration (checkVal, statusVal) $ flip observe durationSample

  return report

registerGHCMetrics :: MonadIO m => m ()
registerGHCMetrics = void $ register ghcMetrics

observeHTTPRequestLatency :: forall site. ParseRoute site => (Route site -> String) -> Middleware
observeHTTPRequestLatency classifyHandler app req respond' = do
  start <- getPOSIXTime
  app req $ \res -> do
    end <- getPOSIXTime
    let method  = decodeUtf8 $ Wai.requestMethod req
        status  = tshow . HTTP.statusCode $ Wai.responseStatus res
        route :: Maybe (Route site)
        route   = parseRoute ( Wai.pathInfo req
                             , over (mapped . _2) (fromMaybe "") . HTTP.queryToQueryText $ Wai.queryString req
                             )
        handler' = pack . classifyHandler <$> route

        labels :: Label3
        labels = (fromMaybe "n/a" handler', method, status)
    withLabel httpRequestLatency labels . flip observe . realToFrac $ end - start

    respond' res

registerReadyMetric :: MonadIO m => m ()
registerReadyMetric = liftIO $ void . register . readyMetric =<< getPOSIXTime

withJobDuration :: (MonadIO m, MonadCatch m) => String -> m a -> m a
withJobDuration job doJob = do
  start <- liftIO getPOSIXTime
  res <- handleAll (return . Left) $ Right <$> doJob
  end <- liftIO getPOSIXTime

  liftIO . withLabel jobDuration (pack job, bool "failure" "success" $ is _Right res) . flip observe . realToFrac $ end - start

  case res of
    Left  exc  -> throwM exc
    Right res' -> return res'