module OpenTelemetry.Trace.SamplerSpec where

import Control.Monad
import Test.Hspec
import qualified OpenTelemetry.Context as Context
import OpenTelemetry.Trace.Core
import OpenTelemetry.Trace.Id
import OpenTelemetry.Trace.Sampler
import qualified OpenTelemetry.Trace.TraceState as TraceState

builtInNonCompositeSamplers :: [Sampler]
builtInNonCompositeSamplers =
  [ alwaysOff
  , alwaysOn
  , traceIdRatioBased 0.99
  ]

-- TODO, these would largely be good candidates for quickcheck
spec :: Spec
spec = describe "Sampler" $ do
  describe "built-in non-composite samplers" $ do
    forM_ builtInNonCompositeSamplers $ \sampler -> do
      specify (show (getDescription sampler) ++ " returns parent tracestate") $ do
        let (Right aTraceId) = baseEncodedToTraceId Base16 "00000000000000000000000000000001"
            (Right parentSpanId) = baseEncodedToSpanId Base16 "000000000000000a"
            remoteSpan = SpanContext
              { traceFlags = defaultTraceFlags
              , isRemote = False
              , traceState = TraceState.insert (TraceState.Key "k") (TraceState.Value "v") TraceState.empty
              , spanId = parentSpanId
              , traceId = aTraceId
              }
            traceParent = wrapSpanContext remoteSpan
            parentContext = Context.insertSpan traceParent Context.empty
        (_, _, ts) <- shouldSample sampler parentContext aTraceId "test" defaultSpanArguments
        ts `shouldBe` traceState remoteSpan


  describe "alwaysOff" $ do
    it "returns Drop" $ do
      let (Right aTraceId) = baseEncodedToTraceId Base16 "00000000000000000000000000000001"
          (Right parentSpanId) = baseEncodedToSpanId Base16 "000000000000000a"
          remoteSpan = SpanContext
            { traceFlags = defaultTraceFlags
            , isRemote = False
            , traceState = TraceState.empty
            , spanId = parentSpanId
            , traceId = aTraceId
            }
          traceParent = wrapSpanContext remoteSpan
          parentContext = Context.insertSpan traceParent Context.empty
      (res, _, _) <- shouldSample alwaysOff parentContext aTraceId "test" defaultSpanArguments
      res `shouldBe` Drop

  describe "alwaysOn" $ do
    it "returns RecordAndSample" $ do
      let (Right aTraceId) = baseEncodedToTraceId Base16 "00000000000000000000000000000001"
          (Right parentSpanId) = baseEncodedToSpanId Base16 "000000000000000a"
          remoteSpan = SpanContext
            { traceFlags = defaultTraceFlags
            , isRemote = False
            , traceState = TraceState.empty
            , spanId = parentSpanId
            , traceId = aTraceId
            }
          traceParent = wrapSpanContext remoteSpan
          parentContext = Context.insertSpan traceParent Context.empty
      (res, _, _) <- shouldSample alwaysOn parentContext aTraceId "test" defaultSpanArguments
      res `shouldBe` RecordAndSample

  describe "traceIdRatioBased" $ do
    it "drops spans that are outside of the sample ratio" $ do
      let sampler = traceIdRatioBased 0.5
          (Right aTraceId) = baseEncodedToTraceId Base16 "ffffffffffffffff0000000000000000"
          (Right parentSpanId) = baseEncodedToSpanId Base16 "0000000000000001"
          remoteSpan = SpanContext
            { traceFlags = defaultTraceFlags
            , isRemote = False
            , traceState = TraceState.empty
            , spanId = parentSpanId
            , traceId = aTraceId
            }
          traceParent = wrapSpanContext remoteSpan
          parentContext = Context.insertSpan traceParent Context.empty
      (res, _, _) <- shouldSample sampler parentContext aTraceId "test" defaultSpanArguments
      res `shouldBe` Drop
    it "samples spans that are within the sample ratio" $ do
      let sampler = traceIdRatioBased 0.5
          (Right aTraceId) = baseEncodedToTraceId Base16 "00000000000000000000000000000000"
          (Right parentSpanId) = baseEncodedToSpanId Base16 "0000000000000001"
          remoteSpan = SpanContext
            { traceFlags = defaultTraceFlags
            , isRemote = False
            , traceState = TraceState.empty
            , spanId = parentSpanId
            , traceId = aTraceId
            }
          traceParent = wrapSpanContext remoteSpan
          parentContext = Context.insertSpan traceParent Context.empty
      (res, _, _) <- shouldSample sampler parentContext aTraceId "test" defaultSpanArguments
      res `shouldBe` RecordAndSample
    it "higher sample ratios sample spans that are also by lower ratios" $ do
      let conservativeSampler = traceIdRatioBased 0.5
          permissiveSampler = traceIdRatioBased 0.75
          (Right aTraceId) = baseEncodedToTraceId Base16 "3fffffffffffffff0000000000000000"
          (Right parentSpanId) = baseEncodedToSpanId Base16 "0000000000000001"
          remoteSpan = SpanContext
            { traceFlags = defaultTraceFlags
            , isRemote = False
            , traceState = TraceState.empty
            , spanId = parentSpanId
            , traceId = aTraceId
            }
          traceParent = wrapSpanContext remoteSpan
          parentContext = Context.insertSpan traceParent Context.empty
      do
        (resConservative, _, _) <- shouldSample permissiveSampler parentContext aTraceId "test" defaultSpanArguments
        (resPermissive, _, _) <- shouldSample conservativeSampler parentContext aTraceId "test" defaultSpanArguments
        resConservative `shouldBe` RecordAndSample
        resPermissive `shouldBe` RecordAndSample