{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -Wall -Werror #-}
{-# OPTIONS_GHC -Wno-missing-fields #-}

-- | A simple logger that reads `LogEntry` from a `TBChan` and dumps JSON-formatted
-- strings to `stdout`.
module Minilang.Log
  ( -- * Types
    LoggerEnv,

    -- * Constructor & Destructor
    newLog,
    stopLogger,
    fakeLogger,

    -- * Logging functions
    logInfo,
    logError,
    withLog,
  )
where

import Control.Concurrent.Async (Async, async, cancel)
import Control.Concurrent.STM (TChan, atomically, newTChan, readTChan, writeTChan)
import Control.Monad (forever)
import Control.Monad.Trans (MonadIO (..))
import Data.Aeson
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as LBS
import Data.Text (Text)
import Data.Time.Clock
  ( UTCTime,
    diffUTCTime,
    getCurrentTime,
  )
import System.IO (Handle, stdout, hFlush)

-- | Environment to control a logger thread.
data LoggerEnv m = LoggerEnv
  { logger :: Maybe Logger,
    loggerId :: Text,
    logInfo :: forall a. (ToJSON a) => a -> m (),
    logError :: forall a. (ToJSON a) => a -> m (),
    withLog :: forall a b. (ToJSON a) => a -> m b -> m b,
    stopLogger :: m ()
  }

type Logger = TChan BS.ByteString

fakeLogger :: (MonadIO m) => LoggerEnv m
fakeLogger = LoggerEnv Nothing "foo" (const $ pure ()) (const $ pure ()) (const id) (pure ())

-- | Starts an asynchronous log-processing thread and returns an initialised `LoggerEnv`.
newLog :: (MonadIO m) => Text -> m (LoggerEnv m)
newLog loggerId = liftIO $ do
  chan <- liftIO $ atomically newTChan
  loggerThread <- async $ runLog chan stdout
  return $
    LoggerEnv
      { logger = Just chan,
        logInfo = logEvent' chan loggerId,
        logError = logError' chan loggerId,
        withLog = withLog' chan loggerId,
        stopLogger = stopLogger' loggerThread,
        loggerId
      }

runLog :: TChan BS.ByteString -> Handle -> IO a
runLog chan hdl = forever $ do
  toLog <- atomically $ readTChan chan
  BS.hPutStr hdl . (<> "\n") $ toLog
  hFlush hdl

stopLogger' :: (MonadIO m) => Async () -> m ()
stopLogger' = liftIO . cancel

logEvent' :: (ToJSON a, MonadIO m) => Logger -> Text -> a -> m ()
logEvent' chan logId message = do
  ts <- liftIO getCurrentTime
  liftIO $
    atomically $
      writeTChan chan $
        LBS.toStrict $
          encode $
            object
              [ "timestamp" .= ts,
                "logger" .= logId,
                "message" .= message
              ]

withLog' :: (ToJSON a, MonadIO m) => Logger -> Text -> a -> m b -> m b
withLog' chan logId message act = do
  startTime <- liftIO getCurrentTime
  logStart chan logId startTime message
  b <- act
  endTime <- liftIO getCurrentTime
  logEnd chan logId startTime endTime message
  pure b

logStart :: (ToJSON a, MonadIO m) => Logger -> Text -> UTCTime -> a -> m ()
logStart chan logId ts command =
  liftIO $
    atomically $
      writeTChan chan $
        LBS.toStrict $
          encode $
            object
              [ "timestamp" .= ts,
                "logger" .= logId,
                "message" .= command
              ]

logEnd :: (ToJSON a, MonadIO m) => Logger -> Text -> UTCTime -> UTCTime -> a -> m ()
logEnd chan logId ts en outcome =
  liftIO $
    atomically $
      writeTChan chan $
        LBS.toStrict $
          encode $
            object
              [ "timestamp" .= en,
                "durationMs" .= (diffUTCTime en ts * 1000),
                "logger" .= logId,
                "message" .= outcome
              ]

logError' :: (ToJSON a, MonadIO m) => Logger -> Text -> a -> m ()
logError' chan logId err = liftIO $ do
  ts <- liftIO getCurrentTime
  liftIO $
    atomically $
      writeTChan chan $
        LBS.toStrict $
          encode $
            object
              [ "timestamp" .= ts,
                "logger" .= logId,
                "error" .= err
              ]