{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE FlexibleInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module OpenTelemetry.Instrumentation.Persistent
  ( wrapSqlBackend
  ) where
import OpenTelemetry.Trace.Core
import OpenTelemetry.Context
import Data.Acquire.Internal
import Data.Maybe (fromMaybe)
import Data.Text (Text)
import Database.Persist.Sql
import Database.Persist.SqlBackend (setConnHooks, emptySqlBackendHooks, MkSqlBackendArgs (connRDBMS), getRDBMS, getConnVault, modifyConnVault)
import Database.Persist.SqlBackend.Internal
import Control.Monad.IO.Class
import System.IO.Unsafe (unsafePerformIO)
import qualified Data.Vault.Strict as Vault
import OpenTelemetry.Attributes (Attributes)
import OpenTelemetry.Resource
import UnliftIO.Exception
import OpenTelemetry.Trace.Monad (MonadTracer(..))
import Control.Monad.Reader
import qualified Data.Text as T
import OpenTelemetry.Context.ThreadLocal (getContext, adjustContext)

instance {-# OVERLAPS #-} MonadTracer m => MonadTracer (ReaderT SqlBackend m) where
  getTracer = lift OpenTelemetry.Trace.Monad.getTracer
instance {-# OVERLAPS #-} MonadTracer m => MonadTracer (ReaderT SqlReadBackend m) where
  getTracer = lift OpenTelemetry.Trace.Monad.getTracer
instance {-# OVERLAPS #-} MonadTracer m => MonadTracer (ReaderT SqlWriteBackend m) where
  getTracer = lift OpenTelemetry.Trace.Monad.getTracer

originalConnectionKey :: Vault.Key SqlBackend
originalConnectionKey = unsafePerformIO Vault.newKey
{-# NOINLINE originalConnectionKey #-}

insertOriginalConnection :: SqlBackend -> SqlBackend -> SqlBackend
insertOriginalConnection conn original = modifyConnVault (Vault.insert originalConnectionKey original) conn

lookupOriginalConnection :: SqlBackend -> Maybe SqlBackend
lookupOriginalConnection = Vault.lookup originalConnectionKey . getConnVault

connectionLevelAttributesKey :: Vault.Key [(Text, Attribute)]
connectionLevelAttributesKey = unsafePerformIO Vault.newKey
{-# NOINLINE connectionLevelAttributesKey #-}

-- | Wrap a 'SqlBackend' with appropriate tracing context and attributes
-- so that queries are tracked appropriately in the tracing hierarchy.
wrapSqlBackend
  :: MonadIO m
  => [(Text, Attribute)]
  -- ^ Attributes that are specific to providers like MySQL, PostgreSQL, etc.
  -> SqlBackend
  -> m SqlBackend
wrapSqlBackend attrs conn_ = do
  tp <- getGlobalTracerProvider
  let conn = Data.Maybe.fromMaybe conn_ (lookupOriginalConnection conn_)
  -- TODO add schema to tracerOptions?
  let t = makeTracer tp "hs-opentelemetry-persistent" tracerOptions
  let hooks = emptySqlBackendHooks
        { hookGetStatement = \conn sql stmt -> do
            pure $ Statement
              { stmtQuery = \ps -> do
                  ctxt <- getContext
                  let spanCreator = do
                        s <- createSpan
                          t
                          ctxt
                          sql
                          (defaultSpanArguments { kind = Client, attributes = ("db.statement", toAttribute sql) : attrs })
                        adjustContext (insertSpan s)
                        pure (lookupSpan ctxt, s)
                      spanCleanup (parent, s) = do
                        s `endSpan` Nothing
                        adjustContext $ \ctx ->
                          maybe ctx (`insertSpan` ctx) parent

                  (p, child) <- mkAcquire spanCreator spanCleanup

                  annotateBasics child conn
                  case stmtQuery stmt ps of
                    Acquire stmtQueryAcquireF -> Acquire $ \f ->
                      handleAny
                        (\(SomeException err) -> do
                          recordException child [] Nothing err
                          endSpan child Nothing
                          throwIO err
                        )
                        (stmtQueryAcquireF f)

              , stmtExecute = \ps -> do
                inSpan' t sql (defaultSpanArguments { kind = Client, attributes = ("db.statement", toAttribute sql) : attrs }) $ \s -> do
                  annotateBasics s conn
                  stmtExecute stmt ps
              , stmtReset = stmtReset stmt
              , stmtFinalize = stmtFinalize stmt
              }
        }

  let conn' = conn
        { connHooks = hooks
        , connBegin = \f mIso -> do
            let statement = "begin transaction" <> case mIso of
                  Nothing -> mempty
                  Just ReadUncommitted -> " isolation level read uncommitted"
                  Just ReadCommitted -> " isolation level read committed"
                  Just RepeatableRead -> " isolation level repeatable read"
                  Just Serializable -> " isolation level serializable"
            let attrs' = ("db.statement", toAttribute statement) : attrs
            inSpan' t statement (defaultSpanArguments { kind = Client, attributes = attrs' }) $ \s -> do
              annotateBasics s conn
              connBegin conn f mIso
        , connCommit = \f -> do
            inSpan' t "commit" (defaultSpanArguments { kind = Client, attributes = ("db.statement", toAttribute ("commit" :: Text)): attrs }) $ \s -> do
              annotateBasics s conn
              connCommit conn f
        , connRollback = \f -> do
            inSpan' t "rollback" (defaultSpanArguments { kind = Client, attributes = ("db.statement", toAttribute ("rollback" :: Text)): attrs }) $ \s -> do
              annotateBasics s conn
              connRollback conn f
        , connClose = do
            inSpan' t "close connection" (defaultSpanArguments { kind = Client, attributes = attrs }) $ \s -> do
              annotateBasics s conn
              connClose conn
        }
  pure $ insertOriginalConnection conn' conn

annotateBasics :: MonadIO m => Span -> SqlBackend -> m ()
annotateBasics span conn = do
  addAttributes span
    [ ("db.system", toAttribute $ getRDBMS conn)
    ]