module Webserver.Data.Telemetry where

----------------------------------------------------------------------------------------------------

import qualified Data.Aeson            as Aeson
import           Data.Functor          (fmap)
import qualified Data.IORef            as IORef
import qualified Data.Map.Strict       as Map
import           Data.Time.Clock       (UTCTime)
import           Data.Time.Clock.POSIX (utcTimeToPOSIXSeconds)
import           Data.Tuple            (snd)
import qualified Data.Vault.Lazy       as Vault
import           GHC.Generics          (Generic)
import           GHC.Real              (realToFrac, truncate)
import           Network.HTTP.Client   (Manager)
import           Webserver.Control.Nil (Nil (..))

----------------------------------------------------------------------------------------------------

-- | 32 hex chars (16 bytes) identifying a trace across spans
newtype TraceId = TraceId {unTraceId :: Text}
  deriving newtype (Eq, Ord, Show, Aeson.ToJSON)

instance Nil TraceId where
  nil = TraceId ""

-- | 16 hex chars (8 bytes) identifying a single span
newtype SpanId = SpanId {unSpanId :: Text}
  deriving newtype (Eq, Ord, Show, Aeson.ToJSON)

instance Nil SpanId where
  nil = SpanId ""

----------------------------------------------------------------------------------------------------

data SpanKind
  = Internal
  | Server
  | Client
  deriving stock (Eq, Show, Generic)

instance Aeson.ToJSON SpanKind where
  toJSON :: SpanKind -> Aeson.Value
  toJSON Internal = Aeson.Number 1
  toJSON Server   = Aeson.Number 2
  toJSON Client   = Aeson.Number 3

data SpanStatus
  = Unset
  | Ok
  | SpanError !Text
  deriving stock (Eq, Show, Generic)

instance Aeson.ToJSON SpanStatus where
  toJSON :: SpanStatus -> Aeson.Value
  toJSON Unset = Aeson.object [("code", Aeson.Number 0)]
  toJSON Ok = Aeson.object [("code", Aeson.Number 1)]
  toJSON (SpanError msg) =
    Aeson.object
      [ ("code", Aeson.Number 2)
      , ("message", Aeson.String msg)
      ]

data Span = Span
  { traceId    :: !TraceId
  , spanId     :: !SpanId
  , parentId   :: !SpanId
  , name       :: !Text
  , kind       :: !SpanKind
  , startTime  :: !UTCTime
  , endTime    :: !UTCTime
  , attributes :: !(Map.Map Text Text)
  , status     :: !SpanStatus
  }
  deriving stock (Show, Generic)

instance Aeson.ToJSON Span where
  toJSON :: Span -> Aeson.Value
  toJSON s =
    Aeson.object
      [ ("traceId", Aeson.toJSON s.traceId)
      , ("spanId", Aeson.toJSON s.spanId)
      , ("parentSpanId", Aeson.toJSON s.parentId)
      , ("name", Aeson.String s.name)
      , ("kind", Aeson.toJSON s.kind)
      , ("startTimeUnixNano", Aeson.String <| utcToNanoText s.startTime)
      , ("endTimeUnixNano", Aeson.String <| utcToNanoText s.endTime)
      , ("attributes", Aeson.toJSON <| attrsToOtlp s.attributes)
      , ("status", Aeson.toJSON s.status)
      ]

----------------------------------------------------------------------------------------------------

-- | An incrementing counter metric backed by an IORef
data Counter = Counter
  { counterName  :: !Text
  , counterValue :: !(IORef.IORef Int)
  }

-- | A histogram metric with configurable buckets, backed by IORefs
data Histogram = Histogram
  { histName    :: !Text
  , histBuckets :: ![Double]
  , histCounts  :: !(IORef.IORef (Map.Map Double Int))
  , histSum     :: !(IORef.IORef Double)
  , histCount   :: !(IORef.IORef Int)
  }

----------------------------------------------------------------------------------------------------

data Severity
  = Trace
  | Debug
  | Info
  | Warn
  | Error
  | Fatal
  deriving stock (Eq, Show, Generic)

