{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE UnboxedSums #-}

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

= Overview

A 'Context' carries trace state (the current 'Span') and 'Baggage' across
API boundaries and through your application. You rarely interact with it
directly; the @inSpan@ functions and propagators manage it for you.

= When you need Context directly

* __Custom propagation__: extracting\/injecting context from non-HTTP transports
* __Manual span parenting__: creating a span as a child of a specific context
* __Baggage access__: reading or modifying baggage entries

= Quick example

@
import OpenTelemetry.Context
import OpenTelemetry.Context.ThreadLocal

-- Read the current span from thread-local context:
ctx <- getContext
case lookupSpan ctx of
  Just span -> addAttribute span "custom.key" (toAttribute "value")
  Nothing   -> pure ()

-- Attach baggage to the current context:
adjustContext (insertBaggage myBaggage)
@

= Thread-local context

In most applications, context is stored in a thread-local variable managed
by "OpenTelemetry.Context.ThreadLocal". The @inSpan@ functions automatically
push and pop spans from this thread-local context.

= Spec reference

<https://opentelemetry.io/docs/specs/otel/context/>
-}
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.Text (Text)
import qualified Data.Vault.Strict as V
import OpenTelemetry.Baggage (Baggage)
import OpenTelemetry.Context.Types
import OpenTelemetry.Internal.Trace.Types
import OpenTelemetry.Internal.UnpackedMaybe
import Unsafe.Coerce (unsafeCoerce)
import Prelude hiding (lookup)


-- | @since 0.0.1.0
newKey :: (MonadIO m) => Text -> m (Key a)
newKey :: forall (m :: * -> *) a. MonadIO m => Text -> m (Key a)
newKey Text
n = IO (Key a) -> m (Key a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Text -> Key a -> Key a
forall a. Text -> Key a -> Key a
Key Text
n (Key a -> Key a) -> IO (Key a) -> IO (Key a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (Key a)
forall a. IO (Key a)
V.newKey)


-- | @since 0.0.1.0
class HasContext s where
  contextL :: Lens' s Context


-- | @since 0.0.1.0
empty :: Context
empty :: Context
empty = UMaybe Any -> UMaybe Baggage -> Vault -> Context
Context UMaybe Any
forall a. UMaybe a
UNothing UMaybe Baggage
forall a. UMaybe a
UNothing Vault
V.empty
{-# INLINE empty #-}


-- | @since 0.0.1.0
lookup :: Key a -> Context -> Maybe a
lookup :: forall a. Key a -> Context -> Maybe a
lookup (Key Text
_ Key a
k) (Context UMaybe Any
_ UMaybe Baggage
_ Vault
v) = Key a -> Vault -> Maybe a
forall a. Key a -> Vault -> Maybe a
V.lookup Key a
k Vault
v
{-# INLINE lookup #-}


-- | @since 0.0.1.0
insert :: Key a -> a -> Context -> Context
insert :: forall a. Key a -> a -> Context -> Context
insert (Key Text
_ Key a
k) a
x (Context UMaybe Any
s UMaybe Baggage
b Vault
v) = UMaybe Any -> UMaybe Baggage -> Vault -> Context
Context UMaybe Any
s UMaybe Baggage
b (Key a -> a -> Vault -> Vault
forall a. Key a -> a -> Vault -> Vault
V.insert Key a
k a
x Vault
v)
{-# INLINE insert #-}


-- | @since 0.0.1.0
adjust :: (a -> a) -> Key a -> Context -> Context
adjust :: forall a. (a -> a) -> Key a -> Context -> Context
adjust a -> a
f (Key Text
_ Key a
k) (Context UMaybe Any
s UMaybe Baggage
b Vault
v) = UMaybe Any -> UMaybe Baggage -> Vault -> Context
Context UMaybe Any
s UMaybe Baggage
b ((a -> a) -> Key a -> Vault -> Vault
forall a. (a -> a) -> Key a -> Vault -> Vault
V.adjust a -> a
f Key a
k Vault
v)


-- | @since 0.0.1.0
delete :: Key a -> Context -> Context
delete :: forall a. Key a -> Context -> Context
delete (Key Text
_ Key a
k) (Context UMaybe Any
s UMaybe Baggage
b Vault
v) = UMaybe Any -> UMaybe Baggage -> Vault -> Context
Context UMaybe Any
s UMaybe Baggage
b (Key a -> Vault -> Vault
forall a. Key a -> Vault -> Vault
V.delete Key a
k Vault
v)


-- | @since 0.0.1.0
union :: Context -> Context -> Context
union :: Context -> Context -> Context
union = Context -> Context -> Context
forall a. Semigroup a => a -> a -> a
(<>)


-- Span operations: O(1) via dedicated unboxed slot, no vault/hash overhead

-- | @since 0.0.1.0
lookupSpan :: Context -> Maybe Span
lookupSpan :: Context -> Maybe Span
lookupSpan (Context UMaybe Any
s UMaybe Baggage
_ Vault
_) = case UMaybe Any
s of
  UMaybe Any
UNothing -> Maybe Span
forall a. Maybe a
Nothing
  UJust Any
x -> Span -> Maybe Span
forall a. a -> Maybe a
Just (Any -> Span
forall a b. a -> b
unsafeCoerce Any
x)
{-# INLINE lookupSpan #-}


-- | @since 0.0.1.0
insertSpan :: Span -> Context -> Context
insertSpan :: Span -> Context -> Context
insertSpan !Span
span (Context UMaybe Any
_ UMaybe Baggage
b Vault
v) = UMaybe Any -> UMaybe Baggage -> Vault -> Context
Context (Any -> UMaybe Any
forall a. a -> UMaybe a
UJust (Span -> Any
forall a b. a -> b
unsafeCoerce Span
span)) UMaybe Baggage
b Vault
v
{-# INLINE insertSpan #-}


-- | @since 0.0.1.0
removeSpan :: Context -> Context
removeSpan :: Context -> Context
removeSpan (Context UMaybe Any
_ UMaybe Baggage
b Vault
v) = UMaybe Any -> UMaybe Baggage -> Vault -> Context
Context UMaybe Any
forall a. UMaybe a
UNothing UMaybe Baggage
b Vault
v
{-# INLINE removeSpan #-}


-- Baggage operations: O(1) via dedicated unboxed slot

-- | @since 0.0.1.0
lookupBaggage :: Context -> Maybe Baggage
lookupBaggage :: Context -> Maybe Baggage
lookupBaggage (Context UMaybe Any
_ UMaybe Baggage
b Vault
_) = UMaybe Baggage -> Maybe Baggage
forall a. UMaybe a -> Maybe a
toBaseMaybe UMaybe Baggage
b
{-# INLINE lookupBaggage #-}


-- | @since 0.0.1.0
insertBaggage :: Baggage -> Context -> Context
insertBaggage :: Baggage -> Context -> Context
insertBaggage Baggage
bag (Context UMaybe Any
s UMaybe Baggage
_ Vault
v) = UMaybe Any -> UMaybe Baggage -> Vault -> Context
Context UMaybe Any
s (Baggage -> UMaybe Baggage
forall a. a -> UMaybe a
UJust Baggage
bag) Vault
v
{-# INLINE insertBaggage #-}


-- | @since 0.0.1.0
removeBaggage :: Context -> Context
removeBaggage :: Context -> Context
removeBaggage (Context UMaybe Any
s UMaybe Baggage
_ Vault
v) = UMaybe Any -> UMaybe Baggage -> Vault -> Context
Context UMaybe Any
s UMaybe Baggage
forall a. UMaybe a
UNothing Vault
v
{-# INLINE removeBaggage #-}