module Webserver.Control.Telemetry (
  initTelemetry,
  shutdownTelemetry,
  recordSpan,
  incCounter,
  observeHistogram,
  recordLog,
  generateTraceId,
  generateSpanId,
  telemetryMiddleware,
  getTraceContext,
  withChildSpan,
) where

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

import           Control.Concurrent       (forkIO, threadDelay)
import           Control.Exception        (SomeException, catch, onException,
                                           throwIO, try)
import           Control.Monad            (forever, mapM, unless, void)
import qualified Data.Aeson               as Aeson
import qualified Data.ByteString          as BS
import qualified Data.ByteString.Builder  as Builder
import qualified Data.ByteString.Lazy     as BSL
import           Data.Functor             (fmap)
import qualified Data.IORef               as IORef
import qualified Data.List                as List
import qualified Data.Map.Strict          as Map
import qualified Data.Maybe               as Maybe
import qualified Data.Text                as Text
import qualified Data.Text.Encoding       as TSE
import qualified Data.Text.Encoding.Error as TEE
import           Data.Time.Clock          (UTCTime, diffUTCTime, getCurrentTime)
import qualified Data.Vault.Lazy          as Vault
import           GHC.Real                 (mod, realToFrac)
import qualified Network.HTTP.Client      as HTTP
import qualified Network.HTTP.Types       as Http
import qualified Network.Wai              as Wai
import qualified Network.Wai.Internal     as WaiI
import qualified System.Environment       as Env
import qualified System.IO                as SysIO
import           System.Random            (randomIO)
import           Text.Printf              (printf)
import           Webserver.Control.Nil    (Nil (..))
import           Webserver.Data.Telemetry

----------------------------------------------------------------------------------------------------
-- ID generation
----------------------------------------------------------------------------------------------------

-- | Generate a 32 hex-char trace ID from 16 random bytes
generateTraceId :: IO TraceId
generateTraceId = do
  bytes <- generateRandomHex 16
  return <| TraceId bytes

-- | Generate a 16 hex-char span ID from 8 random bytes
generateSpanId :: IO SpanId
generateSpanId = do
  bytes <- generateRandomHex 8
  return <| SpanId bytes

-- | Generate n random bytes as hex text
generateRandomHex :: Int -> IO Text
generateRandomHex n = do
  bytes <- mapM (\_ -> randomIO :: IO Int) ([1 .. n] :: [Int])
  return
    <| ( bytes
           |> fmap (\b -> printf "%02x" (b `mod` 256) :: String)
           |> List.concat
           |> Text.pack
       )

----------------------------------------------------------------------------------------------------
-- Lifecycle
----------------------------------------------------------------------------------------------------

{- | Initialize the telemetry subsystem: create state, HTTP manager, Vault key,
  start background flush threads, register default metrics
-}
initTelemetry :: IO TelemetryState
initTelemetry = do
  endpoint <-
    Env.lookupEnv "OTEL_EXPORTER_OTLP_ENDPOINT"
      <&> \case
        Just e -> Text.pack e
        Nothing -> "http://localhost:4318"

  svcName <-
    Env.lookupEnv "OTEL_SERVICE_NAME"
      <&> \case
        Just s -> Text.pack s
        Nothing -> "webserver"

  let cfg =
        TelemetryConfig
          { otlpEndpoint = endpoint
          , serviceName = svcName
          , flushInterval = 5000000
          , maxBufferSize = 10000
          , samplingRate = 1.0
          }

  spanBuf <- IORef.newIORef []
  logBuf <- IORef.newIORef []
  cntrs <- IORef.newIORef Map.empty
  hists <- IORef.newIORef Map.empty
  mgr <- HTTP.newManager HTTP.defaultManagerSettings
  vaultKey <- Vault.newKey

  let st =
        TelemetryState
          { spanBuffer = spanBuf
          , logBuffer = logBuf
          , counters = cntrs
          , histograms = hists
          , httpManager = mgr
          , config = cfg
          , traceKey = vaultKey
          }

  -- Register default metrics
  registerCounter st "http_requests_total"
  registerHistogram st "http_request_duration_seconds" defaultDurationBuckets

  -- Start background flush threads
  void <| forkIO <| flushLoop st flushSpans
  void <| forkIO <| flushLoop st flushMetrics
  void <| forkIO <| flushLoop st flushLogs

  return st

