-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2022 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 Gundeck.Push.Websocket
  ( push,
    bulkPush,
    MonadBulkPush (..),
  )
where

import Bilge hiding (trace)
import Bilge.RPC
import Bilge.Retry (rpcHandlers)
import Control.Arrow ((&&&))
import Control.Exception (ErrorCall (ErrorCall))
import Control.Lens (view, (%~), (^.), _2)
import Control.Monad.Catch (MonadMask, MonadThrow, catch, throwM, try)
import Control.Retry
import Data.Aeson (eitherDecode, encode)
import Data.ByteString.Conversion
import Data.ByteString.Lazy qualified as L
import Data.Id
import Data.List.NonEmpty (NonEmpty)
import Data.Map qualified as Map
import Data.Misc (Milliseconds (..))
import Data.Set qualified as Set
import Data.Time.Clock.POSIX
import Gundeck.Monad
import Gundeck.Presence.Data qualified as Presence
import Gundeck.Util
import Imports
import Network.HTTP.Client (HttpExceptionContent (..))
import Network.HTTP.Client.Internal qualified as Http
import Network.HTTP.Types (StdMethod (POST), status200, status410)
import Network.URI qualified as URI
import Prometheus qualified as Prom
import System.Logger.Class (val, (+++), (~~))
import System.Logger.Class qualified as Log
import UnliftIO (handleAny, mapConcurrently)
import Wire.API.Internal.BulkPush
import Wire.API.Internal.Notification
import Wire.API.Presence

class (Monad m, MonadThrow m, Log.MonadLogger m) => MonadBulkPush m where
  mbpBulkSend :: URI -> BulkPushRequest -> m (URI, Either SomeException BulkPushResponse)
  mbpDeleteAllPresences :: [Presence] -> m ()
  mbpPosixTime :: m Milliseconds
  mbpMapConcurrently :: (Traversable t) => (a -> m b) -> t a -> m (t b)
  mbpMonitorBadCannons :: (URI, (SomeException, [Presence])) -> m ()

instance MonadBulkPush Gundeck where
  mbpBulkSend = bulkSend
  mbpDeleteAllPresences = runWithAdditionalRedis . Presence.deleteAll
  mbpPosixTime = posixTime
  mbpMapConcurrently = mapConcurrently
  mbpMonitorBadCannons = monitorBadCannons

-- | Send a 'Notification's to associated 'Presence's.  Send at most one request to each Cannon.
-- Return the lists of 'Presence's successfully reached for each resp. 'Notification'.
bulkPush :: forall m. (MonadBulkPush m) => [(Notification, [Presence])] -> m [(NotificationId, [Presence])]
-- REFACTOR: make presences lists (and notification list) non-empty where applicable?  are there
-- better types to express more of our semantics / invariants?  (what about duplicates in presence
-- lists?)
bulkPush notifs = do
  let reqs = fanOut notifs
  flbck <- flowBack <$> (uncurry mbpBulkSend `mbpMapConcurrently` reqs)
  let -- lookup by 'URI' can fail iff we screwed up URI handling in this module.
      presencesByCannon = mkPresencesByCannon . mconcat $ snd <$> notifs
      -- lookup by 'PushTarget' can fail iff Cannon sends an invalid key.
      presenceByPushTarget = mkPresenceByPushTarget . mconcat $ snd <$> notifs
  badCannons :: [(URI, (SomeException, [Presence]))] <-
    forM (flowBackBadCannons flbck) $ \(uri, e) -> (uri,) . (e,) <$> presencesByCannon uri
  prcsGone :: [Presence] <-
    presenceByPushTarget `mapM` flowBackLostPrcs flbck
  successes :: [(NotificationId, Presence)] <-
    (\(nid, trgt) -> (nid,) <$> presenceByPushTarget trgt) `mapM` flowBackDelivered flbck
  (\info -> mbpMonitorBadCannons info >> logBadCannons info) `mapM_` badCannons
  logPrcsGone `mapM_` prcsGone
  logSuccesses `mapM_` successes
  mbpDeleteAllPresences =<< do
    now <- mbpPosixTime
    let deletions = prcsGone <> (filter dead . mconcat $ snd . snd <$> badCannons)
        dead prc = now - createdAt prc > 10 * posixDay
        posixDay = Ms (round (1000 * posixDayLength))
    pure deletions
  pure (groupAssoc successes)

