-- | This module defines the type of 'TraceParent' which can be parsed
module HotelCalifornia.Tracing.TraceParent
    ( spanContextFromEnvironment
    , baggageFromEnvironment
    , spanContextToEnvironment
    , setParentSpanFromEnvironment
    ) where

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.Baggage (Baggage)
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
import qualified OpenTelemetry.Propagator.W3CBaggage as W3CBaggage

-- | 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

-- | This function looks up the @BAGGAGE@ environment variable and returns a
-- @'Maybe' 'Baggage'@ constructed from that.
baggageFromEnvironment :: IO (Maybe Baggage)
baggageFromEnvironment :: IO (Maybe Baggage)
baggageFromEnvironment = do
    Maybe ByteString
mBaggageBytes <- String -> IO (Maybe ByteString)
lookupEnvBS String
"BAGGAGE"

    let mBaggage :: Maybe Baggage
mBaggage = do
            ByteString
baggageBytes <- Maybe ByteString
mBaggageBytes
            ByteString -> Maybe Baggage
W3CBaggage.decodeBaggage ByteString
baggageBytes

    Maybe Baggage -> IO (Maybe Baggage)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Baggage
mBaggage

lookupEnvBS :: String -> IO (Maybe BS.ByteString)
lookupEnvBS :: String -> IO (Maybe ByteString)
lookupEnvBS String
str = (String -> ByteString) -> Maybe String -> Maybe ByteString
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Text -> ByteString
TE.encodeUtf8 (Text -> ByteString) -> (String -> Text) -> String -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Text
Text.pack) (Maybe String -> Maybe ByteString)
-> IO (Maybe String) -> IO (Maybe ByteString)
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)
-- , ( "BAGGAGE", traceParent)
-- ]
-- @
spanContextToEnvironment :: Span -> IO [(String, String)]
spanContextToEnvironment :: Span -> IO [(String, String)]
spanContextToEnvironment Span
span_ = do
    (ByteString
traceParent, ByteString
traceState) <- Span -> IO (ByteString, ByteString)
encodeSpanContext Span
span_

    Context
context <- IO Context
forall (m :: * -> *). MonadIO m => m Context
getContext

    let baggageVariables :: [(String, String)]
baggageVariables =
            case Context -> Maybe Baggage
Ctxt.lookupBaggage Context
context of
                Just Baggage
baggage ->
                    [(String
"BAGGAGE", ByteString -> String
BS8.unpack (Baggage -> ByteString
W3CBaggage.encodeBaggage Baggage
baggage))]
                Maybe Baggage
Nothing ->
                    []

    [(String, String)] -> IO [(String, String)]
forall a. a -> IO a
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)
            ]
        [(String, String)] -> [(String, String)] -> [(String, String)]
forall a. Semigroup a => a -> a -> a
<>  [(String, String)]
baggageVariables
        )

-- | 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
    Maybe Baggage
mBaggage <- IO (Maybe Baggage)
baggageFromEnvironment

    let insertSpanContext :: Context -> Context
insertSpanContext =
            case Maybe SpanContext
mSpanContext of
                Maybe SpanContext
Nothing ->
                    Context -> Context
forall a. a -> a
id
                Just SpanContext
spanContext ->
                    Span -> Context -> Context
Ctxt.insertSpan (SpanContext -> Span
wrapSpanContext SpanContext
spanContext{ isRemote = True })

    let insertBaggage :: Context -> Context
insertBaggage =
            case Maybe Baggage
mBaggage of
                Maybe Baggage
Nothing      -> Context -> Context
forall a. a -> a
id
                Just Baggage
baggage -> Baggage -> Context -> Context
Ctxt.insertBaggage Baggage
baggage

    (Context -> Context) -> IO ()
forall (m :: * -> *). MonadIO m => (Context -> Context) -> m ()
adjustContext (Context -> Context
insertBaggage (Context -> Context) -> (Context -> Context) -> Context -> Context
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Context -> Context
insertSpanContext)