{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE RecordWildCards #-}

-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2025 Wire Swiss GmbH <opensource@wire.com>
--
-- This program is free software: you can redistribute it and/or modify it under
-- the terms of the GNU Affero General Public License as published by the Free
-- Software Foundation, either version 3 of the License, or (at your option) any
-- later version.
--
-- This program is distributed in the hope that it will be useful, but WITHOUT
-- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
-- FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
-- details.
--
-- You should have received a copy of the GNU Affero General Public License along
-- with this program. If not, see <https://www.gnu.org/licenses/>.

module Wire.BackendNotificationPusher where

import Control.Arrow
import Control.Monad.Catch
import Control.Retry
import Data.Aeson qualified as A
import Data.Domain
import Data.Id
import Data.List.NonEmpty qualified as NE
import Data.Map qualified as M
import Data.Map.Strict qualified as Map
import Data.Set qualified as Set
import Data.Text qualified as Text
import Imports
import Network.AMQP qualified as Q
import Network.AMQP.Extended
import Network.AMQP.Lifted qualified as QL
import Network.AMQP.Types qualified as QL
import Network.RabbitMqAdmin hiding (adminClient)
import Network.RabbitMqAdmin qualified as RabbitMqAdmin
import Network.Wai.Utilities.Error
import Prometheus
import Servant.Client qualified as Servant
import System.Logger.Class qualified as Log
import UnliftIO
import Wire.API.Federation.API
import Wire.API.Federation.BackendNotifications
import Wire.API.Federation.Client
import Wire.API.Federation.Error
import Wire.API.Federation.Version
import Wire.API.RawJson
import Wire.BackgroundWorker.Env
import Wire.BackgroundWorker.Options
import Wire.BackgroundWorker.Util

startPushingNotifications ::
  MVar () ->
  Q.Channel ->
  Domain ->
  AppT IO Q.ConsumerTag
startPushingNotifications runningFlag chan domain = do
  lift $ ensureQueue chan domain._domainText
  QL.consumeMsgs'
    chan
    (routingKey domain._domainText)
    Q.Ack
    (void . pushNotification runningFlag domain)
    ( \consumerTag ->
        Log.err
          ( Log.msg (Log.val "Got basic.cancel message for queue. Cancel queue consumer.")
              . Log.field "domain" (domainText domain)
              . Log.field "consumer" consumerTag
          )
    )
    (QL.FieldTable M.empty)

pushNotification :: (RabbitMQEnvelope e) => MVar () -> Domain -> (Q.Message, e) -> AppT IO (Async ())
pushNotification runningFlag targetDomain (msg, envelope) = do
  cfg <- asks (.backendNotificationsConfig)
  -- Jittered exponential backoff with 10ms as starting delay and 300s as max
  -- delay. When 300s is reached, every retry will happen after 300s.
  --
  -- FUTUREWORK: Pull these numbers into config.s
  let policy = capDelay cfg.pushBackoffMaxWait $ fullJitterBackoff cfg.pushBackoffMinWait

      logErr :: Bool -> SomeException -> RetryStatus -> AppT IO ()
      logErr willRetry (SomeException e) rs = do
        Log.err $
          Log.msg (Log.val "Exception occurred while pushing notification")
            . Log.field "error" (displayException e)
            . Log.field "domain" (domainText targetDomain)
            . Log.field "willRetry" willRetry
            . Log.field "retryCount" rs.rsIterNumber
        metrics <- asks backendNotificationMetrics
        withLabel metrics.errorCounter (domainText targetDomain) incCounter
        withLabel metrics.stuckQueuesGauge (domainText targetDomain) (flip setGauge (if willRetry then 1 else 0))

      skipChanThreadKilled :: RetryStatus -> Handler (AppT IO) Bool
      skipChanThreadKilled _ = Handler $ \(_ :: Q.ChanThreadKilledException) -> pure False

      shouldRetry :: SomeException -> AppT IO Bool
      shouldRetry e = do
        case fromException @FederatorClientError e of
          Just (FederatorClientError err) -> do
            case err.innerError of
              -- if a message cannot be parsed by the receiving backend, we should not retry because we cannot recover from this
              Just innerError | isBadRequest innerError && innerError.label == "bad-request" -> do
                Log.err $ Log.msg (Log.val "Dropping notification due to bad request")
                lift $ reject envelope False
                pure False
              _ -> pure True
          _ -> pure True

      handlers =
        skipAsyncExceptions
          <> [ skipChanThreadKilled,
               logRetries shouldRetry logErr
             ]
  -- The revcovering policy where it can loop forever effectively blocks the consumer thread.
  -- This isn't a problem for single active consumer with single message delivery, however it
  -- does cause problems when trying to deregister consumers from the channel. This is because
  -- the internal mechanism to remove a consumer goes via the same notification handling code
  -- as messages from the Rabbit server. If the thread is tied up in the recovery code we
  -- can't cancel the consumer, and the calling code will block until the cancellation message
  -- can be processed.
  -- Luckily, we can async this loop and carry on as usual due to how we have the channel setup.
  async $
    recovering policy handlers $
      const $
        -- Ensure that the mvars are reset correctly.
        -- takeMVar also has the nice feature of being a second layer of protection
        -- against lazy thread updates in `amqp`. If this somehow gets called while
        -- we are trying to cleanup workers for a shutdown, this will call will block
        -- and prevent the message from being sent out as we are tearing down resources.
        -- This removes one way that a message might be delivered twice.
        UnliftIO.bracket_ (takeMVar runningFlag) (putMVar runningFlag ()) go
  where
    go :: AppT IO ()
    go = case A.eitherDecode @(PayloadBundle 'Brig) (Q.msgBody msg) of
      Left e -> do
        case A.eitherDecode @BackendNotification (Q.msgBody msg) of
          Left eBN -> do
            Log.err $
              Log.msg
                ( Log.val "Cannot parse a queued message as a notification "
                    <> "nor as a bundle; the message will be ignored"
                )
                . Log.field "domain" (domainText targetDomain)
                . Log.field "error-notification" eBN
                . Log.field
                  "error-bundle"
                  e
            -- FUTUREWORK: This rejects the message without any requeueing. This is
            -- dangerous as it could happen that a new type of notification is
            -- introduced and an old instance of this worker is running, in which case
            -- the notification will just get dropped. On the other hand not dropping
            -- this message blocks the whole queue. Perhaps there is a better way to
            -- deal with this.
            lift $ reject envelope False
          Right notif -> do
            -- FUTUREWORK: Drop support for parsing it as a
            -- single notification as soon as we can guarantee
            -- that the message queue does not contain any
            -- 'BackendNotification's anymore.
            ceFederator <- asks (.federatorInternal)
            ceHttp2Manager <- asks http2Manager
            let ceOriginDomain = notif.ownDomain
                ceTargetDomain = targetDomain
                ceOriginRequestId = fromMaybe (RequestId defRequestId) notif.requestId
                cveEnv = FederatorClientEnv {..}
                cveVersion = Just V0 -- V0 is assumed for non-versioned queue messages
                fcEnv = FederatorClientVersionedEnv {..}
            sendNotificationIgnoringVersionMismatch fcEnv notif.targetComponent notif.path notif.body
            lift $ ack envelope
            metrics <- asks backendNotificationMetrics
            withLabel metrics.pushedCounter (domainText targetDomain) incCounter
            withLabel metrics.stuckQueuesGauge (domainText targetDomain) (flip setGauge 0)
      Right bundle -> do
        federator <- asks (.federatorInternal)
        manager <- asks http2Manager
        let env =
              FederatorClientEnv
                { ceOriginDomain = ownDomain . NE.head $ bundle.notifications,
                  ceTargetDomain = targetDomain,
                  ceFederator = federator,
                  ceHttp2Manager = manager,
                  ceOriginRequestId =
                    fromMaybe (RequestId defRequestId) . (.requestId) . NE.head $ bundle.notifications
                }
        remoteVersions :: Set Int <-
          liftIO
            -- use versioned client with no version set: since we are manually
            -- performing version negotiation, we don't want the client to
            -- negotiate a version for us
            ( runVersionedFederatorClient @'Brig (unversionedEnv env) $
                fedClientIn @'Brig @"api-version" ()
            )
            >>= \case
              Left e -> do
                Log.err $
                  Log.msg (Log.val "Failed to get supported API versions")
                    . Log.field "domain" (domainText targetDomain)
                    . Log.field "error" (displayException e)
                throwM e
              Right vi -> pure . Set.fromList . vinfoSupported $ vi

        -- compute the best usable version in a notification
        let bestVersion = bodyVersions >=> flip latestCommonVersion remoteVersions
        case pairedMaximumOn bestVersion (toList (notifications bundle)) of
          (_, Nothing) ->
            Log.fatal $
              Log.msg (Log.val "No federation API version in common, the notification will be ignored")
                . Log.field "domain" (domainText targetDomain)
          (notif, cveVersion) -> do
            ceFederator <- asks (.federatorInternal)
            ceHttp2Manager <- asks http2Manager
            let ceOriginDomain = notif.ownDomain
                ceTargetDomain = targetDomain
                ceOriginRequestId = fromMaybe (RequestId defRequestId) notif.requestId
                cveEnv = FederatorClientEnv {..}
                fcEnv = FederatorClientVersionedEnv {..}
            sendNotificationIgnoringVersionMismatch fcEnv notif.targetComponent notif.path notif.body
            lift $ ack envelope
            metrics <- asks backendNotificationMetrics
            withLabel metrics.pushedCounter (domainText targetDomain) incCounter
            withLabel metrics.stuckQueuesGauge (domainText targetDomain) (flip setGauge 0)

sendNotificationIgnoringVersionMismatch ::
  FederatorClientVersionedEnv ->
  Component ->
  Text ->
  RawJson ->
  AppT IO ()
sendNotificationIgnoringVersionMismatch env comp path body =
  liftIO (sendNotification env comp path body) >>= \case
    Left (FederatorClientVersionNegotiationError v) -> do
      Log.fatal $
        Log.msg (Log.val "Federator version negotiation error")
          . Log.field "domain" (domainText env.cveEnv.ceTargetDomain)
          . Log.field "error" (show v)
      pure ()
    Left e -> throwM e
    Right () -> pure ()

-- | Find the pair that maximises b.
pairedMaximumOn :: (Ord b) => (a -> b) -> [a] -> (a, b)
pairedMaximumOn f = maximumBy (compare `on` snd) . map (id &&& f)

-- FUTUREWORK: Recosider using 1 channel for many consumers. It shouldn't matter
-- for a handful of remote domains.
-- Consumers is passed in explicitly so that cleanup code has a reference to the consumer tags.
startPusher :: RabbitMqAdmin.AdminAPI (Servant.AsClientT IO) -> IORef (Map Domain (Q.ConsumerTag, MVar ())) -> Q.Channel -> AppT IO ()
startPusher adminClient consumersRef chan = do
  markAsWorking BackendNotificationPusher
  -- This ensures that we receive notifications 1 by 1 which ensures they are
  -- delivered in order.
  lift $ Q.qos chan 0 1 False
  -- Make sure threads aren't dangling if/when this async thread is killed
  let cleanup :: (Exception e, MonadThrow m, MonadIO m) => e -> m ()
      cleanup e = do
        consumers <- liftIO $ readIORef consumersRef
        traverse_ (liftIO . Q.cancelConsumer chan . fst) $ Map.elems consumers
        throwM e
  timeBeforeNextRefresh <- asks (.backendNotificationsConfig.remotesRefreshInterval)
  -- If this thread is cancelled, catch the exception, kill the consumers, and carry on.
  -- FUTUREWORK?:
  -- If this throws an exception on the Chan / in the forever loop, the exception will
  -- bubble all the way up and kill the pod. Kubernetes should restart the pod automatically.
  flip
    UnliftIO.catches
    [ Handler $ cleanup @SomeException,
      Handler $ cleanup @SomeAsyncException
    ]
    $ forever
    $ do
      remotes <- getRemoteDomains adminClient
      ensureConsumers consumersRef chan remotes
      threadDelay timeBeforeNextRefresh

ensureConsumers :: IORef (Map Domain (Q.ConsumerTag, MVar ())) -> Q.Channel -> [Domain] -> AppT IO ()
ensureConsumers consumers chan domains = do
  keys' <- Set.fromList . Map.keys <$> readIORef consumers
  let domains' = Set.fromList domains
      droppedDomains = Set.difference keys' domains'
  -- Loop over all of the new domains. We can check for existing consumers and add new ones.
  traverse_ (ensureConsumer consumers chan) domains
  -- Loop over all of the dropped domains. These need to be cancelled as they are no longer
  -- on the domain list.
  traverse_ (cancelConsumer consumers chan) droppedDomains

cancelConsumer :: IORef (Map Domain (Q.ConsumerTag, MVar ())) -> Q.Channel -> Domain -> AppT IO ()
cancelConsumer consumers chan domain = do
  Log.info $ Log.msg (Log.val "Stopping consumer") . Log.field "domain" (domainText domain)
  -- The ' version of atomicModifyIORef is strict in the function update and is useful
  -- for not leaking memory.
  atomicModifyIORef' consumers (\c -> (Map.delete domain c, Map.lookup domain c))
    >>= liftIO . traverse_ (Q.cancelConsumer chan . fst)

ensureConsumer :: IORef (Map Domain (Q.ConsumerTag, MVar ())) -> Q.Channel -> Domain -> AppT IO ()
ensureConsumer consumers chan domain = do
  consumerExists <- Map.member domain <$> readIORef consumers
  unless consumerExists $ do
    Log.info $ Log.msg (Log.val "Starting consumer") . Log.field "domain" (domainText domain)
    -- Build an MVar for the consumer. This is used as a flag for when a consumer callback is running.
    -- The cleanup code that runs when the service receives a SIGTERM or SIGINT will wait on these MVars
    -- to allow current messages to finish processing before we close AMQP connections.
    runningFlag <- newMVar ()
    tag <- startPushingNotifications runningFlag chan domain
    oldTag <- atomicModifyIORef consumers $ \c -> (Map.insert domain (tag, runningFlag) c, Map.lookup domain c)
    -- This isn't strictly nessacary. `unless consumerExists` won't
    -- let us come down this path if there is an old consumer.
    liftIO $ forM_ oldTag $ Q.cancelConsumer chan . fst

getRemoteDomains :: RabbitMqAdmin.AdminAPI (Servant.AsClientT IO) -> AppT IO [Domain]
getRemoteDomains adminClient = do
  -- Jittered exponential backoff with 10ms as starting delay and 60s as max
  -- cumulative delay. When this is reached, the operation fails.
  --
  -- FUTUREWORK: Pull these numbers into config
  let policy = limitRetriesByCumulativeDelay 60_000_000 $ fullJitterBackoff 10000
      logErrr willRetry (SomeException e) rs =
        Log.err $
          Log.msg (Log.val "Exception occurred while refreshing domains")
            . Log.field "error" (displayException e)
            . Log.field "willRetry" willRetry
            . Log.field "retryCount" rs.rsIterNumber
      handlers =
        skipAsyncExceptions
          <> [logRetries (const $ pure True) logErrr]
  recovering policy handlers $ const $ go [] 1
  where
    go :: [Domain] -> Int -> AppT IO [Domain]
    go domains pageNumber = do
      vhost <- asks rabbitmqVHost
      queuesPage <- liftIO $ listQueuesByVHost adminClient vhost "^backend-notifications\\..*" True 100 pageNumber
      let notifQueuesSuffixes = mapMaybe (\q -> Text.stripPrefix "backend-notifications." q.name) queuesPage.items
      newDomains <- catMaybes <$> traverse (\d -> either (\e -> logInvalidDomain d e >> pure Nothing) (pure . Just) $ mkDomain d) notifQueuesSuffixes
      let domainsSoFar = newDomains <> domains
      if queuesPage.page >= queuesPage.pageCount
        then pure domainsSoFar
        else go domainsSoFar (pageNumber + 1)
    logInvalidDomain d e =
      Log.warn $
        Log.msg (Log.val "Found invalid domain in a backend notifications queue name")
          . Log.field "queue" ("backend-notifications." <> d)
          . Log.field "error" e

startWorker :: AmqpEndpoint -> AppT IO CleanupAction
startWorker rabbitmqOpts = do
  env <- ask
  -- These are used in the POSIX signal handlers, so we need to make
  -- cross thread references that we can use to cancel consumers and
  -- wait for current processing steps to finish.
  chanRef <- newIORef Nothing
  consumersRef <- newIORef mempty
  let -- cleanup the refs when channels die
      -- This is so we aren't trying to close consumers
      -- that don't exist when the service is shutdown.
      clearRefs = do
        atomicWriteIORef chanRef Nothing
        atomicWriteIORef consumersRef mempty
  case env.rabbitmqAdminClient of
    Nothing ->
      Log.info $
        Log.msg $
          Log.val "RabbitMQ admin client not available, skipping backend notification pusher."
    Just client -> do
      markAsNotWorking BackendNotificationPusher
      -- We can fire and forget this thread because it keeps respawning itself using the 'onConnectionClosedHandler'.
      void $
        async $
          liftIO $
            openConnectionWithRetries env.logger rabbitmqOpts (Just "backend-notification-pusher") $
              RabbitMqHooks
                { -- The exception handling in `openConnectionWithRetries` won't open a new
                  -- connection on an explicit close call.
                  onNewChannel = \chan -> do
                    atomicWriteIORef chanRef $ pure chan
                    runAppT env $ startPusher client consumersRef chan,
                  onChannelException = \_ -> do
                    clearRefs
                    runAppT env $ markAsNotWorking BackendNotificationPusher,
                  onConnectionClose = do
                    clearRefs
                    runAppT env $ markAsNotWorking BackendNotificationPusher
                }
  pure $ runAppT env $ cleanupWorker chanRef consumersRef

cleanupWorker :: IORef (Maybe Q.Channel) -> IORef (Map Domain (Q.ConsumerTag, MVar ())) -> AppT IO ()
cleanupWorker notifChanRef notifConsumersRef = do
  Log.info $ Log.msg (Log.val "Cancelling the notification pusher thread")
  readIORef notifChanRef >>= traverse_ \chan -> do
    Log.info $ Log.msg (Log.val "Got channel")
    readIORef notifConsumersRef >>= \m -> for_ (Map.assocs m) \(domain, (consumer, runningFlag)) -> do
      Log.info $ Log.msg (Log.val "Cancelling consumer") . Log.field "Domain" domain._domainText
      -- Remove the consumer from the channel so it isn't called again
      liftIO $ Q.cancelConsumer chan consumer
      -- Take from the mvar. This will only unblock when the consumer callback isn't running.
      -- This allows us to wait until the currently running tasks are completed, and new ones
      -- won't be scheduled because we've already removed the callback from the channel.
      -- If, for some reason, a consumer is invoked after us cancelling it, taking this MVar
      -- will block that thread from trying to push out the notification. At this point, we're
      -- relying on Rabbit to requeue the message for us as we won't be able to ACK or NACK it.
      -- This helps prevent message redelivery to endpoint services during the brief window between
      -- receiving a message from rabbit, and the signal handler shutting down the AMQP connection
      -- before notification delivery has finalised.
      Log.info $ Log.msg $ Log.val "Taking MVar. Waiting for current operation to finish"
      takeMVar runningFlag
    -- Close the channel. `extended` will then close the connection, flushing messages to the server.
    Log.info $ Log.msg $ Log.val "Closing RabbitMQ channel"
    liftIO $ Q.closeChannel chan
