-- | This module defines the type of 'TraceParent' which can be parsed
module HotelCalifornia.Tracing.TraceParent where

import Data.Foldable (for_)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as BS8
import qualified Data.Text.Encoding as TE
import qualified Data.Text as Text
import OpenTelemetry.Propagator.W3CTraceContext
import OpenTelemetry.Trace.Core (SpanContext, isRemote, wrapSpanContext, Span)
import System.Environment
import OpenTelemetry.Context.ThreadLocal
import qualified OpenTelemetry.Context as Ctxt

-- | This function looks up the @TRACEPARENT@ and @TRACECONTEXT@ environment
-- variables and returns a @'Maybe' 'SpanContext'@ constructed from them.
spanContextFromEnvironment :: IO (Maybe SpanContext)
spanContextFromEnvironment :: IO (Maybe SpanContext)
spanContextFromEnvironment = do
    Maybe ByteString
mtraceParent <- String -> IO (Maybe ByteString)
lookupEnvBS String
"TRACEPARENT"
    Maybe ByteString
mtraceContext <- String -> IO (Maybe ByteString)
lookupEnvBS String
"TRACESTATE"
    pure $ Maybe ByteString -> Maybe ByteString -> Maybe SpanContext
decodeSpanContext Maybe ByteString
mtraceParent Maybe ByteString
mtraceContext
  where
    lookupEnvBS :: String -> IO (Maybe BS.ByteString)
    lookupEnvBS :: String -> IO (Maybe ByteString)
lookupEnvBS String
str = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Text -> ByteString
TE.encodeUtf8 forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Text
Text.pack) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> IO (Maybe String)
lookupEnv String
str

-- | This function takes the given 'Span' and converts it into a list of
-- environment variables consisting of:
--
-- @
-- [ ( "TRACEPARENT", traceParent)
-- , ( "TRACESTATE", traceState)
-- ]
-- @
spanContextToEnvironment :: Span -> IO [(String, String)]
spanContextToEnvironment :: Span -> IO [(String, String)]
spanContextToEnvironment Span
spanContext = do
    (ByteString
traceParent, ByteString
traceState) <- Span -> IO (ByteString, ByteString)
encodeSpanContext Span
spanContext
    forall (f :: * -> *) a. Applicative f => a -> f a
pure
        [ (String
"TRACEPARENT", ByteString -> String
BS8.unpack ByteString
traceParent)
        , (String
"TRACESTATE", ByteString -> String
BS8.unpack ByteString
traceState)
        ]


-- | This function should be called after you've initialized and attached the
-- thread local 'Context'.
setParentSpanFromEnvironment :: IO ()
setParentSpanFromEnvironment :: IO ()
setParentSpanFromEnvironment = do
    Maybe SpanContext
mspanContext <- IO (Maybe SpanContext)
spanContextFromEnvironment
    forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ Maybe SpanContext
mspanContext \SpanContext
spanContext -> do
        forall (m :: * -> *). MonadIO m => (Context -> Context) -> m ()
adjustContext forall a b. (a -> b) -> a -> b
$ Span -> Context -> Context
Ctxt.insertSpan (SpanContext -> Span
wrapSpanContext (SpanContext
spanContext {isRemote :: Bool
isRemote = Bool
True}))