-- | Final flush of all buffers before shutdown
shutdownTelemetry :: TelemetryState -> IO ()
shutdownTelemetry st = do
  flushSpans st
  flushMetrics st
  flushLogs st

-- | Default histogram buckets for HTTP request duration (in seconds)
defaultDurationBuckets :: [Double]
defaultDurationBuckets = [0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10]

----------------------------------------------------------------------------------------------------
-- Metric registration
----------------------------------------------------------------------------------------------------

-- | Register a named counter (idempotent — existing counter is kept)
registerCounter :: TelemetryState -> Text -> IO ()
registerCounter st cName = do
  ref <- IORef.newIORef 0
  let counter = Counter {counterName = cName, counterValue = ref}
  IORef.atomicModifyIORef' st.counters \m ->
    case Map.lookup cName m of
      Just _  -> (m, ())
      Nothing -> (Map.insert cName counter m, ())

-- | Register a named histogram with given buckets (idempotent)
registerHistogram :: TelemetryState -> Text -> [Double] -> IO ()
registerHistogram st hName buckets = do
  countsRef <- IORef.newIORef Map.empty
  sumRef <- IORef.newIORef 0
  countRef <- IORef.newIORef 0
  let hist =
        Histogram
          { histName = hName
          , histBuckets = buckets
          , histCounts = countsRef
          , histSum = sumRef
          , histCount = countRef
          }
  IORef.atomicModifyIORef' st.histograms \m ->
    case Map.lookup hName m of
      Just _  -> (m, ())
      Nothing -> (Map.insert hName hist m, ())

----------------------------------------------------------------------------------------------------
-- Recording (non-blocking, IORef atomic)
----------------------------------------------------------------------------------------------------

