{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE RecordWildCards #-}

-- 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 Cannon.WS
  ( Env (..),
    WS,
    env,
    runWS,
    drain,
    close,
    mkWebSocket,
    setRequestId,
    registerLocal,
    unregisterLocal,
    isRemoteRegistered,
    registerRemote,
    sendMsgIO,
    Clock,
    mkClock,
    getClock,
    getTime,
    Websocket,
    connection,
    connIdent,
    Key,
    mkKey,
    key2bytes,
    client,
    logKey,
    sendMsg,
  )
where

import Bilge hiding (trace)
import Bilge.RPC
import Bilge.Retry
import Cannon.Dict (Dict)
import Cannon.Dict qualified as D
import Cannon.Options (DrainOpts, WSOpts, gracePeriodSeconds, millisecondsBetweenBatches, minBatchSize)
import Cannon.RabbitMq
import Cassandra (ClientState)
import Conduit
import Control.Concurrent.Timeout
import Control.Lens ((^.))
import Control.Monad.Catch
import Control.Retry
import Data.Aeson hiding (Error, Key)
import Data.Binary.Builder qualified as B
import Data.ByteString.Char8 (pack)
import Data.ByteString.Conversion
import Data.ByteString.Lazy qualified as L
import Data.Hashable
import Data.Id (ClientId, ConnId (..), UserId, defRequestId)
import Data.List.Extra (chunksOf)
import Data.Text.Encoding (decodeUtf8)
import Data.Timeout (TimeoutUnit (..), (#))
import Data.Unique
import Imports hiding (threadDelay)
import Network.AMQP qualified as Q
import Network.HTTP.Types.Method
import Network.HTTP.Types.Status
import Network.Wai.Utilities.Error
import Network.WebSockets hiding (Request)
import System.Logger qualified as Logger
import System.Logger.Class hiding (Error, Settings, close, (.=))
import System.Random.MWC (GenIO, uniform)
import UnliftIO.Async (async, cancel, pooledMapConcurrentlyN_)
import Wire.API.Presence

-----------------------------------------------------------------------------
-- Key

newtype Key = Key
  { _key :: (ByteString, ByteString)
  }
  deriving (Eq, Show, Hashable, Ord)

mkKey :: UserId -> ConnId -> Key
mkKey u c = Key (toByteString' u, fromConnId c)

instance ToByteString Key where
  builder = B.fromByteString . key2bytes

key2bytes :: Key -> ByteString
key2bytes (Key (u, c)) = u <> "." <> c

keyUserBytes :: Key -> ByteString
keyUserBytes = fst . _key

keyConnBytes :: Key -> ByteString
keyConnBytes = snd . _key

-----------------------------------------------------------------------------
-- Websocket

data Websocket = Websocket
  { connection :: Connection,
    connIdent :: !Word
  }

mkWebSocket :: Connection -> WS Websocket
mkWebSocket c = do
  g <- WS $ asks rand
  Websocket c <$> liftIO (uniform g)

-----------------------------------------------------------------------------
-- Clock

-- | A clock that counts the number of seconds since its creation
-- to measure the uptime or lifetime of websockets.
newtype Clock = Clock (IORef Word64)

mkClock :: IO Clock
mkClock = do
  r <- newIORef 0
  void . forkIO . forever $ do
    threadDelay (1 # Second)
    modifyIORef' r (+ 1)
  pure $ Clock r

getClock :: WS Clock
getClock = WS $ asks clock

-- | Get the number of seconds elapsed since the clock was created.
getTime :: Clock -> IO Word64
getTime (Clock r) = readIORef r

-----------------------------------------------------------------------------
-- WS Monad

data Env = Env
  { externalHostname :: !ByteString,
    portnum :: !Word16,
    upstream :: !Request,
    reqId :: !RequestId,
    logg :: !Logger,
    manager :: !Manager,
    websockets :: !(Dict Key Websocket),
    rabbitConnections :: !(Dict Unique Q.Connection),
    rand :: !GenIO,
    clock :: !Clock,
    drainOpts :: DrainOpts,
    wsOpts :: WSOpts,
    cassandra :: ClientState,
    pool :: RabbitMqPool,
    notificationTTL :: Int
  }

setRequestId :: RequestId -> Env -> Env
setRequestId rid e = e {reqId = rid}

newtype WS a = WS
  { _conn :: ReaderT Env IO a
  }
  deriving
    ( Functor,
      Applicative,
      Monad,
      MonadIO,
      MonadThrow,
      MonadCatch,
      MonadMask,
      MonadReader Env,
      MonadUnliftIO
    )

instance MonadLogger WS where
  log l m = WS $ do
    g <- asks logg
    r <- field "request" . unRequestId <$> asks reqId
    liftIO $ Logger.log g l (r . m)

instance MonadHttp WS where
  handleRequestWithCont req handler = do
    manager <- asks manager
    liftIO $ withResponse req manager handler

instance HasRequestId WS where
  getRequestId = WS $ asks reqId

env ::
  ByteString ->
  Word16 ->
  ByteString ->
  Word16 ->
  Logger ->
  Manager ->
  Dict Key Websocket ->
  Dict Unique Q.Connection ->
  GenIO ->
  Clock ->
  DrainOpts ->
  WSOpts ->
  ClientState ->
  RabbitMqPool ->
  Int ->
  Env
env externalHostname portnum gundeckHost gundeckPort logg manager websockets rabbitConnections rand clock drainOpts wsOpts cassandra pool notificationTTL =
  let upstream = (Bilge.host gundeckHost . Bilge.port gundeckPort $ empty)
      reqId = RequestId defRequestId
   in Env {..}

runWS :: (MonadIO m) => Env -> WS a -> m a
runWS e m = liftIO $ runReaderT (_conn m) e

registerLocal :: Key -> Websocket -> WS ()
registerLocal k c = do
  trace $ client (key2bytes k) . msg (val "register")
  d <- WS $ asks websockets
  D.insert k c d

unregisterLocal :: Key -> Websocket -> WS Bool
unregisterLocal k c = do
  trace $ client (key2bytes k) . msg (val "unregister")
  d <- WS $ asks websockets
  D.removeIf (maybe False ((connIdent c ==) . connIdent)) k d

registerRemote :: Key -> Maybe ClientId -> WS ()
registerRemote k c = do
  let kb = key2bytes k
  debug $ client kb . msg (val "register-remote")
  e <- WS ask
  i <- regInfo k c
  void $
    recovering retry3x rpcHandlers $
      const $
        rpc' "gundeck" (upstream e) (method POST . path "/i/presences" . i . expect2xx)
  debug $ client kb . msg (val "registered")

isRemoteRegistered :: UserId -> ConnId -> WS Bool
isRemoteRegistered u c = do
  e <- WS ask
  rs <-
    recovering retry3x rpcHandlers $
      const $
        rpc' "gundeck" (upstream e) (method GET . paths ["/i/presences", toByteString' u] . expect2xx)
  cs <- map connId <$> parseResponse (mkError status502 "server-error") rs
  pure $ c `elem` cs

sendMsgIO :: (WebSocketsData a) => a -> Websocket -> IO ()
sendMsgIO m c =
  recoverAll retry3x $ const $ sendBinaryData (connection c) m

sendMsg :: (WebSocketsData a) => a -> Key -> Websocket -> WS ()
sendMsg message k c = do
  traceLog message
  liftIO $ sendMsgIO message c
  where
    traceLog :: (WebSocketsData a) => a -> WS ()
    traceLog m = trace $ client kb . msg (logMsg m)

    logMsg :: (WebSocketsData a) => a -> Builder
    logMsg m = val "sendMsgConduit: \"" +++ L.take 129 (toLazyByteString m) +++ val "...\""

    kb = key2bytes k

-- | Closes all websockets connected to this instance of cannon.
--
-- This function is not tested anywhere as it is difficult to write an automated
-- test for. Some pointers on testing this function:
--
-- 1. Set values in cannon.integration.yaml for drainOpts such that it drains
-- "slowly", something like:
--
-- @
-- {gracePeriodSeconds: 1, millisecondsBetweenBatches: 500, minBatchSize: 5}
-- @
--
-- This will ensure that if there 10 or more websockets open, they get drained
-- in 2 batches of n/2.
--
-- 2. Write a test in brig or galley using 'bracketRN' function from
-- tasty-cannon. This function doesn't require users to exist. Just pass it n
-- UserIds and threadDelay for a long-ish time.
--
-- 3. During this threadDelay, send either SIGINT or SIGTERM to the cannon
-- process and use cannon logs to determine what is going on.
--
-- Example test, which worked at the time of writing this comment:
--
-- @
-- testCannonDrain :: Cannon -> Http ()
-- testCannonDrain cannon = do
--   users <- replicateM 50 randomId
--   WS.bracketRN cannon users $ \_websockets -> do
--     putStrLn "-------------------> Before delay"
--     threadDelay 100_000_000
--     putStrLn "-------------------> After delay"
--   putStrLn "-------------------> After bracket"
-- @
--
-- Use @pkill -INT -f cannon.integration.yaml@ to send SIGINT to the cannon
-- process.
drain :: WS ()
drain = do
  opts <- asks drainOpts
  websockets <- asks websockets
  numberOfConns <- fromIntegral <$> D.size websockets
  let maxNumberOfBatches = (opts ^. gracePeriodSeconds * 1000) `div` (opts ^. millisecondsBetweenBatches)
      computedBatchSize = numberOfConns `div` maxNumberOfBatches
      batchSize = max (opts ^. minBatchSize) computedBatchSize
  conns <- D.toList websockets
  info $
    msg (val "draining all websockets")
      . field "numberOfConns" numberOfConns
      . field "computedBatchSize" computedBatchSize
      . field "minBatchSize" (opts ^. minBatchSize)
      . field "batchSize" batchSize
      . field "maxNumberOfBatches" maxNumberOfBatches

  -- Sleeps for the grace period + 1 second. If the sleep completes, it means
  -- that draining didn't finish, and we should log that.
  timeoutAction <- async $ do
    -- Allocate 1 second more than the grace period to allow for overhead of
    -- spawning threads.
    liftIO $ threadDelay ((opts ^. gracePeriodSeconds) # Second + 1 # Second)
    err $ msg (val "Drain grace period expired") . field "gracePeriodSeconds" (opts ^. gracePeriodSeconds)

  for_ (chunksOf (fromIntegral batchSize) conns) $ \batch -> do
    -- 16 was chosen with a roll of a fair dice.
    void . async $ pooledMapConcurrentlyN_ 16 (uncurry close) batch
    liftIO $ threadDelay ((opts ^. millisecondsBetweenBatches) # MilliSecond)
  cancel timeoutAction
  info $ msg (val "Draining complete")

close :: Key -> Websocket -> WS ()
close k c = do
  let kb = key2bytes k
  debug $ client kb . msg (val "close websocket")
  liftIO $ sendClose (connection c) ("close" :: ByteString)

regInfo :: Key -> Maybe ClientId -> WS (Request -> Request)
regInfo k c = do
  e <- WS ask
  let h = externalHostname e
      p = portnum e
      r = "http://" <> h <> ":" <> pack (show p) <> "/i/push/"
  pure . Bilge.json . object $
    [ "user_id" .= decodeUtf8 (keyUserBytes k),
      "device_id" .= decodeUtf8 (keyConnBytes k),
      "resource" .= decodeUtf8 (r <> keyUserBytes k <> "/" <> keyConnBytes k),
      "client_id" .= c
    ]

client :: ByteString -> Msg -> Msg
client = field "client"

logKey :: Key -> Msg -> Msg
logKey k =
  field "user" (keyUserBytes k) . field "conn" (keyConnBytes k)

retry3x :: RetryPolicy
retry3x = limitRetries 3 <> exponentialBackoff 100000