-- | log all cannons with response status @/= 200@.
monitorBadCannons ::
  (Prom.MonadMonitor m) =>
  (uri, (error, [Presence])) ->
  m ()
monitorBadCannons (_uri, (_err, prcs)) =
  void $ Prom.addCounter pushWsUnreachableCounter (fromIntegral $ length prcs)

{-# NOINLINE pushWsUnreachableCounter #-}
pushWsUnreachableCounter :: Prom.Counter
pushWsUnreachableCounter =
  Prom.unsafeRegister $
    Prom.counter
      Prom.Info
        { Prom.metricName = "push_ws_unreachable",
          Prom.metricHelp = "Number of times websocket pushes were not pushed due cannon being unreachable"
        }

logBadCannons :: (Log.MonadLogger m) => (URI, (SomeException, [Presence])) -> m ()
logBadCannons (uri, (err, prcs)) = do
  forM_ prcs $ \prc ->
    Log.warn $
      logPresence prc
        ~~ Log.field "created_at" (ms $ createdAt prc)
        ~~ Log.field "cannon_uri" (show uri)
        ~~ Log.field "resource_target" (show $ resource prc)
        ~~ Log.field "http_exception" (intercalate " | " . lines . displayException $ err)
        ~~ Log.msg (val "WebSocket presence unreachable: ")

logPrcsGone :: (Log.MonadLogger m) => Presence -> m ()
logPrcsGone prc = Log.debug $ logPresence prc ~~ Log.msg (val "WebSocket presence gone")

logSuccesses :: (Log.MonadLogger m) => (a, Presence) -> m ()
logSuccesses (_, prc) = Log.debug $ logPresence prc ~~ Log.msg (val "WebSocket push success")

fanOut :: [(Notification, [Presence])] -> [(URI, BulkPushRequest)]
fanOut =
  fmap (_2 %~ (mkBulkPushRequest . groupByNotification))
    . groupByURI
    . mconcat
    . fmap pullUri
  where
    mkBulkPushRequest :: [(Notification, [Presence])] -> BulkPushRequest
    mkBulkPushRequest = BulkPushRequest . fmap (_2 %~ fmap mkPushTarget)
    groupByNotification :: [(Notification, Presence)] -> [(Notification, [Presence])]
    groupByNotification = groupAssoc' (compare `on` ntfId)
    groupByURI :: [(Notification, (URI, Presence))] -> [(URI, [(Notification, Presence)])]
    groupByURI = groupAssoc . fmap (\(notif, (uri, prc)) -> (uri, (notif, prc)))
    pullUri :: (notif, [Presence]) -> [(notif, (URI, Presence))]
    pullUri (notif, prcs) = (notif,) . (bulkresource &&& id) <$> prcs

bulkSend ::
  forall m.
  ( MonadMask m,
    HasRequestId m,
    MonadHttp m,
    MonadUnliftIO m,
    Log.MonadLogger m
  ) =>
  URI ->
  BulkPushRequest ->
  m (URI, Either SomeException BulkPushResponse)
bulkSend uri req = (uri,) <$> ((Right <$> bulkSend' uri req) `catch` (pure . Left))

bulkSend' ::
  forall m.
  ( MonadUnliftIO m,
    MonadMask m,
    HasRequestId m,
    MonadHttp m,
    Log.MonadLogger m
  ) =>
  URI ->
  BulkPushRequest ->
  m BulkPushResponse
bulkSend' uri bulkPushRequest = do
  forM_ (fromBulkPushRequest bulkPushRequest) $ \(notification, targets) ->
    Log.debug $
      Log.msg ("Bulk sending notification to Cannon." :: Text)
        . Log.field "ntf_id" (show (ntfId notification))
        . Log.field "user_ids" (show (map ptUserId targets))
        . Log.field "conn_ids" (show (map ptConnId targets))

  let jsbody = encode bulkPushRequest
  req <-
    check
      . method POST
      . contentJson
      . lbytes jsbody
      . timeout 3000 -- ms
      <$> Http.setUri empty (fromURI uri)
  try (submit req) >>= \case
    Left e -> throwM (e :: SomeException)
    Right r -> decodeBulkResp $ responseBody r
  where
    submit req = recovering (limitRetries 1) rpcHandlers $ const (rpc' "cannon" req id)
    check req =
      req
        { Http.checkResponse = \rq rs ->
            when (responseStatus rs /= status200) $
              let ex = StatusCodeException (rs {responseBody = ()}) mempty
               in throwM $ HttpExceptionRequest rq ex
        }
    decodeBulkResp :: Maybe L.ByteString -> m BulkPushResponse
    decodeBulkResp Nothing = throwM $ ErrorCall "missing response body from cannon"
    decodeBulkResp (Just lbs) = either err pure $ eitherDecode lbs
      where
        err = throwM . ErrorCall . ("bad response body from cannon: " <>)

-- | NOTE: 'PushTarget's may occur several times both in the "lost" and in the "delivered" list.
-- This happens iff there are several 'Notifcation's for the same 'PushTarget', and some of them are
-- delivered while others aren't.
data FlowBack = FlowBack
  { -- | list of cannons that failed to respond with status 200
    flowBackBadCannons :: [(URI, SomeException)],
    -- | 401 inside the body (for one presence)
    flowBackLostPrcs :: [PushTarget],
    flowBackDelivered :: [(NotificationId, PushTarget)]
  }

flowBack :: [(URI, Either SomeException BulkPushResponse)] -> FlowBack
flowBack rawresps = FlowBack broken gone delivered
  where
    broken :: [(URI, SomeException)]
    broken =
      lefts' rawresps
    gone :: [PushTarget]
    gone =
      map (snd . snd)
        . filter
          ( \(st, _) -> case st of
              PushStatusOk -> False
              PushStatusGone -> True
          )
        $ responsive
    delivered :: [(NotificationId, PushTarget)]
    delivered =
      map snd
        . filter
          ( \(st, _) -> case st of
              PushStatusOk -> True
              PushStatusGone -> False
          )
        $ responsive
    responsive :: [(PushStatus, (NotificationId, PushTarget))]
    responsive =
      map (\(n, t, s) -> (s, (n, t)))
        . mconcat
        . fmap fromBulkPushResponse
        . rights
        $ snd <$> rawresps
    lefts' :: [(c, Either a b)] -> [(c, a)]
    lefts' [] = []
    lefts' ((c, Left x) : xs) = (c, x) : lefts' xs
    lefts' ((_, Right _) : xs) = lefts' xs

{-# INLINE mkPresencesByCannon #-}
mkPresencesByCannon :: (MonadThrow m) => [Presence] -> URI -> m [Presence]
mkPresencesByCannon prcs uri = maybe (throwM err) pure $ Map.lookup uri mp
  where
    err = ErrorCall "internal error in Gundeck: invalid URL in bulkpush result"
    mp :: Map URI [Presence]
    mp = foldl' collect mempty $ (bulkresource &&& id) <$> prcs
    collect :: Map URI [Presence] -> (URI, Presence) -> Map URI [Presence]
    collect mp' (uri', prc) = Map.alter (go prc) uri' mp'
    go :: Presence -> Maybe [Presence] -> Maybe [Presence]
    go prc Nothing = Just [prc]
    go prc (Just prcs') = Just $ prc : prcs'

{-# INLINE mkPresenceByPushTarget #-}
mkPresenceByPushTarget :: (MonadThrow m) => [Presence] -> PushTarget -> m Presence
mkPresenceByPushTarget prcs ptarget = maybe (throwM err) pure $ Map.lookup ptarget mp
  where
    err = ErrorCall "internal error in Cannon: invalid PushTarget in bulkpush response"
    mp :: Map PushTarget Presence
    mp = Map.fromList $ (mkPushTarget &&& id) <$> prcs

{-# INLINE bulkresource #-}
bulkresource :: Presence -> URI
bulkresource = URI . (\x -> x {URI.uriPath = "/i/bulkpush"}) . fromURI . resource

-- TODO: a Map-based implementation would be faster for sufficiently large inputs.  do we want to
-- take the time and benchmark the difference?  move it to types-common?
{-# INLINE groupAssoc #-}
groupAssoc :: (Ord a) => [(a, b)] -> [(a, [b])]
groupAssoc = groupAssoc' compare

-- TODO: Also should we give 'Notification' an 'Ord' instance?
{-# INLINE groupAssoc' #-}
groupAssoc' :: (Eq a) => (a -> a -> Ordering) -> [(a, b)] -> [(a, [b])]
groupAssoc' cmp =
  fmap
    ( \case
        xs@(x : _) -> (fst x, snd <$> xs)
        [] -> error "impossible: list elements returned by groupBy are never empty."
    )
    . groupBy ((==) `on` fst)
    . sortBy (cmp `on` fst)

{-# INLINE mkPushTarget #-}
mkPushTarget :: Presence -> PushTarget
mkPushTarget pre = PushTarget (userId pre) (connId pre)

-----------------------------------------------------------------------------
-- old, multi-request push.

push ::
  Notification ->
  NonEmpty NotificationTarget ->
  Maybe UserId -> -- Origin user.
  Maybe ConnId -> -- Origin device connection.
  Set ConnId -> -- Only target these connections.
  Gundeck [Presence]
push notif (toList -> tgts) originUser originConn conns = do
  pp <- handleAny noPresences listPresences
  (ok, gone) <- foldM onResult ([], []) =<< send notif pp
  runWithAdditionalRedis $ Presence.deleteAll gone
  pure ok
  where
    listPresences =
      excludeOrigin
        . filterByConnection
        . concat
        . filterByClient
        . zip tgts
        <$> runWithDefaultRedis (Presence.listAll (view targetUser <$> tgts))
    noPresences exn = do
      Log.err $
        Log.field "error" (displayException exn)
          ~~ Log.msg (val "Failed to get presences.")
      pure []
    filterByClient = map $ \(tgt, ps) ->
      let cs = tgt ^. targetClients
       in if null cs
            then ps
            else filter (maybe True (`elem` cs) . clientId) ps
    filterByConnection =
      if Set.null conns
        then id
        else filter ((`Set.member` conns) . connId)
    excludeOrigin =
      let neqUser p = originUser /= Just (userId p)
          neqConn p = originConn /= Just (connId p)
       in filter (\p -> neqUser p || neqConn p)
    onResult (ok, gone) (PushSuccess p) = do
      Log.debug $ logPresence p ~~ Log.msg (val "WebSocket push success")
      pure (p : ok, gone)
    onResult (ok, gone) (PushGone p) = do
      Log.debug $ logPresence p ~~ Log.msg (val "WebSocket presence gone")
      pure (ok, p : gone)
    onResult (ok, gone) (PushFailure p _) = do
      Prom.incCounter pushWsUnreachableCounter
      Log.info $
        logPresence p
          ~~ Log.field "created_at" (ms $ createdAt p)
          ~~ Log.msg (val "WebSocket presence unreachable: " +++ toByteString (resource p))
      now <- posixTime
      if now - createdAt p > 10 * posixDay
        then pure (ok, p : gone)
        else pure (ok, gone)
    posixDay = Ms (round (1000 * posixDayLength))

-----------------------------------------------------------------------------
-- Internal

-- | Not to be confused with 'PushStatus': 'PushResult' is in internal to Gundeck, carries a
-- 'Presence', and can express HTTP errors.
data PushResult
  = PushSuccess Presence
  | PushGone Presence
  | PushFailure Presence SomeException

send :: Notification -> [Presence] -> Gundeck [PushResult]
send n pp =
  let js = encode n
   in zipWith eval pp <$> mapAsync (fn js) pp
  where
    fn js p = do
      req <- Http.setUri empty (fromURI (resource p))
      recovering x1 rpcHandlers $
        const $
          rpc' "cannon" (check req) $
            method POST
              . contentJson
              . lbytes js
              . timeout 3000 -- ms
    check r =
      r
        { Http.checkResponse = \rq rs ->
            unless (responseStatus rs `elem` [status200, status410]) $
              let ex = StatusCodeException (rs {responseBody = ()}) mempty
               in throwM $ HttpExceptionRequest rq ex
        }
    eval p (Left e) = PushFailure p e
    eval p (Right r) = if statusCode r == 200 then PushSuccess p else PushGone p
    x1 = limitRetries 1

logPresence :: Presence -> Log.Msg -> Log.Msg
logPresence p =
  Log.field "user_id" (toByteString (userId p))
    ~~ Log.field "conn_id" (toByteString (connId p))