instance Aeson.ToJSON Severity where
  toJSON :: Severity -> Aeson.Value
  toJSON Trace = Aeson.Number 1
  toJSON Debug = Aeson.Number 5
  toJSON Info  = Aeson.Number 9
  toJSON Warn  = Aeson.Number 13
  toJSON Error = Aeson.Number 17
  toJSON Fatal = Aeson.Number 21

severityText :: Severity -> Text
severityText Trace = "TRACE"
severityText Debug = "DEBUG"
severityText Info  = "INFO"
severityText Warn  = "WARN"
severityText Error = "ERROR"
severityText Fatal = "FATAL"

data LogEntry = LogEntry
  { logTime     :: !UTCTime
  , logSeverity :: !Severity
  , logBody     :: !Text
  , logAttrs    :: !(Map.Map Text Text)
  , logTraceId  :: !TraceId
  , logSpanId   :: !SpanId
  }
  deriving stock (Show, Generic)

instance Aeson.ToJSON LogEntry where
  toJSON :: LogEntry -> Aeson.Value
  toJSON le =
    Aeson.object
      [ ("timeUnixNano", Aeson.String <| utcToNanoText le.logTime)
      , ("severityNumber", Aeson.toJSON le.logSeverity)
      , ("severityText", Aeson.String <| severityText le.logSeverity)
      , ("body", Aeson.object [("stringValue", Aeson.String le.logBody)])
      , ("attributes", Aeson.toJSON <| attrsToOtlp le.logAttrs)
      , ("traceId", Aeson.toJSON le.logTraceId)
      , ("spanId", Aeson.toJSON le.logSpanId)
      ]

----------------------------------------------------------------------------------------------------

-- | Wraps spans into OTLP JSON export format
data OtlpTraceExport = OtlpTraceExport !Text ![Span]

instance Aeson.ToJSON OtlpTraceExport where
  toJSON :: OtlpTraceExport -> Aeson.Value
  toJSON (OtlpTraceExport svcName spans) =
    Aeson.object
      [
        ( "resourceSpans"
        , jsonArray
            [ Aeson.object
                [ ("resource", resourceAttrs svcName)
                ,
                  ( "scopeSpans"
                  , jsonArray
                      [ Aeson.object
                          [ ("scope", Aeson.object [("name", Aeson.String "webserver")])
                          , ("spans", Aeson.toJSON spans)
                          ]
                      ]
                  )
                ]
            ]
        )
      ]

-- | Wraps metrics into OTLP JSON export format
data OtlpMetricExport = OtlpMetricExport !Text ![OtlpMetric]

instance Aeson.ToJSON OtlpMetricExport where
  toJSON :: OtlpMetricExport -> Aeson.Value
  toJSON (OtlpMetricExport svcName metrics) =
    Aeson.object
      [
        ( "resourceMetrics"
        , jsonArray
            [ Aeson.object
                [ ("resource", resourceAttrs svcName)
                ,
                  ( "scopeMetrics"
                  , jsonArray
                      [ Aeson.object
                          [ ("scope", Aeson.object [("name", Aeson.String "webserver")])
                          , ("metrics", Aeson.toJSON metrics)
                          ]
                      ]
                  )
                ]
            ]
        )
      ]

-- | Wraps logs into OTLP JSON export format
data OtlpLogExport = OtlpLogExport !Text ![LogEntry]

instance Aeson.ToJSON OtlpLogExport where
  toJSON :: OtlpLogExport -> Aeson.Value
  toJSON (OtlpLogExport svcName logs) =
    Aeson.object
      [
        ( "resourceLogs"
        , jsonArray
            [ Aeson.object
                [ ("resource", resourceAttrs svcName)
                ,
                  ( "scopeLogs"
                  , jsonArray
                      [ Aeson.object
                          [ ("scope", Aeson.object [("name", Aeson.String "webserver")])
                          , ("logRecords", Aeson.toJSON logs)
                          ]
                      ]
                  )
                ]
            ]
        )
      ]

-- | Pure snapshot of a metric at flush time (no IORefs)
data OtlpMetric
  = OtlpCounter !Text !Int !UTCTime !(Map.Map Text Text)
  | OtlpHistogram !Text ![Double] !(Map.Map Double Int) !Double !Int !UTCTime !(Map.Map Text Text)

