{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

{-# OPTIONS -Wno-orphans #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE InstanceSigs #-}
module Control.Trace.Core (
  Tracer (..),
  MonadTrace (..),
  MonadMaskInstance (..),
  logMsg,
  traceMsg,
  traceIf,
  Contravariant,
  (>$<),
) where

import Control.Monad (when)
import Control.Monad.Catch (
  ExitCase (..),
  MonadCatch,
  MonadMask (generalBracket),
  MonadThrow,
 )
import Control.Monad.IO.Class (
  MonadIO (..),
 )
import Data.Coerce
import Data.Functor.Contravariant (
  Contravariant (contramap),
  (>$<),
 )
import GHC.Stack (
  HasCallStack,
  withFrozenCallStack,
 )

-- | A contravariant tracing abstraction
data Tracer msg = Tracer
  { -- | Log a message
    logMsg_ :: msg -> IO ()
  , -- | Starts a trace and returns an action to end it
    traceMsg_
      :: forall a. HasCallStack => msg -> IO (ExitCase a -> IO ())
  }

logMsg :: (HasCallStack, MonadIO m) => Tracer msg -> msg -> m ()
logMsg logger msg = withFrozenCallStack $ liftIO $ logMsg_ logger msg

traceMsg ::
  (HasCallStack, MonadTrace m) => Tracer msg -> msg -> m a -> m a
traceMsg logger msg act = withFrozenCallStack $
  bracketM (traceMsg_ logger msg) id (const act)

instance Contravariant Tracer where
  contramap f (Tracer logf traceF) = Tracer (logf . f) (traceF . f)

instance Monoid (Tracer msg) where
  mempty = Tracer (\_ -> pure ()) (const $ pure $ const $ pure ())

instance Semigroup (Tracer msg) where
  l1 <> l2 =
    Tracer
      { logMsg_ = \m -> logMsg_ l1 m *> logMsg_ l2 m
      , traceMsg_ = \msg -> do
        end1 <- traceMsg_ l1 msg
        end2 <- traceMsg_ l2 msg
        return (\res -> end2 res >> end1 res)
      }

--------------------------------------------------------------------------------
-- useful combinators

-- | Gate every trace behind a condition
traceIf :: forall msg. IO Bool -> Tracer msg -> Tracer msg
traceIf cond tracer =
  let
    logMsg' msg = do
      value <- cond
      when value $ logMsg_ tracer msg
    traceMsg' :: msg -> IO (ExitCase b -> IO ())
    traceMsg' msg = do
      value <- cond
      if value
        then traceMsg_ tracer msg
        else pure $ pure $ pure ()
  in Tracer logMsg' traceMsg'

--------------------------------------------------------------------------------
-- A Monad for 'bracket'

class MonadIO m => MonadTrace m where
  bracketM :: IO a -> (a -> ExitCase b -> IO ()) -> (a -> m b) -> m b

-- deriving via (MonadMaskInstance IO) instance MonadTrace IO
instance MonadTrace IO where
  bracketM
    :: forall a b . IO a -> (a -> ExitCase b -> IO ()) -> (a -> IO b) -> IO b
  bracketM = coerce (bracketM @(MonadMaskInstance IO) @a @b)

-- | Deriving 'MonadTrace' via 'MonadMask'
newtype MonadMaskInstance m a = MonadMaskInstance (m a)
  deriving
    (Applicative, Functor, Monad, MonadCatch, MonadIO, MonadMask, MonadThrow)

instance (MonadIO m, MonadMask m) => MonadTrace (MonadMaskInstance m) where
  bracketM acquire release =
    fmap fst . generalBracket (liftIO acquire) ((liftIO .) . release)