{-# OPTIONS_GHC -Wno-ambiguous-fields #-}

-- 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 Test.Cells where

import API.Galley
import qualified API.GalleyInternal as I
import Control.Concurrent
import Control.Concurrent.Async
import Control.Concurrent.STM
import qualified Control.Exception as E
import Control.Monad.Codensity
import Control.Monad.Reader
import Control.Retry
import qualified Data.Aeson as A
import Data.IORef
import qualified Data.Map as Map
import Network.AMQP
import Network.AMQP.Extended
import Notifications
import SetupHelpers
import System.Timeout
import Testlib.Prelude

testCellsEvent :: (HasCallStack) => App ()
testCellsEvent = do
  (alice, tid, [bob, chaz, dean, eve]) <- createTeam OwnDomain 5
  conv <- postConversation alice defProteus {team = Just tid} >>= getJSON 201
  q <- watchCellsEventsForTeam tid (convEvents conv)

  bobId <- bob %. "qualified_id"
  chazId <- chaz %. "qualified_id"
  deanId <- dean %. "qualified_id"
  eveId <- eve %. "qualified_id"

  addMembers alice conv def {role = Just "wire_member", users = [bobId]} >>= assertSuccess

  I.setCellsState alice conv "pending" >>= assertSuccess
  addMembers alice conv def {role = Just "wire_member", users = [chazId]} >>= assertSuccess

  do
    event <- getMessage q %. "payload.0"
    event %. "type" `shouldMatch` "conversation.member-join"
    event %. "conversation" `shouldMatch` (conv %. "qualified_id" & objId)
    event %. "qualified_from" `shouldMatch` (alice %. "qualified_id")
    users <- event %. "data.users" & asList
    assertOne users %. "qualified_id" `shouldMatch` chazId

  I.setCellsState alice conv "ready" >>= assertSuccess
  addMembers alice conv def {role = Just "wire_member", users = [deanId]} >>= assertSuccess

  do
    event <- getMessage q %. "payload.0"
    event %. "type" `shouldMatch` "conversation.member-join"
    event %. "conversation" `shouldMatch` (conv %. "qualified_id" & objId)
    event %. "qualified_from" `shouldMatch` (alice %. "qualified_id")
    users <- event %. "data.users" & asList
    assertOne users %. "qualified_id" `shouldMatch` deanId

  I.setCellsState alice conv "disabled" >>= assertSuccess
  addMembers alice conv def {role = Just "wire_member", users = [eveId]} >>= assertSuccess

  assertNoMessage q

testCellsCreationEvent :: (HasCallStack) => App ()
testCellsCreationEvent = do
  (alice, tid, _) <- createTeam OwnDomain 1
  q0 <- watchCellsEventsForTeam tid def
  conv <- postConversation alice defProteus {team = Just tid, cells = True} >>= getJSON 201

  let q = q0 {filter = isNotifConv conv} :: QueueConsumer

  event <- getMessage q %. "payload.0"
  event %. "type" `shouldMatch` "conversation.create"
  event %. "qualified_conversation.id" `shouldMatch` (conv %. "qualified_id.id")
  event %. "qualified_from" `shouldMatch` (alice %. "qualified_id")

  assertNoMessage q

testCellsDeletionEvent :: (HasCallStack) => App ()
testCellsDeletionEvent = do
  (alice, tid, _) <- createTeam OwnDomain 1
  q0 <- watchCellsEventsForTeam tid def
  conv <- postConversation alice defProteus {team = Just tid, cells = True} >>= getJSON 201
  void $ deleteTeamConversation tid conv alice >>= assertSuccess

  let q = q0 {filter = isConvDeleteNotif} :: QueueConsumer

  event <- getMessage q %. "payload.0"
  event %. "type" `shouldMatch` "conversation.delete"
  event %. "conversation" `shouldMatch` (conv %. "qualified_id.id")
  event %. "qualified_conversation" `shouldMatch` (conv %. "qualified_id")
  event %. "qualified_from" `shouldMatch` (alice %. "qualified_id")
  event %. "from" `shouldMatch` (alice %. "qualified_id.id")
  event %. "team" `shouldMatch` tid

  assertNoMessage q

testCellsCreationEventIsSentOnlyOnce :: (HasCallStack) => App ()
testCellsCreationEventIsSentOnlyOnce = do
  (alice, tid, members) <- createTeam OwnDomain 2
  q0 <- watchCellsEventsForTeam tid def
  conv <- postConversation alice defProteus {team = Just tid, cells = True, qualifiedUsers = members} >>= getJSON 201

  let q = q0 {filter = isNotifConv conv} :: QueueConsumer

  event <- getMessage q %. "payload.0"
  event %. "type" `shouldMatch` "conversation.create"
  event %. "qualified_conversation.id" `shouldMatch` (conv %. "qualified_id.id")
  event %. "qualified_from" `shouldMatch` (alice %. "qualified_id")

  assertNoMessage q

testCellsFeatureCheck :: (HasCallStack) => App ()
testCellsFeatureCheck = do
  (alice, tid, _) <- createTeam OwnDomain 1
  I.patchTeamFeature OwnDomain tid "cells" (object ["status" .= "disabled"]) >>= assertSuccess
  conv <- postConversation alice defProteus {team = Just tid} >>= getJSON 201
  bindResponse (I.setCellsState alice conv "ready") $ \resp -> do
    resp.status `shouldMatchInt` 403
    resp.json %. "label" `shouldMatch` "invalid-op"

testCellsEventOnFeatureToggle :: (HasCallStack) => App ()
testCellsEventOnFeatureToggle = do
  (_, tid, _) <- createTeam OwnDomain 1
  q <- watchCellsEventsForTeam tid def
  I.patchTeamFeature OwnDomain tid "cells" (object ["status" .= "disabled"]) >>= assertSuccess
  getMessage q >>= \event -> do
    event %. "payload.0.type" `shouldMatch` "feature-config.update"
    event %. "payload.0.name" `shouldMatch` "cells"
    event %. "payload.0.team" `shouldMatch` (asString tid)
    event %. "payload.0.data.status" `shouldMatch` "disabled"
  I.patchTeamFeature OwnDomain tid "cells" (object ["status" .= "enabled"]) >>= assertSuccess
  getMessage q >>= \event -> do
    event %. "payload.0.type" `shouldMatch` "feature-config.update"
    event %. "payload.0.name" `shouldMatch` "cells"
    event %. "payload.0.team" `shouldMatch` (asString tid)
    event %. "payload.0.data.status" `shouldMatch` "enabled"

testCellsTeamConversationCheck :: (HasCallStack) => App ()
testCellsTeamConversationCheck = do
  alice <- randomUser OwnDomain def
  conv <- postConversation alice defProteus >>= getJSON 201
  bindResponse (I.setCellsState alice conv "ready") $ \resp -> do
    resp.status `shouldMatchInt` 403
    resp.json %. "label" `shouldMatch` "invalid-op"

testCellsIgnoredEvents :: (HasCallStack) => App ()
testCellsIgnoredEvents = do
  (alice, tid, _) <- createTeam OwnDomain 1
  conv <- postConversation alice defProteus {team = Just tid} >>= getJSON 201
  I.setCellsState alice conv "ready" >>= assertSuccess
  q <- watchCellsEventsForTeam tid (convEvents conv)
  void $ updateMessageTimer alice conv 1000 >>= getBody 200
  assertNoMessage q

--------------------------------------------------------------------------------
-- Utilities

connectToCellsQueue :: ServiceMap -> TChan Message -> Codensity App ()
connectToCellsQueue sm messages = do
  queueName <- lift $ asks (.cellsEventQueue)

  env <- lift ask
  let opts = (demoteOpts env.rabbitMQConfig :: AmqpEndpoint) {vHost = sm.rabbitMqVHost}
  let createConnection :: IO Connection
      createConnection =
        recovering
          (limitRetries 5)
          skipAsyncExceptions
          ( const $ do
              connOpts <-
                mkConnectionOpts opts Nothing
              liftIO $ openConnection'' connOpts
          )
  conn <-
    hoistCodensity
      $ Codensity
      $ E.bracket createConnection closeConnection

  chan <-
    hoistCodensity
      $ Codensity
      $ E.bracket (openChannel conn) closeChannel

  handler <- lift $ appToIOKleisli $ \(m, e) -> do
    liftIO $ atomically $ writeTChan messages m
    liftIO $ ackMsg chan e.envDeliveryTag False

  void
    . hoistCodensity
    $ Codensity
    $ E.bracket
      (consumeMsgs chan (fromString queueName) Ack handler)
      (cancelConsumer chan)

getNextMessage :: QueueConsumer -> App Value
getNextMessage q = do
  m <- liftIO $ atomically $ readTChan q.chan
  v <- either assertFailure pure $ A.eitherDecode m.msgBody
  ok <- q.filter v
  if ok
    then pure v
    else getNextMessage q

getMessageMaybe :: QueueConsumer -> App (Maybe Value)
getMessageMaybe q = do
  timeOutSeconds <- asks (.timeOutSeconds)
  next <- appToIO (getNextMessage q)
  liftIO $ timeout (timeOutSeconds * 1000000) next

getMessage :: QueueConsumer -> App Value
getMessage q = getMessageMaybe q >>= assertJust "Cells queue timeout"

assertNoMessage :: QueueConsumer -> App ()
assertNoMessage f =
  getMessageMaybe f >>= \case
    Nothing -> pure ()
    Just m -> do
      j <- prettyJSON m
      assertFailure $ "Expected no message, got:\n" <> j

--------------------------------------------------------------------------------
-- Queue watcher

data QueueConsumer = QueueConsumer
  { chan :: TChan Message,
    filter :: Value -> App Bool
  }

startQueueWatcher :: ServiceMap -> App QueueWatcher
startQueueWatcher sm = do
  broadcast <- liftIO newBroadcastTChanIO
  readyVar <- liftIO $ newEmptyMVar
  doneVar <- liftIO $ newEmptyMVar

  startIO <- appToIO $ lowerCodensity $ do
    void $ connectToCellsQueue sm broadcast
    liftIO $ putMVar readyVar ()
    liftIO $ takeMVar doneVar

  void $ liftIO $ async startIO
  liftIO $ takeMVar readyVar

  pure QueueWatcher {doneVar, broadcast}

ensureWatcher :: String -> App QueueWatcher
ensureWatcher domain = do
  watchersLock <- asks (.cellsEventWatchersLock)
  watchersRef <- asks (.cellsEventWatchers)
  serviceMaps <- asks (.serviceMap)
  sm <- assertOne $ Map.lookup domain serviceMaps

  start <- appToIO (startQueueWatcher sm)

  liftIO
    $ E.bracket
      (putMVar watchersLock ())
      (\_ -> tryTakeMVar watchersLock)
    $ \_ -> do
      watchers <- liftIO $ readIORef watchersRef
      case Map.lookup domain watchers of
        Nothing -> do
          watcher <- start
          let watchers' = Map.insert domain watcher watchers
          writeIORef watchersRef watchers'
          pure watcher
        Just watcher -> pure watcher

data WatchCellsEvents = WatchCellsEvents
  { domain :: Either Domain String,
    filter :: Value -> App Bool
  }

instance Default WatchCellsEvents where
  def =
    WatchCellsEvents
      { domain = Left OwnDomain,
        filter = const (pure True)
      }

convEvents :: (MakesValue conv) => conv -> WatchCellsEvents
convEvents conv = def {filter = isNotifConv conv}

watchCellsEvents :: WatchCellsEvents -> App QueueConsumer
watchCellsEvents opts = do
  domain <- either (asString . make) pure opts.domain
  watcher <- ensureWatcher domain
  chan <- liftIO $ atomically $ dupTChan watcher.broadcast
  pure QueueConsumer {filter = opts.filter, chan}

watchCellsEventsForTeam :: String -> WatchCellsEvents -> App QueueConsumer
watchCellsEventsForTeam tid opts = do
  q <- watchCellsEvents opts
  let isEventForTeam v = fieldEquals @Value v "payload.0.team" tid
  -- the cells event queue is shared by tests
  -- let's hope this filter reduces the risk of tests interfering with each other
  pure $ q {filter = isEventForTeam}