instance Aeson.ToJSON OtlpMetric where
  toJSON :: OtlpMetric -> Aeson.Value
  toJSON (OtlpCounter mName val ts attrs) =
    Aeson.object
      [ ("name", Aeson.String mName)
      ,
        ( "sum"
        , Aeson.object
            [ ("isMonotonic", Aeson.Bool True)
            ,
              ( "dataPoints"
              , jsonArray
                  [ Aeson.object
                      [ ("asInt", Aeson.String <| showAsText val)
                      , ("timeUnixNano", Aeson.String <| utcToNanoText ts)
                      , ("attributes", Aeson.toJSON <| attrsToOtlp attrs)
                      ]
                  ]
              )
            ]
        )
      ]
  toJSON (OtlpHistogram mName buckets counts hSum hCount ts attrs) =
    let sortedBuckets = Map.toAscList counts
        bucketCounts = fmap snd sortedBuckets
     in Aeson.object
          [ ("name", Aeson.String mName)
          ,
            ( "histogram"
            , Aeson.object
                [
                  ( "dataPoints"
                  , jsonArray
                      [ Aeson.object
                          [ ("timeUnixNano", Aeson.String <| utcToNanoText ts)
                          , ("count", Aeson.String <| showAsText hCount)
                          , ("sum", Aeson.Number <| realToFrac hSum)
                          , ("bucketCounts", Aeson.toJSON <| fmap showAsText bucketCounts)
                          , ("explicitBounds", Aeson.toJSON buckets)
                          , ("attributes", Aeson.toJSON <| attrsToOtlp attrs)
                          ]
                      ]
                  )
                ]
            )
          ]

----------------------------------------------------------------------------------------------------

-- | Configuration for the telemetry subsystem
data TelemetryConfig = TelemetryConfig
  { otlpEndpoint  :: !Text
  , serviceName   :: !Text
  , flushInterval :: !Int
  , maxBufferSize :: !Int
  , samplingRate  :: !Double
  }
  deriving stock (Eq, Show, Generic)

-- | Central mutable state for the telemetry subsystem
data TelemetryState = TelemetryState
  { spanBuffer  :: !(IORef.IORef [Span])
  , logBuffer   :: !(IORef.IORef [LogEntry])
  , counters    :: !(IORef.IORef (Map.Map Text Counter))
  , histograms  :: !(IORef.IORef (Map.Map Text Histogram))
  , httpManager :: !Manager
  , config      :: !TelemetryConfig
  , traceKey    :: !(Vault.Key TraceContext)
  }

-- | Per-request context stored in WAI Request Vault
data TraceContext = TraceContext
  { tcTraceId :: !TraceId
  , tcSpanId  :: !SpanId
  }
  deriving stock (Eq, Show, Generic)

----------------------------------------------------------------------------------------------------

-- | Encode a list of Aeson Values as a JSON array (avoids OverloadedLists ambiguity)
jsonArray :: [Aeson.Value] -> Aeson.Value
jsonArray = Aeson.toJSON @[Aeson.Value]

-- | Convert a UTCTime to nanoseconds-since-epoch text (OTLP format)
utcToNanoText :: UTCTime -> Text
utcToNanoText t =
  let posix = utcTimeToPOSIXSeconds t
      nanos = truncate (posix * 1000000000) :: Int
   in showAsText nanos

-- | Convert a Map of key-value attributes to OTLP attribute array format
attrsToOtlp :: Map.Map Text Text -> [Aeson.Value]
attrsToOtlp =
  Map.toList
    <. fmap \(k, v) ->
      Aeson.object
        [ ("key", Aeson.String k)
        , ("value", Aeson.object [("stringValue", Aeson.String v)])
        ]

-- | Build OTLP resource attributes from a service name
resourceAttrs :: Text -> Aeson.Value
resourceAttrs svcName =
  Aeson.object
    [
      ( "attributes"
      , jsonArray
          [ Aeson.object
              [ ("key", Aeson.String "service.name")
              , ("value", Aeson.object [("stringValue", Aeson.String svcName)])
              ]
          ]
      )
    ]