content: {-# LANGUAGE DerivingStrategies #-}
  {-# LANGUAGE ImportQualifiedPost #-}
  {-# LANGUAGE NamedFieldPuns #-}
  {-# LANGUAGE OverloadedLists #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE PatternSynonyms #-}
  
  module OpenTelemetry.Instrumentation.Tasty.Tests (tests) where
  
  import Control.Concurrent (newEmptyMVar, putMVar, takeMVar)
  import Control.Concurrent.Async (async)
  import Control.Exception (bracket)
  import Data.Functor (void)
  import Data.IORef (atomicModifyIORef, newIORef, readIORef)
  import Data.Map (Map)
  import Data.Map qualified as Map
  import Data.Maybe (fromMaybe)
  import Data.Set qualified as Set
  import Data.Text (Text)
  import OpenTelemetry.Instrumentation.Tasty (instrumentTestTree)
  import OpenTelemetry.Processor.Span (ShutdownResult (ShutdownSuccess), SpanProcessor, spanProcessorForceFlush, spanProcessorOnEnd, spanProcessorOnStart, spanProcessorShutdown, pattern SpanProcessor)
  import OpenTelemetry.Trace (ImmutableSpan (ImmutableSpan, spanName, spanParent), createTracerProvider, defaultSpanArguments, emptyTracerProviderOptions, getGlobalTracerProvider, inSpan, makeTracer, setGlobalTracerProvider, shutdownTracerProvider, tracerOptions)
  import OpenTelemetry.Trace.Core (unsafeReadSpan)
  import Test.Tasty (DependencyType (AllFinish), TestTree, sequentialTestGroup, testGroup, withResource)
  import Test.Tasty.HUnit (testCase, (@?=))
  import Test.Tasty.Ingredients (tryIngredients)
  import Test.Tasty.Ingredients.Basic (Quiet (Quiet), consoleTestReporter)
  import Test.Tasty.Options (setOption)
  
  
  tests :: TestTree
  -- each of these tests sets and tears down the global tracer provider,
  -- so they have to run sequentially
  tests =
    sequentialTestGroup
      "OpenTelemetry.Instrumentation.Tasty"
      AllFinish
      [ basic
      , simpleNesting
      , branching
      , testWithSpan
      , parallelism
      , resources
      ]
  
  
  basic :: TestTree
  basic = testCase "basic" $ do
    (_, trees) <- spanTrees $ do
      testTree <- instrumentTestTree $ testCase "hello" $ pure ()
      runTests $ testTree
    trees @?= [Leaf "hello"]
  
  
  simpleNesting :: TestTree
  simpleNesting = testCase "nested groups make a simple tree" $ do
    (_, trees) <- spanTrees $ do
      testTree <-
        instrumentTestTree $
          testGroup
            "g1"
            [ testGroup "g2" [testCase "t1" (pure ())]
            ]
      runTests $ testTree
    trees @?= [Branch "g1" [Branch "g2" [Leaf "t1"]]]
  
  
  branching :: TestTree
  branching = testCase "nested groups make a branching tree" $ do
    (_, trees) <- spanTrees $ do
      testTree <-
        instrumentTestTree $
          testGroup
            "g1"
            [ testGroup "g2" [testCase "t1" (pure ()), testCase "t2" (pure ())]
            , testGroup "g3" [testCase "t3" (pure ())]
            ]
      runTests $ testTree
    trees @?= [Branch "g1" [Branch "g2" [Leaf "t1", Leaf "t2"], Branch "g3" [Leaf "t3"]]]
  
  
  testWithSpan :: TestTree
  testWithSpan = testCase "test that has a span itself" $ do
    (_, trees) <- spanTrees $ do
      testTree <- instrumentTestTree $ testCase "hello" $ do
        tp <- getGlobalTracerProvider
        let tracer = makeTracer tp "test" tracerOptions
        inSpan tracer "inner" defaultSpanArguments $ pure ()
      runTests $ testTree
    trees @?= [Branch "hello" [Leaf "inner"]]
  
  
  parallelism :: TestTree
  parallelism = testCase "parallelism works" $ do
    block <- newEmptyMVar
    (_, trees) <- spanTrees $ do
      testTree <-
        instrumentTestTree $
          testGroup
            "g1"
            -- Force t1 to wait for t2 before it can even begin, so we
            -- should definitely not start the span for t1 until then
            [ testGroup
                "g2"
                [ withResource (takeMVar block) (const $ pure ()) $ \_ ->
                    testCase "t1" (pure ())
                ]
            , testGroup "g3" [testCase "t2" (putMVar block ())]
            ]
      runTests $ testTree
    trees @?= [Branch "g1" [Branch "g2" [Leaf "acquire", Leaf "t1", Leaf "release"], Branch "g3" [Leaf "t2"]]]
  
  
  resources :: TestTree
  resources = testCase "spans for resource setup and teardown" $ do
    (_, trees) <- spanTrees $ do
      tp <- getGlobalTracerProvider
      let tracer = makeTracer tp "test" tracerOptions
      let acquire = inSpan tracer "myAcquire" defaultSpanArguments $ pure ()
      let release _ = inSpan tracer "myRelease" defaultSpanArguments $ pure ()
      testTree <- instrumentTestTree $ withResource acquire release $ \_ -> testCase "hello" $ pure ()
      runTests $ testTree
    trees @?= [Leaf "hello", Branch "acquire" [Leaf "myAcquire"], Branch "release" [Leaf "myRelease"]]
  
  
  data Tree a = Branch a (Set.Set (Tree a)) | Leaf a
    deriving stock (Show, Eq, Ord)
  
  
  spanTrees :: IO a -> IO (a, Set.Set (Tree Text))
  spanTrees act = do
    (processor, readSpans) <- recordingProcessor
    res <- bracket (setup processor) shutdownTracerProvider $ \_ -> do
      act
    spans <- readSpans
    trees <- spansToTrees spans
    pure (res, Set.fromList trees)
    where
      setup processor = do
        tp <- createTracerProvider [processor] emptyTracerProviderOptions
        setGlobalTracerProvider tp
        pure tp
  
  
  toChildMap :: [ImmutableSpan] -> ([Text], Map Text [Text]) -> IO ([Text], Map Text [Text])
  toChildMap [] acc = pure acc
  toChildMap (ImmutableSpan {spanParent, spanName} : spans) (accRoots, accChildren) = case spanParent of
    Just parent -> do
      ImmutableSpan {spanName = parentName} <- unsafeReadSpan parent
      let existingChildren = fromMaybe mempty $ Map.lookup parentName accChildren
      toChildMap spans (accRoots, Map.insert parentName (existingChildren ++ [spanName]) accChildren)
    Nothing -> toChildMap spans (spanName : accRoots, accChildren)
  
  
  toTree :: (Ord a) => Map a [a] -> a -> Tree a
  toTree childMap node = case Map.lookup node childMap of
    Nothing -> Leaf node
    Just children ->
      if null children
        then Leaf node
        else Branch node $ Set.fromList $ fmap (toTree childMap) children
  
  
  spansToTrees :: [ImmutableSpan] -> IO [Tree Text]
  spansToTrees spans = do
    (roots, childMap) <- toChildMap spans mempty
    let trees = fmap (toTree childMap) roots
    pure trees
  
  
  recordingProcessor :: IO (SpanProcessor, IO [ImmutableSpan])
  recordingProcessor = do
    spans <- newIORef []
  
    let processor =
          SpanProcessor
            { spanProcessorOnStart = mempty
            , spanProcessorOnEnd = \spanRef -> do
                immutableSpan <- readIORef spanRef
                atomicModifyIORef spans $ \soFar -> (soFar ++ [immutableSpan], ())
            , spanProcessorShutdown = async $ pure ShutdownSuccess
            , spanProcessorForceFlush = mempty
            }
  
    setGlobalTracerProvider
      =<< createTracerProvider [processor] emptyTracerProviderOptions
  
    pure (processor, readIORef spans)
  
  
  runTests :: TestTree -> IO ()
  runTests t = case tryIngredients [consoleTestReporter] (setOption (Quiet True) mempty) t of
    Just act -> void act
    Nothing -> error "no ingredient"
  