{-|
Module: IHP.PGListener
Description: Event listener handling pg_notify messages
Copyright: (c) digitally induced GmbH, 2021

This module is solving the problem, that previously IHP was using one database connection
per running @LISTEN ..;@ statement. A @PGListener@ provides one central object to listen on
postgres channels, without manually dealing with connection management.
-}
module IHP.PGListener
( Channel
, Callback
, Subscription (..)
, PGListener (..)
, init
, stop
, withPGListener
, subscribe
, subscribeJSON
, unsubscribe
) where

import IHP.Prelude hiding (init)
import IHP.ModelSupport
import qualified Data.Set as Set
import qualified Database.PostgreSQL.Simple as PG
import qualified Database.PostgreSQL.Simple.Types as PG
import qualified Database.PostgreSQL.Simple.Notification as PG
import qualified Data.UUID.V4 as UUID
import Control.Concurrent.MVar (MVar)
import qualified Control.Concurrent.MVar as MVar
import Data.HashMap.Strict as HashMap
import qualified Control.Concurrent.Async as Async
import qualified Data.List as List
import qualified Data.Aeson as Aeson
import qualified IHP.Log as Log
import qualified Control.Exception.Safe as Exception
import qualified Control.Concurrent.Chan.Unagi as Queue
import qualified Control.Concurrent

-- TODO: How to deal with timeout of the connection?

-- | The channel is like the event name
--
-- It's used in the postgres NOTIFY call:
--
-- > NOTIFY channel [ , payload ]
--
type Channel = ByteString

-- | An event callback receives the postgres notification object and can do IO
type Callback = PG.Notification -> IO ()

-- | Returned by a call to 'subscribe'
data Subscription = Subscription
    { id :: !UUID
    , reader :: !(Async ())
    , inChan :: !(Queue.InChan PG.Notification)
    , channel :: !Channel
    }

-- | The main datatype of the service. Keeps tracks of all channels we're watching on, as well as all open subscriptions
--
-- Use 'init' to create a new object and 'stop' to deallocate it.
data PGListener = PGListener
    { modelContext :: !ModelContext
    , listeningTo :: !(MVar (Set Channel))
    , listenTo :: !(MVar Channel)
    , subscriptions :: !(IORef (HashMap Channel [Subscription]))
    , notifyLoopAsync :: !(Async ())
    }

-- | Creates a new 'PGListener' object
--
-- > let modelContext = ..
-- > pgListener <- PGListener.init modelContext
--
-- This will start a new async listening for postgres notifications. This will take one connection
-- from the database pool and keep it blocked until 'stop' is called.
--
init :: ModelContext -> IO PGListener
init modelContext = do
    listeningTo <- MVar.newMVar Set.empty
    subscriptions <- newIORef HashMap.empty
    listenTo <- MVar.newEmptyMVar

    let ?modelContext = modelContext
    notifyLoopAsync <- async (notifyLoop listeningTo listenTo subscriptions)
    pure PGListener { modelContext, listeningTo, subscriptions, listenTo, notifyLoopAsync }

-- | Stops the database listener async and puts the database connection used back into the database pool
--
-- > PGListener.stop pgListener
--
stop :: PGListener -> IO ()
stop PGListener { notifyLoopAsync } = do
    cancel notifyLoopAsync

withPGListener :: ModelContext -> (PGListener -> IO a) -> IO a
withPGListener modelContext =
    Exception.bracket (init modelContext) stop

-- | After you subscribed to a channel, the provided callback will be called whenever there's a new
-- notification on the channel.
--
-- > pgListener <- PGListener.init
-- >
-- > let callback notification = do
-- >         let payload :: Text = cs (notification.notificationData)
-- >         putStrLn ("Received notification: " <> payload)
-- > 
-- > subscription <- pgListener |> PGListener.subscribe "my_channel" callback
--
-- The @callback@ function will now be called whenever @NOTIFY "my_channel", "my payload"@ is executed on the postgres server.
--
-- When the subscription is not used anymore, call 'unsubscribe' to stop the callback from being called anymore.
--
subscribe :: Channel -> Callback -> PGListener -> IO Subscription
subscribe channel callback pgListener = do
    id <- UUID.nextRandom
    listenToChannelIfNeeded channel pgListener

    -- We use a queue here to guarantee that the messages are processed in the right order
    -- while keeping high performance.
    --
    -- A naive implementation might be just kicking of an async for each message. But in that case
    -- the messages might be delivered to the final consumer out of order.
    (inChan, outChan) <- Queue.newChan
    
    let
        -- We need to log any exception, otherwise there might be silent errors
        logException :: SomeException -> IO ()
        logException exception = logError pgListener ("Error in pg_notify handler: " <> cs (displayException exception))

    reader <- async $ forever do
            message <- Queue.readChan outChan
            callback message `Exception.catch` logException
    let subscription = Subscription { .. }

    modifyIORef' (pgListener.subscriptions) (HashMap.insertWith mappend channel [subscription] )

    pure subscription

-- | Like 'subscribe' but decodes the notification payload from JSON and passes the decoded data structure to the callback
--
-- When JSON parsing fails, this will ignore the notification.
--
-- > pgListener <- PGListener.init
-- >
-- > let callback (jsonObject :: Aeson.Value) = do
-- >         putStrLn ("Received notification: " <> tshow jsonObject)
-- > 
-- > subscription <- pgListener |> PGListener.subscribeJSON "my_json_channel" callback
--
-- The @callback@ function will now be called whenever @NOTIFY "my_json_channel", "{\"hello\":\"world\"}"@ is executed on the postgres server.
subscribeJSON :: Aeson.FromJSON jsonValue => Channel -> (jsonValue -> IO ()) -> PGListener -> IO Subscription
subscribeJSON channel callback pgListener = subscribe channel callback' pgListener
    where
        callback' notification = do
            let payload = (notification.notificationData)
            case Aeson.decodeStrict' payload of
                Just payload -> callback payload
                Nothing -> logError pgListener ("PGListener.subscribeJSON: Failed to parse " <> tshow payload)

