module Yam.Middleware.Trace(
  -- * Trace Middleware
    MonadTracer(..)
  , MonadTracing(..)
  , TraceConfig(..)
  , hTraceId
  , hParentTraceId
  , hSpanId
  , hSampled
  , traceMiddleware
  ) where

import qualified Data.HashMap.Lazy as HM
import           Data.Opentracing
import           Data.Salak
import qualified Data.Text         as T
import qualified Data.Vault.Lazy   as L
import           System.IO.Unsafe  (unsafePerformIO)
import           Yam.Logger
import           Yam.Middleware
import           Yam.Types

data TraceConfig = TraceConfig
  { enabled :: Bool
  , method  :: TraceNotifyType
  } deriving (Eq, Show)

data TraceNotifyType
  = NoTracer
  deriving (Eq, Show)

instance FromProperties TraceNotifyType where
  fromProperties = fromProperties >=> go
    where
      go :: Property -> Return TraceNotifyType
      go _ = return NoTracer

instance FromProperties TraceConfig where
  fromProperties p = TraceConfig
    <$> p .?> "enabled" .?= enabled def
    <*> p .?> "type"    .?= method  def

instance Default TraceConfig where
  def = TraceConfig True NoTracer

notifier :: TraceNotifyType -> Span -> App ()
notifier _ _ = return ()

{-# NOINLINE spanContextKey #-}
spanContextKey :: L.Key SpanContext
spanContextKey = unsafePerformIO newKey

{-# NOINLINE spanKey #-}
spanKey :: L.Key Span
spanKey = unsafePerformIO newKey

instance MonadTracer App where
  askSpanContext = requireAttr spanContextKey

instance MonadTracing App where
  runInSpan name notify action = do
    s <- askAttr spanKey
    n <- case s of
      Just sp -> newChildSpan name sp
      _       -> newSpan name
    notify n
    a <- withAttr spanKey n $ action n
    finishSpan n >>= notify
    return a

hTraceId :: HeaderName
hTraceId = "X-B3-TraceId"

hParentTraceId :: HeaderName
hParentTraceId = "X-B3-ParentSpanId"

hSpanId :: HeaderName
hSpanId = "X-B3-SpanId"

hSampled :: HeaderName
hSampled = "X-B3-Sampled"

parseSpan :: RequestHeaders -> Env -> IO Env
parseSpan headers env =
  let sc = fromMaybe (SpanContext "" HM.empty) $ getAttr spanContextKey env
  in case Prelude.lookup hTraceId headers of
      Just tid -> let sc' = sc { traceId = tid }
                  in return $ env
                      & setAttr spanContextKey      sc'
                      & go (fromMaybe (traceId sc') $ Prelude.lookup hSpanId headers) sc'
      _        -> do
        c <- newContext
        return $ setAttr spanContextKey c env
  where
    go spanId context env' =
      let name = "-"
          startTime  = undefined
          finishTime = Nothing
          tags       = HM.empty
          logs       = HM.empty
          references = []
      in setAttr spanKey Span{..} env'

traceMw :: Env -> (Span -> App ()) -> Middleware
traceMw env' notify app req resH = do
  env <- parseSpan (requestHeaders req) env'
  runApp env $
    runInSpan (decodeUtf8 (requestMethod req) <> " /" <> T.intercalate "/" (pathInfo req)) notify $ \s@Span{..} -> do
      let SpanContext{..} = context
          tid = decodeUtf8 $ traceId <> "," <> spanId
          v   = L.insert extensionLogKey tid (vault req)
          v'  = L.insert spanKey s v
          rh' = resH . mapResponseHeaders (\hs -> (hTraceId, traceId):(hSpanId, spanId):hs)
          c e = do
            runApp env { reqAttributes = Just v} (logError $ showText e)
            rh' $ whenException e
      liftIO (app req {vault = v'} rh' `catch` c)

traceMiddleware :: TraceConfig -> AppMiddleware
traceMiddleware TraceConfig{..}
  = AppMiddleware $ \env f -> if enabled
    then f (env, traceMw env $ notifier method)
    else f (env, id)