{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}

-----------------------------------------------------------------------------

-----------------------------------------------------------------------------

{- |
 Module      :  OpenTelemetry.Context
 Copyright   :  (c) Ian Duncan, 2021
 License     :  BSD-3
 Description :  Carrier for execution-scoped values across API boundaries
 Maintainer  :  Ian Duncan
 Stability   :  experimental
 Portability :  non-portable (GHC extensions)

 The ability to correlate events across service boundaries is one of the principle concepts behind distributed tracing. To find these correlations, components in a distributed system need to be able to collect, store, and transfer metadata referred to as context.

 A context will often have information identifying the current span and trace, and can contain arbitrary correlations as key-value pairs.

 Propagation is the means by which context is bundled and transferred in and across services, often via HTTP headers.

 Together, context and propagation represent the engine behind distributed tracing.
-}
module OpenTelemetry.Context (
  Key (keyName),
  newKey,
  Context,
  HasContext (..),
  empty,
  lookup,
  insert,
  -- , insertWith
  adjust,
  delete,
  union,
  insertSpan,
  lookupSpan,
  removeSpan,
  insertBaggage,
  lookupBaggage,
  removeBaggage,
) where

import Control.Monad.IO.Class
import Data.Maybe
import Data.Text (Text)
import qualified Data.Vault.Strict as V
import OpenTelemetry.Baggage (Baggage)
import OpenTelemetry.Context.Types
import OpenTelemetry.Internal.Trace.Types
import System.IO.Unsafe
import Prelude hiding (lookup)


newKey :: MonadIO m => Text -> m (Key a)
newKey n = liftIO (Key n <$> V.newKey)


class HasContext s where
  contextL :: Lens' s Context


empty :: Context
empty = Context V.empty


lookup :: Key a -> Context -> Maybe a
lookup (Key _ k) (Context v) = V.lookup k v


insert :: Key a -> a -> Context -> Context
insert (Key _ k) x (Context v) = Context $ V.insert k x v


-- insertWith
--   :: (a -> a -> a)
--   -- ^ new value -> old value -> result
--   -> Key a -> a -> Context -> Context
-- insertWith f (Key _ k) x (Context v) = Context $ case V.lookup k of
--   Nothing -> V.insert k x v
--   Just ox -> V.insert k (f x ox) v

adjust :: (a -> a) -> Key a -> Context -> Context
adjust f (Key _ k) (Context v) = Context $ V.adjust f k v


delete :: Key a -> Context -> Context
delete (Key _ k) (Context v) = Context $ V.delete k v


union :: Context -> Context -> Context
union (Context v1) (Context v2) = Context $ V.union v1 v2


spanKey :: Key Span
spanKey = unsafePerformIO $ newKey "span"
{-# NOINLINE spanKey #-}


lookupSpan :: Context -> Maybe Span
lookupSpan = lookup spanKey


insertSpan :: Span -> Context -> Context
insertSpan = insert spanKey


removeSpan :: Context -> Context
removeSpan = delete spanKey


baggageKey :: Key Baggage
baggageKey = unsafePerformIO $ newKey "baggage"
{-# NOINLINE baggageKey #-}


lookupBaggage :: Context -> Maybe Baggage
lookupBaggage = lookup baggageKey


insertBaggage :: Baggage -> Context -> Context
insertBaggage b c = case lookup baggageKey c of
  Nothing -> insert baggageKey b c
  Just b' -> insert baggageKey (b <> b') c


removeBaggage :: Context -> Context
removeBaggage = delete baggageKey