content: {-# LANGUAGE DeriveFunctor #-}
  {-# LANGUAGE DuplicateRecordFields #-}
  {-# LANGUAGE GeneralizedNewtypeDeriving #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE OverloadedStrings #-}
  
  module OpenTelemetry.Common where
  
  import Control.Monad
  import qualified Data.ByteString as BS
  import qualified Data.HashMap.Strict as HM
  import Data.Hashable
  import Data.IORef (modifyIORef, newIORef, readIORef)
  import Data.List (sortOn)
  import Data.String
  import qualified Data.Text as T
  import Data.Word
  import GHC.Generics
  import GHC.Int (Int8)
  import OpenTelemetry.SpanContext
  import System.Clock
  
  type Timestamp = Word64
  
  newtype SpanName = SpanName T.Text deriving (Show, Eq, Generic)
  
  newtype TagName = TagName T.Text deriving (Show, Eq, Generic, Hashable)
  
  newtype TagVal = TagVal T.Text deriving (Show, Eq, Generic)
  
  newtype EventName = EventName T.Text deriving (Show, Eq, Generic)
  
  newtype EventVal = EventVal T.Text deriving (Show, Eq, Generic)
  
  instance IsString TagName where
    fromString = TagName . T.pack
  
  data TagValue
    = StringTagValue !TagVal
    | BoolTagValue !Bool
    | IntTagValue !Int
    | DoubleTagValue !Double
    deriving (Eq, Show)
  
  class ToTagValue a where
    toTagValue :: a -> TagValue
  
  instance ToTagValue String where
    toTagValue = StringTagValue . TagVal . T.pack
  
  instance ToTagValue TagVal where
    toTagValue = StringTagValue
  
  instance ToTagValue T.Text where
    toTagValue = StringTagValue . TagVal
  
  instance ToTagValue Bool where
    toTagValue = BoolTagValue
  
  instance ToTagValue Int where
    toTagValue = IntTagValue
  
  data Span = Span
    { spanContext :: {-# UNPACK #-} !SpanContext,
      spanOperation :: T.Text,
      spanThreadId :: Word32,
      spanDisplayThreadId :: Word32,
      spanStartedAt :: !Timestamp,
      spanFinishedAt :: !Timestamp,
      spanTags :: !(HM.HashMap TagName TagValue),
      spanEvents :: [SpanEvent],
      spanStatus :: !SpanStatus,
      spanParentId :: Maybe SpanId,
      spanNanosecondsSpentInGC :: !Word64
    }
    deriving (Show, Eq)
  
  -- | Reflects the constructors of 'OpenTelemetry.Metrics_Internal.Instrument'
  data InstrumentType
    = CounterType
    | UpDownCounterType
    | ValueRecorderType
    | SumObserverType
    | UpDownSumObserverType
    | ValueObserverType
    deriving (Show, Eq, Enum, Generic)
  
  instance Hashable InstrumentType
  
  data CaptureInstrument = CaptureInstrument
    { instrumentType :: !InstrumentType,
      instrumentName :: !BS.ByteString
    }
    deriving (Show, Eq, Generic)
  
  instance Hashable CaptureInstrument
  
  -- | Based on https://github.com/open-telemetry/opentelemetry-proto/blob/1a931b4b57c34e7fd8f7dddcaa9b7587840e9c08/opentelemetry/proto/metrics/v1/metrics.proto#L96-L107
  data Metric = Metric
    { instrument :: !CaptureInstrument,
      datapoints :: ![MetricDatapoint Int]
    }
    deriving (Show, Eq)
  
  data AggregatedMetric = AggregatedMetric
    { instrument :: !CaptureInstrument,
      datapoint :: !(MetricDatapoint Int)
    }
    deriving (Show, Eq)
  
  data MetricDatapoint a = MetricDatapoint
    { timestamp :: !Timestamp,
      value :: !a
    }
    deriving (Show, Eq, Functor)
  
  spanTraceId :: Span -> TraceId
  spanTraceId Span {spanContext = SpanContext _ tid} = tid
  
  spanId :: Span -> SpanId
  spanId Span {spanContext = SpanContext sid _} = sid
  
  data SpanEvent = SpanEvent
    { spanEventTimestamp :: !Timestamp,
      spanEventKey :: !EventName,
      spanEventValue :: !EventVal
    }
    deriving (Show, Eq)
  
  data SpanStatus = OK
    deriving (Show, Eq)
  
  data Event
    = Event T.Text Timestamp
    deriving (Show, Eq)
  
  data SpanProcessor = SpanProcessor
    { onStart :: Span -> IO (),
      onEnd :: Span -> IO ()
    }
  
  data OpenTelemetryConfig = OpenTelemetryConfig
    { otcSpanExporter :: Exporter Span
    }
  
  data ExportResult
    = ExportSuccess
    | ExportFailedRetryable
    | ExportFailedNotRetryable
    deriving (Show, Eq)
  
  data Exporter thing = Exporter
    { export :: [thing] -> IO ExportResult,
      shutdown :: IO ()
    }
  
  readInstrumentTag :: Int8 -> Maybe InstrumentType
  readInstrumentTag 1 = Just CounterType
  readInstrumentTag 2 = Just UpDownCounterType
  readInstrumentTag 3 = Just ValueRecorderType
  readInstrumentTag 4 = Just SumObserverType
  readInstrumentTag 5 = Just UpDownSumObserverType
  readInstrumentTag 6 = Just ValueObserverType
  readInstrumentTag _ = Nothing
  
  additive :: InstrumentType -> Bool
  additive CounterType = True
  additive UpDownCounterType = True
  additive ValueRecorderType = False
  additive SumObserverType = True
  additive UpDownSumObserverType = True
  additive ValueObserverType = False
  
  noopExporter :: Exporter whatever
  noopExporter = Exporter (const (pure ExportFailedNotRetryable)) (pure ())
  
  aggregated :: Exporter AggregatedMetric -> IO (Exporter Metric)
  aggregated (Exporter export shutdown) = do
    -- We keep a mutable map of latest metric values. When a new datapoint comes
    -- in, it either replaces or gets added to the current value, based on whether
    -- the instrument is additive.
    currentValuesRef <- newIORef HM.empty
    return $
      Exporter
        { export = \metrics -> do
            forM_ metrics $ \(Metric instrument datapoints) -> do
              forM_ (sortOn timestamp datapoints) $ \dp@(MetricDatapoint ts value) ->
                modifyIORef currentValuesRef $
                  if additive (instrumentType instrument)
                    then
                      HM.alter
                        ( \case
                            Nothing -> Just dp
                            Just (MetricDatapoint _ oldValue) -> Just (MetricDatapoint ts $ oldValue + value)
                        )
                        instrument
                    else HM.insert instrument dp
  
            -- Read the latest value for each instrument
            currentValues <- readIORef currentValuesRef
            export [AggregatedMetric i (currentValues HM.! i) | Metric i _ <- metrics],
          shutdown
        }
  
  now64 :: IO Timestamp
  now64 = do
    TimeSpec secs nsecs <- getTime Realtime
    pure $! fromIntegral secs * 1_000_000_000 + fromIntegral nsecs
  