content: {-# LANGUAGE ImportQualifiedPost #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE NamedFieldPuns #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TemplateHaskell #-}
  
  module OpenTelemetry.Instrumentation.Tasty (instrumentTestTree, instrumentTestTreeWithTracer) where
  
  import Control.Exception (bracket)
  import Data.Tagged (Tagged, retag)
  import Data.Text qualified as T
  import OpenTelemetry.Context (insertSpan, lookupSpan, removeSpan)
  import OpenTelemetry.Context.ThreadLocal (adjustContext, getContext)
  import OpenTelemetry.Trace.Core (Span, SpanStatus (Error, Ok), Tracer, addAttribute, createSpan, defaultSpanArguments, detectInstrumentationLibrary, endSpan, getGlobalTracerProvider, inSpan, makeTracer, setStatus, tracerOptions)
  import Test.Tasty (TestTree, withResource)
  import Test.Tasty.Options (OptionDescription)
  import Test.Tasty.Providers (IsTest (run, testOptions))
  import Test.Tasty.Runners (Outcome (Failure, Success), ResourceSpec (ResourceSpec), Result (Result, resultDescription, resultOutcome), TestTree (After, AskOptions, PlusTestOptions, SingleTest, TestGroup, WithResource))
  
  
  {- | A test case with a wrapper function that can do some IO around the
  test. We use the wrapper to set up spans appropriately.
  -}
  data WrappedTest t = WrappedTest
    {wrapper :: forall a. IO a -> IO a, innerTest :: t}
  
  
  instance IsTest t => IsTest (WrappedTest t) where
    run opts (WrappedTest {wrapper, innerTest}) progress =
      wrapper $ do
        ctx <- getContext
        let mspan = lookupSpan ctx
        res@Result {resultOutcome, resultDescription} <- run opts innerTest progress
        case mspan of
          Just s -> do
            addAttribute s "result.description" (T.pack resultDescription)
            case resultOutcome of
              Success -> do
                setStatus s $ Ok
              Failure reason -> do
                setStatus s $ Error $ T.pack $ show reason
          Nothing -> pure ()
  
        pure res
    testOptions = retag (testOptions :: Tagged t [OptionDescription])
  
  
  -- | Transform a 'TestTree' into one that emits spans around tests and test groups.
  instrumentTestTree :: TestTree -> IO TestTree
  instrumentTestTree t = do
    provider <- getGlobalTracerProvider
    let tracer = makeTracer provider $detectInstrumentationLibrary tracerOptions
    pure $ instrumentTestTreeWithTracer tracer t
  
  
  -- | See 'instrumentTestTree'.
  instrumentTestTreeWithTracer :: Tracer -> TestTree -> TestTree
  instrumentTestTreeWithTracer tracer = instrumentTestTree' tracer (pure Nothing)
  
  
  instrumentTestTree'
    :: Tracer
    -> IO (Maybe Span)
    -> TestTree
    -> TestTree
  instrumentTestTree' tracer = go
    where
      -- See Note [Test parallelism] for why we pass around 'getParentSpan'
      go :: IO (Maybe Span) -> TestTree -> TestTree
      go getParentSpan = \case
        TestGroup name tests ->
          -- We use 'withResource' to associate the creation and destruction of the
          -- group span with the beginning and end of the group itself. This way
          -- 'tasty' manages the lifetime of the span for us.
          withResource
            (mkSpan (T.pack name))
            (\s -> endSpan s Nothing)
            $ \getGroupSpan ->
              let getParentSpan' = Just <$> getGroupSpan
              in TestGroup name (fmap (go getParentSpan') tests)
        SingleTest name t ->
          SingleTest name $
            WrappedTest {wrapper = withNamedSpan name, innerTest = t}
        WithResource (ResourceSpec acquire release) f ->
          -- Add spans for resource acquisition and release
          -- Nit: currently we don't create a span for the top-level itself, so
          -- if you acquire outside a test group then the spans will be detached.
          -- We could add a span for the top level, although it's maybe a little odd.
          let newResourceSpec = ResourceSpec (withNamedSpan "acquire" acquire) (fmap (withNamedSpan "release") release)
          in WithResource newResourceSpec (go' . f)
        PlusTestOptions modifier t -> PlusTestOptions modifier (go' t)
        AskOptions f -> AskOptions (go' . f)
        After d e t -> After d e $ go' t
        where
          go' = go getParentSpan
          mkSpan name = do
            ctx <- getContext
            -- See Note [Test parallelism]
            parentSpan <- getParentSpan
            -- This does not modify the thread-local context, just locally
            ctx' <- case parentSpan of
              Just s -> pure $ insertSpan s ctx
              Nothing -> pure ctx
            createSpan tracer ctx' name defaultSpanArguments
          withNamedSpan :: String -> (forall a. IO a -> IO a)
          withNamedSpan name act = do
            -- See Note [Test parallelism]
            parentSpan <- getParentSpan
            let wrapper = case parentSpan of
                  Just ps -> withParentSpan ps
                  Nothing -> id
            wrapper $ inSpan tracer (T.pack name) defaultSpanArguments act
  
  
  -- Possibly should upstream this to the SDK?
  
  {- | Given a span, produces a wrapper function that sets the given span
  as the installed span in the context.
  -}
  withParentSpan :: Span -> (forall a. IO a -> IO a)
  withParentSpan parentSpan act =
    bracket setup teardown $ \_ -> act
    where
      setup = do
        ctx <- getContext
        adjustContext (insertSpan parentSpan)
        pure (lookupSpan ctx, ctx)
      teardown (originalParentSpan, _ctx) = do
        adjustContext $ \ctx -> maybe (removeSpan ctx) (`insertSpan` ctx) originalParentSpan
  
  {- Note [Test parallelism]
  Tasty runs tests in parallel by default, and we don't want to disturb that.
  However, that means that any individual test case may be running on a random thread at tasty's
  discretion, and so we can't rely on the thread-local context to link up spans.
  
  Our solution is just to track (an action to access) the parent span manually as we traverse
  the tree, so we can connect them up manually.
  -}
  