-- | Stops the callback of a subscription from receiving further notifications
--
-- > pgListener <- PGListener.init
-- >
-- > subscription <- pgListener |> PGListener.subscribe "my_channel" callback
-- > doSomethingExpensive
-- > pgListener |> PGListener.unsubscribe subscription
--
unsubscribe :: Subscription -> PGListener -> IO ()
unsubscribe subscription@(Subscription { .. }) pgListener = do
    let
        deleteById :: [Subscription] -> [Subscription]
        deleteById = List.deleteBy (\a b -> a.id == b.id) subscription
    modifyIORef' (pgListener.subscriptions) (HashMap.adjust deleteById channel)
    uninterruptibleCancel reader
    pure ()     

-- | Runs a @LISTEN ..;@ statements on the postgres connection, if not already listening on that channel
listenToChannelIfNeeded :: Channel -> PGListener -> IO ()
listenToChannelIfNeeded channel pgListener = do
    listeningTo <- MVar.readMVar (pgListener.listeningTo)
    let alreadyListening = channel `Set.member` listeningTo

    unless alreadyListening do
        MVar.putMVar (pgListener.listenTo) channel
            

            

-- | The main loop that is receiving events from the database and triggering callbacks
--
-- Todo: What happens when the connection dies?
notifyLoop :: (?modelContext :: ModelContext) => MVar (Set Channel) -> MVar Channel -> IORef (HashMap Channel [Subscription]) -> IO ()
notifyLoop listeningToVar listenToVar subscriptions = do
    -- Wait until the first LISTEN is requested before taking a database connection from the pool
    MVar.readMVar listenToVar

    let innerLoop = do
            withDatabaseConnection \databaseConnection -> do

                -- If listeningTo already contains channels, this means that previously the database connection
                -- died, so we're restarting here. Therefore we need to replay all LISTEN calls to restore the previous state
                listeningTo <- MVar.readMVar listeningToVar
                forEach listeningTo (listenToChannel databaseConnection)

                -- This loop reads channels from the 'listenToVar' and then triggers a LISTEN statement on
                -- the current database connections
                let listenLoop = forever do
                        channel <- MVar.takeMVar listenToVar
                        
                        MVar.modifyMVar_ listeningToVar \listeningTo -> do
                            let alreadyListening = channel `Set.member` listeningTo

                            if alreadyListening
                                then pure listeningTo
                                else do
                                    listenToChannel databaseConnection channel
                                    pure (Set.insert channel listeningTo)

                Async.withAsync listenLoop \listenLoopAsync -> do
                    forever do
                        notification <- PG.getNotification databaseConnection
                        let channel = notification.notificationChannel

                        allSubscriptions <- readIORef subscriptions
                        let channelSubscriptions = allSubscriptions
                                |> HashMap.lookup channel
                                |> fromMaybe []

                        forEach channelSubscriptions \subscription -> do
                            let inChan = subscription.inChan
                            Queue.writeChan inChan notification

    -- Initial delay (in microseconds)
    let initialDelay = 500 * 1000
    -- Max delay (in microseconds)
    let maxDelay = 60 * 1000 * 1000
    -- This outer loop restarts the listeners if the database connection dies (e.g. due to a timeout)
    let retryLoop delay isFirstError = do
            result <- Exception.tryAny innerLoop
            case result of
                Left error -> do
                    let ?context = ?modelContext -- Log onto the modelContext logger
                    if isFirstError then do
                        Log.info ("PGListener is going to restart, loop failed with exception: " <> (displayException error) <> ". Retrying immediately.")
                        retryLoop delay False -- Retry with no delay interval on first error, but will increase delay interval in subsequent retries 
                    else do
                        let increasedDelay = delay * 2 -- Double current delay
                        let nextDelay = min increasedDelay maxDelay -- Picks whichever delay is lowest of increasedDelay * 2 or maxDelay
                        Log.info ("PGListener is going to restart, loop failed with exception: " <> (displayException error) <> ". Retrying in " <> cs (printTimeToNextRetry delay) <> ".")
                        Control.Concurrent.threadDelay delay -- Sleep for the current delay
                        retryLoop nextDelay False -- Retry with longer interval
                Right _ -> 
                    retryLoop initialDelay True -- If all went well, re-run with no sleeping and reset current delay to the initial value
    retryLoop initialDelay True

printTimeToNextRetry :: Int -> Text
printTimeToNextRetry microseconds
    | microseconds >= 1000000000 =  show (microseconds `div` 1000000000) <> " min"
    | microseconds >= 1000000 =  show (microseconds `div` 1000000) <> " s"
    | microseconds >= 1000 = show (microseconds `div` 1000) <> " ms"
    | otherwise = show microseconds <> " µs"

listenToChannel :: PG.Connection -> Channel -> IO ()
listenToChannel databaseConnection channel = do
    PG.execute databaseConnection "LISTEN ?" [PG.Identifier (cs channel)]
    pure ()

logError :: PGListener -> Text -> IO ()
logError pgListener message = let ?context = pgListener.modelContext in Log.error message