-- | Append a completed span to the buffer, dropping oldest if at capacity
recordSpan :: TelemetryState -> Span -> IO ()
recordSpan st s =
  IORef.atomicModifyIORef' st.spanBuffer \buf ->
    let buf' = s : buf
     in if List.length buf' > st.config.maxBufferSize
          then (List.take st.config.maxBufferSize buf', ())
          else (buf', ())

-- | Increment a named counter by 1
incCounter :: TelemetryState -> Text -> IO ()
incCounter st cName = do
  cntrs <- IORef.readIORef st.counters
  case Map.lookup cName cntrs of
    Nothing -> return ()
    Just counter ->
      IORef.atomicModifyIORef' counter.counterValue \n -> (n + 1, ())

-- | Record a value observation into a named histogram
observeHistogram :: TelemetryState -> Text -> Double -> IO ()
observeHistogram st hName val = do
  hists <- IORef.readIORef st.histograms
  case Map.lookup hName hists of
    Nothing -> return ()
    Just hist -> do
      let bucket = findBucket hist.histBuckets val
      IORef.atomicModifyIORef' hist.histCounts \m ->
        (Map.insertWith (+) bucket 1 m, ())
      IORef.atomicModifyIORef' hist.histSum \s -> (s + val, ())
      IORef.atomicModifyIORef' hist.histCount \c -> (c + 1, ())

-- | Append a log entry to the buffer
recordLog :: TelemetryState -> Severity -> Text -> Map.Map Text Text -> TraceId -> SpanId -> IO ()
recordLog st sev body attrs tId sId = do
  now <- getCurrentTime
  let entry =
        LogEntry
          { logTime = now
          , logSeverity = sev
          , logBody = body
          , logAttrs = attrs
          , logTraceId = tId
          , logSpanId = sId
          }
  IORef.atomicModifyIORef' st.logBuffer \buf ->
    let buf' = entry : buf
     in if List.length buf' > st.config.maxBufferSize
          then (List.take st.config.maxBufferSize buf', ())
          else (buf', ())

----------------------------------------------------------------------------------------------------
-- Background flush
----------------------------------------------------------------------------------------------------

-- | Run a flush function in a loop with the configured interval
flushLoop :: TelemetryState -> (TelemetryState -> IO ()) -> IO ()
flushLoop st flushFn = forever do
  threadDelay st.config.flushInterval
  flushFn st
    `catch` \(e :: SomeException) ->
      SysIO.hPutStrLn SysIO.stderr <| "Telemetry flush error: " <> show e

-- | Atomically drain the span buffer and POST to /v1/traces
flushSpans :: TelemetryState -> IO ()
flushSpans st = do
  spans <- IORef.atomicModifyIORef' st.spanBuffer ([],)
  unless (List.null spans) do
    let payload = OtlpTraceExport st.config.serviceName spans
        url = Text.unpack st.config.otlpEndpoint <> "/v1/traces"
    postOtlpJson st.httpManager url payload

-- | Snapshot all counters and histograms, POST to /v1/metrics
flushMetrics :: TelemetryState -> IO ()
flushMetrics st = do
  now <- getCurrentTime
  cntrs <- IORef.readIORef st.counters
  hists <- IORef.readIORef st.histograms

  counterMetrics <- mapM (snapshotCounter now) (Map.elems cntrs)
  histMetrics <- mapM (snapshotHistogram now) (Map.elems hists)

  let allMetrics = counterMetrics <> histMetrics
  unless (List.null allMetrics) do
    let payload = OtlpMetricExport st.config.serviceName allMetrics
        url = Text.unpack st.config.otlpEndpoint <> "/v1/metrics"
    postOtlpJson st.httpManager url payload

-- | Atomically drain the log buffer and POST to /v1/logs
flushLogs :: TelemetryState -> IO ()
flushLogs st = do
  logs <- IORef.atomicModifyIORef' st.logBuffer ([],)
  unless (List.null logs) do
    let payload = OtlpLogExport st.config.serviceName logs
        url = Text.unpack st.config.otlpEndpoint <> "/v1/logs"
    postOtlpJson st.httpManager url payload

-- | Snapshot a counter to a pure OtlpMetric value
snapshotCounter :: UTCTime -> Counter -> IO OtlpMetric
snapshotCounter now counter = do
  val <- IORef.readIORef counter.counterValue
  return <| OtlpCounter counter.counterName val now Map.empty

-- | Snapshot a histogram to a pure OtlpMetric value
snapshotHistogram :: UTCTime -> Histogram -> IO OtlpMetric
snapshotHistogram now hist = do
  counts <- IORef.readIORef hist.histCounts
  hSum <- IORef.readIORef hist.histSum
  hCount <- IORef.readIORef hist.histCount
  return <| OtlpHistogram hist.histName hist.histBuckets counts hSum hCount now Map.empty

-- | POST OTLP JSON payload to an endpoint, silently dropping on failure
postOtlpJson :: (Aeson.ToJSON a) => HTTP.Manager -> String -> a -> IO ()
postOtlpJson mgr url payload = do
  initReq <- HTTP.parseRequest url
  let req =
        initReq
          { HTTP.method = "POST"
          , HTTP.requestBody = HTTP.RequestBodyLBS <| Aeson.encode payload
          , HTTP.requestHeaders =
              [("Content-Type", "application/json")]
          }
  void (HTTP.httpLbs req mgr)
    `catch` \(e :: SomeException) ->
      SysIO.hPutStrLn SysIO.stderr <| "OTLP POST failed (" <> url <> "): " <> show e

----------------------------------------------------------------------------------------------------
-- Histogram helpers
----------------------------------------------------------------------------------------------------

-- | Find the smallest bucket boundary >= the observed value, or +Infinity
findBucket :: [Double] -> Double -> Double
findBucket buckets val =
  case List.filter (>= val) buckets of
    (b : _) -> b
    []      -> 1 / 0

----------------------------------------------------------------------------------------------------
-- WAI Middleware
----------------------------------------------------------------------------------------------------

-- | Telemetry middleware: creates root spans, records metrics and logs per request
telemetryMiddleware :: TelemetryState -> Wai.Middleware
telemetryMiddleware st app request respondWith = do
  -- Check sampling rate
  shouldSample <- sampleRequest st.config.samplingRate
  if not shouldSample
    then app request respondWith
    else do
      -- Generate trace/span IDs and record start time
      tId <- generateTraceId
      sId <- generateSpanId
      startT <- getCurrentTime

      let traceCtx = TraceContext {tcTraceId = tId, tcSpanId = sId}

      -- Extract request metadata from headers
      let method = TSE.decodeUtf8 <| Wai.requestMethod request
          path = "/" <> Text.intercalate "/" (Wai.pathInfo request)
          reqHeaders = Wai.requestHeaders request
          reqContentType =
            List.lookup Http.hContentType reqHeaders
              |> fmap TSE.decodeUtf8
              |> Maybe.fromMaybe ""
          reqContentLength =
            List.lookup Http.hContentLength reqHeaders
              |> fmap TSE.decodeUtf8
              |> Maybe.fromMaybe ""

      -- Buffer request body for methods that carry a body
      let hasBody = case Wai.requestMethod request of
            "GET"     -> False
            "HEAD"    -> False
            "OPTIONS" -> False
            "DELETE"  -> False
            _         -> True
      (request', reqBodyText) <-
        if hasBody
          then bufferRequestBody request
          else return (request, "")

      -- Store TraceContext in WAI Vault
      let vault' = Vault.insert st.traceKey traceCtx (Wai.vault request')
          request'' = request' {Wai.vault = vault'}

      -- Build base request attributes once
      let reqAttrs =
            Map.fromList
              <| [ ("http.method", method)
                 , ("http.path", path)
                 , ("http.request.content_type", reqContentType)
                 , ("http.request.content_length", reqContentLength)
                 , ("http.request.body", redactBody reqBodyText)
                 ]

      -- Increment request counter
      incCounter st "http_requests_total"

      -- Log request received
      recordLog st Info "Request received" reqAttrs tId sId

      -- Wrap respondWith to capture response details
      let respondWith' :: Wai.Response -> IO Wai.ResponseReceived
          respondWith' response = do
            endT <- getCurrentTime
            let duration = realToFrac (diffUTCTime endT startT) :: Double
                statusCode = Http.statusCode <| Wai.responseStatus response
                statusText = showAsText statusCode
                respHeaders = Wai.responseHeaders response
                respContentType =
                  List.lookup Http.hContentType respHeaders
                    |> fmap TSE.decodeUtf8
                    |> Maybe.fromMaybe ""
                respContentLength =
                  List.lookup Http.hContentLength respHeaders
                    |> fmap TSE.decodeUtf8
                    |> Maybe.fromMaybe ""

            -- Layer response attributes on top of request attributes
            let respBody = extractResponseBody response
                respAttrs =
                  reqAttrs
                    |> Map.insert "http.status_code" statusText
                    |> Map.insert "http.duration_s" (showAsText duration)
                    |> Map.insert "http.response.content_type" respContentType
                    |> Map.insert "http.response.content_length" respContentLength
                    |> Map.insert "http.response.body" respBody

            -- Complete root span
            recordSpan
              st
              Span
                { traceId = tId
                , spanId = sId
                , parentId = nil
                , name = method <> " " <> path
                , kind = Server
                , startTime = startT
                , endTime = endT
                , attributes = respAttrs
                , status =
                    if statusCode >= 500
                      then SpanError ("HTTP " <> statusText)
                      else Ok
                }

            -- Record duration histogram
            observeHistogram st "http_request_duration_seconds" duration

            -- Log request completed
            recordLog st Info "Request completed" respAttrs tId sId

            respondWith response

      -- Run app, catching exceptions for error telemetry
      result <- try <| app request'' respondWith'
      case result of
        Right received -> return received
        Left (e :: SomeException) -> do
          endT <- getCurrentTime
          let errMsg = Text.pack <| show e
              errAttrs =
                reqAttrs
                  |> Map.insert "error.message" errMsg

          -- Record error span
          recordSpan
            st
            Span
              { traceId = tId
              , spanId = sId
              , parentId = nil
              , name = method <> " " <> path
              , kind = Server
              , startTime = startT
              , endTime = endT
              , attributes = errAttrs
              , status = SpanError errMsg
              }

          -- Record error log
          recordLog st Error ("Request failed: " <> errMsg) errAttrs tId sId

          -- Re-throw so the app can handle it upstream
          throwIO e

-- | Maximum body size to buffer for logging (4 KiB)
maxLogBodySize :: Int
maxLogBodySize = 4096

-- | Keys whose values should be replaced with "[REDACTED]" in logged request bodies
isSensitiveKey :: Text -> Bool
isSensitiveKey k =
  let lower = Text.toLower k
   in List.any
        (`Text.isInfixOf` lower)
        (["password", "secret", "token"] :: List Text)

{- | Parse a JSON body and replace sensitive field values with "[REDACTED]".
  Non-JSON bodies (e.g. plain text) are returned unchanged.
-}
redactBody :: Text -> Text
redactBody bodyText =
  case Aeson.decode (BSL.fromStrict (TSE.encodeUtf8 bodyText)) of
    Just obj ->
      let redacted =
            Map.mapWithKey
              (\k v -> if isSensitiveKey k then Aeson.String "[REDACTED]" else v)
              (obj :: Map.Map Text Aeson.Value)
       in redacted
            |> Aeson.encode
            |> BSL.toStrict
            |> TSE.decodeUtf8With TEE.lenientDecode
    Nothing -> bodyText

{- | Buffer the request body into memory and return a new request with the body
  re-wrapped (so downstream handlers can still read it), plus the body as text
  for logging (truncated to 'maxLogBodySize').
-}
bufferRequestBody :: Wai.Request -> IO (Wai.Request, Text)
bufferRequestBody request = do
  chunks <- readAllChunks request
  let body = BS.concat chunks
      bodyForLog =
        if BS.length body > maxLogBodySize
          then TSE.decodeUtf8With TEE.lenientDecode (BS.take maxLogBodySize body) <> "...(truncated)"
          else TSE.decodeUtf8With TEE.lenientDecode body
  bodyRef <- IORef.newIORef (Just body)
  let getBody =
        IORef.atomicModifyIORef' bodyRef \case
          Just b -> (Nothing, b)
          Nothing -> (Nothing, BS.empty)
      request' = request {Wai.requestBody = getBody}
  return (request', bodyForLog)

{- | Extract the response body as text for logging.
  For builder/LBS responses, returns the body content (truncated to 'maxLogBodySize').
  For file responses, returns the file path.
  For stream/raw responses, returns a placeholder.
-}
extractResponseBody :: Wai.Response -> Text
extractResponseBody = \case
  WaiI.ResponseFile _ _ path _ ->
    "(file: " <> Text.pack path <> ")"
  WaiI.ResponseBuilder _ _ builder ->
    let body = BSL.toStrict <| Builder.toLazyByteString builder
     in truncateBody body
  WaiI.ResponseStream {} ->
    "(stream)"
  WaiI.ResponseRaw _ _ ->
    "(raw)"

-- | Truncate a ByteString to 'maxLogBodySize' and decode as UTF-8 for logging
truncateBody :: BS.ByteString -> Text
truncateBody body =
  if BS.length body > maxLogBodySize
    then TSE.decodeUtf8With TEE.lenientDecode (BS.take maxLogBodySize body) <> "...(truncated)"
    else TSE.decodeUtf8With TEE.lenientDecode body

-- | Read all chunks from a WAI request body into a list
readAllChunks :: Wai.Request -> IO [BS.ByteString]
readAllChunks request = do
  chunk <- Wai.getRequestBodyChunk request
  if BS.null chunk
    then return []
    else do
      rest <- readAllChunks request
      return (chunk : rest)

-- | Decide whether to sample this request based on sampling rate
sampleRequest :: Double -> IO Bool
sampleRequest rate
  | rate >= 1.0 = return True
  | rate <= 0.0 = return False
  | otherwise = do
      val <- randomIO :: IO Double
      return <| val < rate

----------------------------------------------------------------------------------------------------
-- Handler helpers
----------------------------------------------------------------------------------------------------

-- | Extract the per-request TraceContext from the WAI Vault
getTraceContext :: Wai.Request -> TelemetryState -> Maybe TraceContext
getTraceContext request st = Vault.lookup st.traceKey (Wai.vault request)

-- | Run an action within a child span, automatically recording it
withChildSpan :: TelemetryState -> TraceContext -> Text -> IO a -> IO a
withChildSpan st parentCtx spanName action = do
  childSId <- generateSpanId
  startT <- getCurrentTime
  result <-
    action
      `onException` do
        endT <- getCurrentTime
        recordSpan
          st
          Span
            { traceId = parentCtx.tcTraceId
            , spanId = childSId
            , parentId = parentCtx.tcSpanId
            , name = spanName
            , kind = Internal
            , startTime = startT
            , endTime = endT
            , attributes = Map.empty
            , status = SpanError "exception"
            }
  endT <- getCurrentTime
  recordSpan
    st
    Span
      { traceId = parentCtx.tcTraceId
      , spanId = childSId
      , parentId = parentCtx.tcSpanId
      , name = spanName
      , kind = Internal
      , startTime = startT
      , endTime = endT
      , attributes = Map.empty
      , status = Ok
      }
  return result