{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE LambdaCase #-}
module OpenTelemetry.Instrumentation.Wai 
  ( newOpenTelemetryWaiMiddleware
  , newOpenTelemetryWaiMiddleware'
  , requestContext
  ) where

import qualified Data.Vault.Lazy as Vault
import Network.HTTP.Types
import Network.Wai
import qualified OpenTelemetry.Context as Context
import OpenTelemetry.Context.ThreadLocal
import OpenTelemetry.Propagator
import OpenTelemetry.Trace.Core
import System.IO.Unsafe
import qualified Data.Text.Encoding as T
import qualified Data.Text as T
import Control.Monad
import Network.Socket
import Data.IP (fromHostAddress, fromHostAddress6)
import OpenTelemetry.Attributes (lookupAttribute)
import Control.Exception (bracket)

newOpenTelemetryWaiMiddleware :: IO Middleware
newOpenTelemetryWaiMiddleware = getGlobalTracerProvider >>= newOpenTelemetryWaiMiddleware'

newOpenTelemetryWaiMiddleware'
  :: TracerProvider 
  -> IO Middleware
newOpenTelemetryWaiMiddleware' tp = do
  waiTracer <- getTracer 
    tp
    "opentelemetry-instrumentation-wai" 
    (TracerOptions Nothing)
  pure $ middleware waiTracer
  where
    middleware :: Tracer -> Middleware
    middleware tracer app req sendResp = do
      let propagator = getTracerProviderPropagators $ getTracerTracerProvider tracer
      let parentContextM = do
            ctx <- getContext
            ctxt <- extract propagator (requestHeaders req) ctx
            attachContext ctxt
      let path_ = T.decodeUtf8 $ rawPathInfo req
          -- peer = remoteHost req
      bracket 
        parentContextM
        (\case
          Nothing -> void detachContext
          Just p -> void (attachContext p)
        )
        $ \_ -> do
          inSpan' tracer path_ (defaultSpanArguments { kind = Server }) $ \requestSpan -> do
            ctxt <- getContext
            addAttributes requestSpan
              [ ( "http.method", toAttribute $ T.decodeUtf8 $ requestMethod req)
              -- , ( "http.url",
              --     toAttribute $
              --     T.decodeUtf8
              --     ((if secure req then "https://" else "http://") <> host req <> ":" <> B.pack (show $ port req) <> path req <> queryString req)
              --   )
              , ( "http.target", toAttribute $ T.decodeUtf8 (rawPathInfo req <> rawQueryString req))
              -- , ( "http.host", toAttribute $ T.decodeUtf8 $ host req)
              -- , ( "http.scheme", toAttribute $ TextAttribute $ if secure req then "https" else "http")
              , ( "http.flavor"
                , toAttribute $ case httpVersion req of
                    (HttpVersion major minor) -> T.pack (show major <> "." <> show minor)
                )
              , ( "http.user_agent"
                , toAttribute $ maybe "" T.decodeUtf8 (lookup hUserAgent $ requestHeaders req)
                )
              -- TODO HTTP/3 will require detecting this dynamically
              , ( "net.transport", toAttribute ("ip_tcp" :: T.Text))
              ]

            -- TODO this is warp dependent, probably.
            -- , ( "net.host.ip")
            -- , ( "net.host.port")
            -- , ( "net.host.name")
            addAttributes requestSpan $ case remoteHost req of
              SockAddrInet port addr ->
                [ ("net.peer.port", toAttribute (fromIntegral port :: Int))
                , ("net.peer.ip", toAttribute $ T.pack $ show $ fromHostAddress addr)
                ]
              SockAddrInet6 port _ addr _ ->
                [ ("net.peer.port", toAttribute (fromIntegral port :: Int))
                , ("net.peer.ip", toAttribute $ T.pack $ show $ fromHostAddress6 addr)
                ]
              SockAddrUnix path ->
                [ ("net.peer.name", toAttribute $ T.pack path)
                ]
            let req' = req 
                  { vault = Vault.insert 
                      contextKey 
                      ctxt
                      (vault req) 
                  }
            app req' $ \resp -> do
              ctxt' <- getContext
              hs <- inject propagator (Context.insertSpan requestSpan ctxt') []
              let resp' = mapResponseHeaders (hs ++) resp
              attrs <- spanGetAttributes requestSpan
              forM_ (lookupAttribute attrs "http.route") $ \case
                AttributeValue (TextAttribute route) -> updateName requestSpan route 
                _ -> pure ()

              addAttributes requestSpan
                [ ( "http.status_code", toAttribute $ statusCode $ responseStatus resp)
                ]
              when (statusCode (responseStatus resp) >= 500) $ do
                setStatus requestSpan (Error "")
              respReceived <- sendResp resp'
              ts <- getTimestamp
              endSpan requestSpan (Just ts)
              pure respReceived

contextKey :: Vault.Key Context.Context
contextKey = unsafePerformIO Vault.newKey
{-# NOINLINE contextKey #-}

requestContext :: Request -> Maybe Context.Context
requestContext = 
  Vault.lookup contextKey . 
  vault