{-# LANGUAGE GeneralizedNewtypeDeriving #-}
-- Disabling to stop warnings on HasCallStack
{-# OPTIONS_GHC -Wno-redundant-constraints #-}

-- 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 API.MLS.Util where

import API.Util
import Bilge
import Bilge.Assert
import Control.Arrow ((&&&))
import Control.Error.Util
import Control.Lens (preview, to, view, (.~), (^..), (^?))
import Control.Monad.Catch
import Control.Monad.Cont
import Control.Monad.State (StateT, evalStateT)
import Control.Monad.State qualified as State
import Control.Monad.Trans.Maybe
import Data.Aeson.Lens
import Data.Bifunctor
import Data.Binary.Builder (toLazyByteString)
import Data.Binary.Get
import Data.ByteArray qualified as BA
import Data.ByteString qualified as BS
import Data.ByteString.Base64.URL qualified as B64U
import Data.ByteString.Conversion
import Data.ByteString.Lazy qualified as LBS
import Data.Default
import Data.Domain
import Data.Id
import Data.Json.Util hiding ((#))
import Data.Map qualified as Map
import Data.Monoid
import Data.Qualified
import Data.Set qualified as Set
import Data.Text qualified as T
import Data.Text.Encoding qualified as T
import Data.Time
import Data.UUID qualified as UUID
import Data.UUID.V4 qualified as UUIDV4
import Imports hiding (getFirst, getSymbolicLinkTarget)
import Network.HTTP.Client (setQueryString)
import System.FilePath
import System.IO.Temp
import System.Posix hiding (createDirectory)
import System.Process
import Test.QuickCheck (arbitrary, generate)
import Test.Tasty.Cannon qualified as WS
import Test.Tasty.HUnit
import TestHelpers
import TestSetup
import Web.HttpApiData
import Wire.API.Conversation
import Wire.API.Conversation.Action
import Wire.API.Conversation.Protocol
import Wire.API.Conversation.Role (roleNameWireMember)
import Wire.API.Event.Conversation
import Wire.API.Federation.API.Galley
import Wire.API.MLS.CipherSuite (SignatureSchemeTag (Ed25519))
import Wire.API.MLS.CommitBundle
import Wire.API.MLS.Credential
import Wire.API.MLS.Group.Serialisation qualified as Group
import Wire.API.MLS.KeyPackage
import Wire.API.MLS.Keys
import Wire.API.MLS.LeafNode
import Wire.API.MLS.Message
import Wire.API.MLS.Serialisation
import Wire.API.MLS.SubConversation
import Wire.API.Routes.Public.Galley.MLS
import Wire.API.User.Client
import Wire.API.User.Client.Prekey
import Wire.Options.Galley
import Wire.Options.Galley qualified as Opts
import Wire.Options.Keys

cid2Str :: ClientIdentity -> String
cid2Str cid =
  show cid.ciUser
    <> ":"
    <> T.unpack (clientToText cid.ciClient)
    <> "@"
    <> T.unpack (domainText (ciDomain cid))

postMessage ::
  ( HasCallStack,
    MonadIO m,
    MonadHttp m,
    HasGalley m
  ) =>
  ClientIdentity ->
  ByteString ->
  m ResponseLBS
postMessage sender msg = do
  galleyCall <- viewGalley
  post
    ( galleyCall
        . paths ["mls", "messages"]
        . zUser (ciUser sender)
        . zClient (ciClient sender)
        . zConn "conn"
        . Bilge.content "message/mls"
        . bytes msg
    )

localPostCommitBundle ::
  ( HasCallStack,
    MonadIO m,
    MonadHttp m,
    HasGalley m
  ) =>
  ClientIdentity ->
  ByteString ->
  m ResponseLBS
localPostCommitBundle sender bundle = do
  galleyCall <- viewGalley
  post
    ( galleyCall
        . paths ["mls", "commit-bundles"]
        . zUser (ciUser sender)
        . zClient (ciClient sender)
        . zConn "conn"
        . Bilge.content "message/mls"
        . bytes bundle
    )

remotePostCommitBundle ::
  ( MonadIO m,
    MonadReader TestSetup m
  ) =>
  Remote ClientIdentity ->
  Qualified ConvOrSubConvId ->
  ByteString ->
  m [Event]
remotePostCommitBundle rsender qcs bundle = do
  client <- view tsFedGalleyClient
  let msr =
        MLSMessageSendRequest
          { convOrSubId = qUnqualified qcs,
            sender = ciUser (tUnqualified rsender),
            senderClient = ciClient (tUnqualified rsender),
            rawMessage = Base64ByteString bundle,
            enableOutOfSyncCheck = Nothing
          }
  runFedClient
    @"send-mls-commit-bundle"
    client
    (tDomain rsender)
    msr
    >>= liftIO . \case
      MLSMessageResponseError e ->
        assertFailure $
          "error while receiving commit bundle: " <> show e
      MLSMessageResponseProtocolError e ->
        assertFailure $
          "protocol error while receiving commit bundle: " <> T.unpack e
      MLSMessageResponseProposalFailure e ->
        assertFailure $
          "proposal failure while receiving commit bundle: " <> displayException e
      e@(MLSMessageResponseUnreachableBackends _) ->
        assertFailure $
          "error while receiving commit bundle: " <> show e
      e@(MLSMessageResponseGroupInfoDiagnostics _) ->
        assertFailure $
          "error while receiving commit bundle: " <> show e
      e@(MLSMessageResponseNonFederatingBackends _) ->
        assertFailure $
          "error while receiving commit bundle: " <> show e
      e@(MLSMessageResponseOutOfSyncError _) ->
        assertFailure $
          "error while receiving commit bundle: " <> show e
      MLSMessageResponseUpdates _ -> pure []

postCommitBundle ::
  (HasCallStack) =>
  ClientIdentity ->
  Qualified ConvOrSubConvId ->
  ByteString ->
  TestM [Event]
postCommitBundle sender qcs bundle = do
  loc <- qualifyLocal ()
  foldQualified
    loc
    ( \_ ->
        fmap mmssEvents . responseJsonError
          =<< localPostCommitBundle sender bundle
            <!! const 201 === statusCode
    )
    (\rsender -> remotePostCommitBundle rsender qcs bundle)
    (cidQualifiedUser sender $> sender)

saveRemovalKey :: FilePath -> TestM ()
saveRemovalKey fp = do
  keys <- fromJust <$> view (tsGConf . settings . mlsPrivateKeyPaths)
  keysByPurpose <- liftIO $ loadAllMLSKeys keys
  let pub = (mlsKeysToPublic keysByPurpose.removal).ed25519
  liftIO $ BS.writeFile fp (BA.convert $ unwrapMLSPublicKey pub)

data MLSState = MLSState
  { mlsBaseDir :: FilePath,
    -- | for creating clients
    mlsUnusedPrekeys :: [LastPrekey],
    mlsMembers :: Set ClientIdentity,
    -- | users expected to receive a welcome message after the next commit
    mlsNewMembers :: Set ClientIdentity,
    mlsClientGroupState :: Map ClientIdentity ByteString,
    mlsGroupId :: Maybe GroupId,
    mlsConvId :: Maybe (Qualified ConvOrSubConvId),
    mlsEpoch :: Word64
  }
  deriving (Show)

newtype MLSTest a = MLSTest {unMLSTest :: StateT MLSState TestM a}
  deriving newtype
    ( Functor,
      Applicative,
      Monad,
      MonadThrow,
      MonadHttp,
      MonadIO,
      MonadCatch,
      MonadFail,
      MonadMask,
      State.MonadState MLSState,
      MonadReader TestSetup
    )

instance HasGalley MLSTest where
  viewGalley = MLSTest $ lift viewGalley
  viewGalleyOpts = MLSTest $ lift viewGalleyOpts

instance HasBrig MLSTest where
  viewBrig = MLSTest $ lift viewBrig

instance HasSettingsOverrides MLSTest where
  withSettingsOverrides f (MLSTest action) = MLSTest $
    State.StateT $ \s ->
      withSettingsOverrides f (State.runStateT action s)

liftTest :: TestM a -> MLSTest a
liftTest = MLSTest . lift

runMLSTest :: (HasCallStack) => MLSTest a -> TestM a
runMLSTest (MLSTest m) =
  withSystemTempDirectory "mls" $ \tmp -> do
    saveRemovalKey (tmp </> "removal.key")
    evalStateT
      m
      MLSState
        { mlsBaseDir = tmp,
          mlsUnusedPrekeys = someLastPrekeys,
          mlsMembers = mempty,
          mlsNewMembers = mempty,
          mlsClientGroupState = mempty,
          mlsGroupId = Nothing,
          mlsConvId = Nothing,
          mlsEpoch = 0
        }

data MessagePackage = MessagePackage
  { mpSender :: ClientIdentity,
    mpMessage :: ByteString,
    mpWelcome :: Maybe ByteString,
    mpGroupInfo :: Maybe ByteString
  }
  deriving (Show)

takeLastPrekeyNG :: (HasCallStack) => MLSTest LastPrekey
takeLastPrekeyNG = do
  s <- State.get
  case mlsUnusedPrekeys s of
    (pk : pks) -> do
      State.modify (\s' -> s' {mlsUnusedPrekeys = pks})
      pure pk
    [] -> error "no prekeys left"

toRandomFile :: ByteString -> MLSTest FilePath
toRandomFile bs = do
  p <- randomFileName
  liftIO $ BS.writeFile p bs
  pure p

randomFileName :: MLSTest FilePath
randomFileName = do
  bd <- State.gets mlsBaseDir
  (bd </>) . UUID.toString <$> liftIO UUIDV4.nextRandom

mlscli :: (HasCallStack) => ClientIdentity -> [String] -> Maybe ByteString -> MLSTest ByteString
mlscli qcid args mbstdin = do
  bd <- State.gets mlsBaseDir
  let cdir = bd </> cid2Str qcid

  groupOut <- randomFileName
  let substOut = argSubst "<group-out>" groupOut

  hasState <- hasClientGroupState qcid
  substIn <-
    if hasState
      then do
        gs <- getClientGroupState qcid
        fn <- toRandomFile gs
        pure (argSubst "<group-in>" fn)
      else pure Imports.id

  out <-
    liftIO $
      spawn
        ( proc
            "mls-test-cli"
            ( ["--store", cdir </> "store"]
                <> map (substIn . substOut) args
            )
        )
        mbstdin

  groupOutWritten <- liftIO $ doesFileExist groupOut
  when groupOutWritten $ do
    gs <- liftIO (BS.readFile groupOut)
    setClientGroupState qcid gs
  pure out

argSubst :: String -> String -> String -> String
argSubst from to_ s =
  if s == from then to_ else s

createWireClient :: (HasCallStack) => Qualified UserId -> MLSTest ClientIdentity
createWireClient qusr = do
  lpk <- takeLastPrekeyNG
  clientId <- liftTest $ randomClient (qUnqualified qusr) lpk
  pure $ mkClientIdentity qusr clientId

initMLSClient :: (HasCallStack) => ClientIdentity -> MLSTest ()
initMLSClient cid = do
  bd <- State.gets mlsBaseDir
  createDirectory $ bd </> cid2Str cid
  void $ mlscli cid ["init", cid2Str cid] Nothing

createLocalMLSClient :: Local UserId -> MLSTest ClientIdentity
createLocalMLSClient (tUntagged -> qusr) = do
  qcid <- createWireClient qusr
  initMLSClient qcid

  -- set public key
  pkey <- mlscli qcid ["public-key"] Nothing
  brigCall <- viewBrig
  let update = defUpdateClient {updateClientMLSPublicKeys = Map.singleton Ed25519 pkey}
  put
    ( brigCall
        . paths ["clients", toByteString' . ciClient $ qcid]
        . zUser (ciUser qcid)
        . json update
    )
    !!! const 200 === statusCode
  pure qcid

-- | Create new mls client and register with backend. If the user is remote,
-- this only creates a fake client (see 'createFakeMLSClient').
createMLSClient :: (HasCallStack) => Qualified UserId -> MLSTest ClientIdentity
createMLSClient qusr = do
  loc <- liftTest $ qualifyLocal ()
  foldQualified loc createLocalMLSClient (createFakeMLSClient . tUntagged) qusr

-- | Like 'createMLSClient', but do not actually register client with backend.
createFakeMLSClient :: (HasCallStack) => Qualified UserId -> MLSTest ClientIdentity
createFakeMLSClient qusr = do
  c <- liftIO $ generate arbitrary
  let cid = mkClientIdentity qusr c
  initMLSClient cid
  pure cid

-- | create and upload to backend
uploadNewKeyPackage :: (HasCallStack) => ClientIdentity -> MLSTest (RawMLS KeyPackage)
uploadNewKeyPackage qcid = do
  (kp, _) <- generateKeyPackage qcid

  -- upload key package
  brigCall <- viewBrig
  post
    ( brigCall
        . paths ["mls", "key-packages", "self", toByteString' . ciClient $ qcid]
        . zUser (ciUser qcid)
        . json (KeyPackageUpload [kp])
    )
    !!! const 201 === statusCode
  pure kp

generateKeyPackage :: (HasCallStack) => ClientIdentity -> MLSTest (RawMLS KeyPackage, KeyPackageRef)
generateKeyPackage qcid = do
  kpData <- mlscli qcid ["key-package", "create"] Nothing
  kp <- liftIO $ decodeMLSError kpData
  let ref = fromJust (kpRef' kp)
  pure (kp, ref)

setClientGroupState :: (HasCallStack) => ClientIdentity -> ByteString -> MLSTest ()
setClientGroupState cid g =
  State.modify $ \s ->
    s {mlsClientGroupState = Map.insert cid g (mlsClientGroupState s)}

getClientGroupState :: (HasCallStack) => ClientIdentity -> MLSTest ByteString
getClientGroupState cid = do
  mgs <- State.gets (Map.lookup cid . mlsClientGroupState)
  case mgs of
    Nothing -> liftIO $ assertFailure ("Attempted to get non-existing group state for client " <> show cid)
    Just g -> pure g

hasClientGroupState :: (HasCallStack) => ClientIdentity -> MLSTest Bool
hasClientGroupState cid =
  State.gets (isJust . Map.lookup cid . mlsClientGroupState)

-- | Create a conversation from a provided action and then create a
-- corresponding group.
setupMLSGroupWithConv ::
  (HasCallStack) =>
  MLSTest (Conversation GroupConvType) ->
  ClientIdentity ->
  MLSTest (GroupId, Qualified ConvId)
setupMLSGroupWithConv convAction creator = do
  ownDomain <- liftTest viewFederationDomain
  liftIO $ assertEqual "creator is not local" (ciDomain creator) ownDomain
  conv <- convAction
  let groupId =
        fromJust
          ( asum
              [ preview (to (.protocol) . _ProtocolMLS . to cnvmlsGroupId) conv,
                preview (to (.protocol) . _ProtocolMixed . to cnvmlsGroupId) conv
              ]
          )
  let qcnv = conv.qualifiedId
  createGroup creator (fmap Conv qcnv) groupId
  pure (groupId, qcnv)

-- | Create conversation and corresponding group.
setupMLSGroup :: (HasCallStack) => ClientIdentity -> MLSTest (GroupId, Qualified ConvId)
setupMLSGroup creator = setupMLSGroupWithConv action creator
  where
    action =
      responseJsonError
        =<< liftTest
          ( postConvQualified
              (ciUser creator)
              (Just (ciClient creator))
              defNewMLSConv
          )
          <!! const 201 === statusCode

memberToOtherMember :: Member -> OtherMember
memberToOtherMember m =
  OtherMember
    { omService = m.memService,
      omQualifiedId = m.memId,
      omConvRoleName = m.memConvRoleName
    }

-- | Create self-conversation and corresponding group.
setupMLSSelfGroup :: (HasCallStack) => ClientIdentity -> MLSTest (GroupId, Qualified ConvId)
setupMLSSelfGroup creator = setupMLSGroupWithConv action creator
  where
    action =
      fmap fromOwnConversation
        . responseJsonError
        =<< liftTest
          (getSelfConv (ciUser creator))
          <!! const 200 === statusCode

createGroup :: ClientIdentity -> Qualified ConvOrSubConvId -> GroupId -> MLSTest ()
createGroup cid qcs gid = do
  State.gets mlsGroupId >>= \case
    Just _ -> liftIO $ assertFailure "only one group can be created"
    Nothing -> pure ()
  resetGroup cid qcs gid

resetGroup :: ClientIdentity -> Qualified ConvOrSubConvId -> GroupId -> MLSTest ()
resetGroup cid qcs gid = do
  State.modify $ \s ->
    s
      { mlsGroupId = Just gid,
        mlsConvId = Just qcs,
        mlsMembers = Set.singleton cid,
        mlsEpoch = 0,
        mlsNewMembers = mempty
      }
  resetClientGroup cid gid

resetClientGroup :: ClientIdentity -> GroupId -> MLSTest ()
resetClientGroup cid gid = do
  bd <- State.gets mlsBaseDir
  groupJSON <-
    mlscli
      cid
      [ "group",
        "create",
        "--removal-key",
        bd </> "removal.key",
        T.unpack (toBase64Text (unGroupId gid))
      ]
      Nothing
  setClientGroupState cid groupJSON

getConvId :: MLSTest (Qualified ConvOrSubConvId)
getConvId =
  State.gets mlsConvId
    >>= maybe (liftIO (assertFailure "Uninitialised test conversation")) pure

createSubConv ::
  (HasCallStack) =>
  Qualified ConvId ->
  ClientIdentity ->
  SubConvId ->
  MLSTest (Qualified ConvOrSubConvId)
createSubConv qcnv creator subId = do
  sub <-
    liftTest $
      responseJsonError
        =<< getSubConv (ciUser creator) qcnv subId
          <!! const 200 === statusCode
  let qcs = fmap (flip SubConv subId) qcnv
  resetGroup creator qcs (pscGroupId sub)
  void $ createPendingProposalCommit creator >>= sendAndConsumeCommitBundle
  pure qcs

-- | Create a local group only without a conversation. This simulates creating
-- an MLS conversation on a remote backend.
setupFakeMLSGroup ::
  (HasCallStack) =>
  ClientIdentity ->
  Maybe SubConvId ->
  MLSTest (GroupId, Qualified ConvId)
setupFakeMLSGroup creator mSubId = do
  qcnv <- randomQualifiedId (ciDomain creator)
  let groupId =
        Group.newGroupId RegularConv $
          maybe (Conv <$> qcnv) ((<$> qcnv) . flip SubConv) mSubId
  createGroup creator (fmap Conv qcnv) groupId
  pure (groupId, qcnv)

claimLocalKeyPackages :: (HasCallStack) => ClientIdentity -> Local UserId -> MLSTest KeyPackageBundle
claimLocalKeyPackages qcid lusr = do
  brigCall <- viewBrig
  responseJsonError
    =<< post
      ( brigCall
          . paths ["mls", "key-packages", "claim", toByteString' (tDomain lusr), toByteString' (tUnqualified lusr)]
          . setQueryString [("ciphersuite", Just "0x0001")]
          . zUser (ciUser qcid)
      )
      <!! const 200 === statusCode

-- | Get all test clients of a user by listing the temporary MLS directory.
getUserClients :: (HasCallStack) => Qualified UserId -> MLSTest [ClientIdentity]
getUserClients qusr = do
  bd <- State.gets mlsBaseDir
  files <- getDirectoryContents bd
  let toClient f = do
        cid <- hush . decodeMLS' . T.encodeUtf8 . T.pack $ f
        guard (cidQualifiedUser cid == qusr)
        pure cid
  pure . mapMaybe toClient $ files

-- | Generate one key package for each client of a remote user
claimRemoteKeyPackages :: (HasCallStack) => Remote UserId -> MLSTest KeyPackageBundle
claimRemoteKeyPackages (tUntagged -> qusr) = do
  clients <- getUserClients qusr
  fmap (KeyPackageBundle . Set.fromList) $
    for clients $ \cid -> do
      (kp, ref) <- generateKeyPackage cid
      pure $
        KeyPackageBundleEntry
          { user = qusr,
            client = ciClient cid,
            ref = ref,
            keyPackage = KeyPackageData (raw kp)
          }

-- | Claim key package for a local user, or generate and map key packages for remote ones.
claimKeyPackages ::
  (HasCallStack) =>
  ClientIdentity ->
  Qualified UserId ->
  MLSTest KeyPackageBundle
claimKeyPackages cid qusr = do
  loc <- liftTest $ qualifyLocal ()
  foldQualified loc (claimLocalKeyPackages cid) claimRemoteKeyPackages qusr

bundleKeyPackages :: KeyPackageBundle -> [(ClientIdentity, ByteString)]
bundleKeyPackages bundle =
  let getEntry be =
        ( mkClientIdentity be.user be.client,
          kpData be.keyPackage
        )
   in map getEntry (toList bundle.entries)

-- | Claim keypackages and create a commit/welcome pair on a given client.
-- Note that this alters the state of the group immediately. If we want to test
-- a scenario where the commit is rejected by the backend, we can restore the
-- group to the previous state by using an older version of the group file.
createAddCommit :: (HasCallStack) => ClientIdentity -> [Qualified UserId] -> MLSTest MessagePackage
createAddCommit cid users = do
  kps <- fmap (concatMap bundleKeyPackages) . traverse (claimKeyPackages cid) $ users
  liftIO $ assertBool "no key packages could be claimed" (not (null kps))
  createAddCommitWithKeyPackages cid kps

createExternalCommit ::
  (HasCallStack) =>
  ClientIdentity ->
  Maybe ByteString ->
  Qualified ConvOrSubConvId ->
  MLSTest MessagePackage
createExternalCommit qcid mpgs qcs = do
  bd <- State.gets mlsBaseDir
  pgsFile <- liftIO $ emptyTempFile bd "pgs"
  pgs <- case mpgs of
    Nothing -> liftTest $ getGroupInfo (cidQualifiedUser qcid) qcs
    Just v -> pure v
  commit <-
    mlscli
      qcid
      [ "external-commit",
        "--group-info-in",
        "-",
        "--group-info-out",
        pgsFile,
        "--group-out",
        "<group-out>"
      ]
      (Just pgs)

  State.modify $ \mls ->
    mls
      { mlsNewMembers = Set.singleton qcid
      -- This might be a different client than those that have been in the
      -- group from before.
      }

  newPgs <- liftIO $ BS.readFile pgsFile
  pure $
    MessagePackage
      { mpSender = qcid,
        mpMessage = commit,
        mpWelcome = Nothing,
        mpGroupInfo = Just newPgs
      }

createAddProposals :: (HasCallStack) => ClientIdentity -> [Qualified UserId] -> MLSTest [MessagePackage]
createAddProposals cid users = do
  kps <- fmap (concatMap bundleKeyPackages) . traverse (claimKeyPackages cid) $ users
  traverse (createAddProposalWithKeyPackage cid) kps

-- | Create an application message.
createApplicationMessage ::
  (HasCallStack) =>
  ClientIdentity ->
  String ->
  MLSTest MessagePackage
createApplicationMessage cid messageContent = do
  message <-
    mlscli
      cid
      ["message", "--group-in", "<group-in>", messageContent, "--group-out", "<group-out>"]
      Nothing

  pure $
    MessagePackage
      { mpSender = cid,
        mpMessage = message,
        mpWelcome = Nothing,
        mpGroupInfo = Nothing
      }

createAddCommitWithKeyPackages ::
  (HasCallStack) =>
  ClientIdentity ->
  [(ClientIdentity, ByteString)] ->
  MLSTest MessagePackage
createAddCommitWithKeyPackages qcid clientsAndKeyPackages = do
  bd <- State.gets mlsBaseDir
  welcomeFile <- liftIO $ emptyTempFile bd "welcome"
  giFile <- liftIO $ emptyTempFile bd "gi"

  commit <- runContT (traverse (withTempKeyPackageFile . snd) clientsAndKeyPackages) $ \kpFiles ->
    mlscli
      qcid
      ( [ "member",
          "add",
          "--group",
          "<group-in>",
          "--welcome-out",
          welcomeFile,
          "--group-info-out",
          giFile,
          "--group-out",
          "<group-out>"
        ]
          <> kpFiles
      )
      Nothing

  State.modify $ \mls ->
    mls
      { mlsNewMembers = Set.fromList (map fst clientsAndKeyPackages)
      }

  welcome <- liftIO $ BS.readFile welcomeFile
  gi <- liftIO $ BS.readFile giFile
  pure $
    MessagePackage
      { mpSender = qcid,
        mpMessage = commit,
        mpWelcome = Just welcome,
        mpGroupInfo = Just gi
      }

createAddProposalWithKeyPackage ::
  ClientIdentity ->
  (ClientIdentity, ByteString) ->
  MLSTest MessagePackage
createAddProposalWithKeyPackage cid (_, kp) = do
  prop <- runContT (withTempKeyPackageFile kp) $ \kpFile ->
    mlscli
      cid
      ["proposal", "--group-in", "<group-in>", "--group-out", "<group-out>", "add", kpFile]
      Nothing
  pure
    MessagePackage
      { mpSender = cid,
        mpMessage = prop,
        mpWelcome = Nothing,
        mpGroupInfo = Nothing
      }

createPendingProposalCommit :: (HasCallStack) => ClientIdentity -> MLSTest MessagePackage
createPendingProposalCommit qcid = do
  bd <- State.gets mlsBaseDir
  welcomeFile <- liftIO $ emptyTempFile bd "welcome"
  pgsFile <- liftIO $ emptyTempFile bd "pgs"
  commit <-
    mlscli
      qcid
      [ "commit",
        "--group",
        "<group-in>",
        "--group-out",
        "<group-out>",
        "--welcome-out",
        welcomeFile,
        "--group-info-out",
        pgsFile
      ]
      Nothing

  welcome <- liftIO $ readWelcome welcomeFile
  pgs <- liftIO $ BS.readFile pgsFile
  pure
    MessagePackage
      { mpSender = qcid,
        mpMessage = commit,
        mpWelcome = welcome,
        mpGroupInfo = Just pgs
      }

readWelcome :: FilePath -> IO (Maybe ByteString)
readWelcome fp = runMaybeT $ do
  liftIO (doesFileExist fp) >>= guard
  stat <- liftIO $ getFileStatus fp
  guard $ fileSize stat > 0
  liftIO $ BS.readFile fp

createExternalAddProposal :: (HasCallStack) => ClientIdentity -> MLSTest MessagePackage
createExternalAddProposal joiner = do
  groupId <-
    State.gets mlsGroupId >>= \case
      Nothing -> liftIO $ assertFailure "Creating add proposal for non-existing group"
      Just g -> pure g
  epoch <- State.gets mlsEpoch
  proposal <-
    mlscli
      joiner
      [ "external-proposal",
        "--group-id",
        T.unpack (toBase64Text (unGroupId groupId)),
        "--epoch",
        show epoch,
        "add"
      ]
      Nothing

  State.modify $ \mls ->
    mls
      { mlsNewMembers = mlsNewMembers mls <> Set.singleton joiner
      }
  pure
    MessagePackage
      { mpSender = joiner,
        mpMessage = proposal,
        mpWelcome = Nothing,
        mpGroupInfo = Nothing
      }

consumeWelcome :: (HasCallStack) => ByteString -> MLSTest ()
consumeWelcome welcome = do
  qcids <- State.gets mlsNewMembers
  for_ qcids $ \qcid -> do
    hasState <- hasClientGroupState qcid
    liftIO $ assertBool "Existing clients in a conversation should not consume welcomes" (not hasState)
    void $
      mlscli
        qcid
        [ "group",
          "from-welcome",
          "--group-out",
          "<group-out>",
          "-"
        ]
        (Just welcome)

-- | Make all member clients consume a given message.
consumeMessage :: (HasCallStack) => MessagePackage -> MLSTest ()
consumeMessage msg = do
  mems <- State.gets mlsMembers
  for_ (Set.delete (mpSender msg) mems) $ \cid ->
    consumeMessage1 cid (mpMessage msg)

consumeMessage1 :: (HasCallStack) => ClientIdentity -> ByteString -> MLSTest ()
consumeMessage1 cid msg =
  void $
    mlscli
      cid
      [ "consume",
        "--group",
        "<group-in>",
        "--group-out",
        "<group-out>",
        "-"
      ]
      (Just msg)

-- | Send an MLS message and simulate clients receiving it. If the message is a
-- commit, the 'sendAndConsumeCommitBundle' function should be used instead.
sendAndConsumeMessage :: (HasCallStack) => MessagePackage -> MLSTest [Event]
sendAndConsumeMessage mp = do
  for_ mp.mpWelcome $ \_ -> liftIO $ assertFailure "use sendAndConsumeCommitBundle"
  res <-
    fmap mmssEvents $
      responseJsonError
        =<< postMessage (mpSender mp) (mpMessage mp)
          <!! const 201 === statusCode
  consumeMessage mp
  pure res

mkBundle :: MessagePackage -> Either Text CommitBundle
mkBundle mp = do
  commitB <- first ("Commit: " <>) $ decodeMLS' (mpMessage mp)
  welcomeB <- first ("Welcome: " <>) $ for (mpWelcome mp) $ \m -> do
    w <- decodeMLS' @Message m
    case w.content of
      MessageWelcome welcomeB -> pure welcomeB
      _ -> Left "expected welcome"
  ginfo <- note "group info unavailable" (mpGroupInfo mp)
  ginfoB <- first ("GroupInfo: " <>) $ decodeMLS' ginfo
  pure $ CommitBundle commitB welcomeB ginfoB Nothing

createBundle :: (HasCallStack, MonadIO m) => MessagePackage -> m ByteString
createBundle mp = do
  bundle <-
    either (liftIO . assertFailure . T.unpack) pure $
      mkBundle mp
  pure (encodeMLS' bundle)

sendAndConsumeCommitBundle :: (HasCallStack) => MessagePackage -> MLSTest [Event]
sendAndConsumeCommitBundle mp = do
  qcs <- getConvId
  bundle <- createBundle mp
  resp <- liftTest $ postCommitBundle (mpSender mp) qcs bundle
  consumeMessage mp
  traverse_ consumeWelcome (mpWelcome mp)

  -- increment epoch and add new clients
  State.modify $ \mls ->
    mls
      { mlsEpoch = mlsEpoch mls + 1,
        mlsMembers = mlsMembers mls <> mlsNewMembers mls,
        mlsNewMembers = mempty
      }

  pure resp

mlsBracket ::
  (HasCallStack) =>
  [ClientIdentity] ->
  ([WS.WebSocket] -> MLSTest a) ->
  MLSTest a
mlsBracket clients k = do
  c <- view tsCannon
  WS.bracketAsClientRN c (map (ciUser &&& ciClient) clients) k

readGroupState :: ByteString -> [(ClientIdentity, LeafIndex)]
readGroupState j = do
  (node, n) <- zip (j ^.. key "group" . key "public_group" . key "treesync" . key "tree" . key "leaf_nodes" . _Array . traverse) [0 ..]
  case node ^? key "node" of
    Just leafNode -> do
      identityBytes <- leafNode ^.. key "payload" . key "credential" . key "credential" . key "Basic" . key "identity" . key "vec"
      let identity = BS.pack (identityBytes ^.. _Array . traverse . _Integer . to fromIntegral)
      cid <- case decodeMLS' identity of
        Left _ -> []
        Right x -> pure x
      pure (cid, n)
    Nothing -> []

getClientsFromGroupState ::
  ClientIdentity ->
  Qualified UserId ->
  MLSTest [(ClientIdentity, LeafIndex)]
getClientsFromGroupState cid u = do
  groupState <- readGroupState <$> getClientGroupState cid
  pure $ filter (\(cid', _) -> cidQualifiedUser cid' == u) groupState

clientKeyPair :: ClientIdentity -> MLSTest (ByteString, ByteString)
clientKeyPair cid = do
  bd <- State.gets mlsBaseDir
  credential <-
    liftIO . BS.readFile $
      bd </> cid2Str cid </> "store" </> T.unpack (T.decodeUtf8 (B64U.encode "self"))
  case runGetOrFail
    ((,) <$> parseMLSBytes @VarInt <*> parseMLSBytes @VarInt)
    (LBS.fromStrict credential) of
    Left (_, _, msg) -> liftIO $ assertFailure msg
    Right (_, _, keys) -> pure keys

receiveOnConvUpdated ::
  (MonadReader TestSetup m, MonadIO m) =>
  Qualified ConvId ->
  Qualified UserId ->
  Qualified UserId ->
  m ()
receiveOnConvUpdated conv origUser joiner = do
  client <- view tsFedGalleyClient
  now <- liftIO getCurrentTime
  let cu =
        ConversationUpdate
          { time = now,
            origUserId = origUser,
            convId = qUnqualified conv,
            alreadyPresentUsers = [qUnqualified joiner],
            action =
              SomeConversationAction
                SConversationJoinTag
                ConversationJoin
                  { users = pure joiner,
                    role = roleNameWireMember,
                    joinType = InternalAdd
                  },
            extraConversationData = def
          }
  void $
    runFedClient
      @"on-conversation-updated"
      client
      (qDomain conv)
      cu

getGroupInfo :: (HasCallStack) => Qualified UserId -> Qualified ConvOrSubConvId -> TestM ByteString
getGroupInfo qusr qcs = do
  loc <- qualifyLocal ()
  foldQualified
    loc
    ( \lusr ->
        fmap (LBS.toStrict . fromJust . responseBody) $
          localGetGroupInfo
            (tUnqualified lusr)
            qcs
            <!! const 200 === statusCode
    )
    (\rusr -> remoteGetGroupInfo rusr qcs)
    qusr

localGetGroupInfo ::
  ( HasCallStack,
    MonadIO m,
    MonadHttp m,
    HasGalley m
  ) =>
  UserId ->
  Qualified ConvOrSubConvId ->
  m ResponseLBS
localGetGroupInfo sender qcs = do
  galleyCall <- viewGalley
  case qUnqualified qcs of
    Conv cnv ->
      get
        ( galleyCall
            . paths
              [ "conversations",
                toByteString' (qDomain qcs),
                toByteString' cnv,
                "groupinfo"
              ]
            . zUser sender
            . zConn "conn"
        )
    SubConv cnv sub ->
      get
        ( galleyCall
            . paths
              [ "conversations",
                toByteString' (qDomain qcs),
                toByteString' cnv,
                "subconversations",
                toByteString' sub,
                "groupinfo"
              ]
            . zUser sender
            . zConn "conn"
        )

remoteGetGroupInfo ::
  Remote UserId ->
  Qualified ConvOrSubConvId ->
  TestM ByteString
remoteGetGroupInfo rusr qcs = do
  client <- view tsFedGalleyClient
  GetGroupInfoResponseState (Base64ByteString pgs) <-
    runFedClient
      @"query-group-info"
      client
      (tDomain rusr)
      GetGroupInfoRequest
        { conv = qUnqualified qcs,
          sender = tUnqualified rusr
        }
  pure pgs

getSelfConv ::
  UserId ->
  TestM ResponseLBS
getSelfConv u = do
  g <- viewGalley
  get $
    g
      . paths ["/conversations", "mls-self"]
      . zUser u
      . zConn "conn"
      . zType "access"

withMLSDisabled :: (HasSettingsOverrides m) => m a -> m a
withMLSDisabled = withSettingsOverrides noMLS
  where
    noMLS = Opts.settings . Opts.mlsPrivateKeyPaths .~ Nothing

getSubConv ::
  UserId ->
  Qualified ConvId ->
  SubConvId ->
  TestM ResponseLBS
getSubConv u qcnv sconv = do
  g <- viewGalley
  get $
    g
      . paths
        [ "conversations",
          toByteString' (qDomain qcnv),
          toByteString' (qUnqualified qcnv),
          "subconversations",
          LBS.toStrict (toLazyByteString (toEncodedUrlPiece sconv))
        ]
      . zUser u

deleteSubConv ::
  UserId ->
  Qualified ConvId ->
  SubConvId ->
  MLSReset ->
  TestM ResponseLBS
deleteSubConv u qcnv sconv reset = do
  g <- viewGalley
  delete $
    g
      . paths
        [ "conversations",
          toByteString' (qDomain qcnv),
          toByteString' (qUnqualified qcnv),
          "subconversations",
          LBS.toStrict (toLazyByteString (toEncodedUrlPiece sconv))
        ]
      . zUser u
      . contentJson
      . json reset

leaveSubConv ::
  UserId ->
  ClientId ->
  Qualified ConvId ->
  SubConvId ->
  TestM ResponseLBS
leaveSubConv u c qcnv subId = do
  g <- viewGalley
  delete $
    g
      . paths
        [ "conversations",
          toByteString' (qDomain qcnv),
          toByteString' (qUnqualified qcnv),
          "subconversations",
          toHeader subId,
          "self"
        ]
      . zUser u
      . zClient c

remoteLeaveCurrentConv ::
  Remote ClientIdentity ->
  Qualified ConvId ->
  SubConvId ->
  TestM ()
remoteLeaveCurrentConv rcid qcnv subId = do
  client <- view tsFedGalleyClient
  let lscr =
        LeaveSubConversationRequest
          { lscrUser = ciUser $ tUnqualified rcid,
            lscrClient = ciClient $ tUnqualified rcid,
            lscrConv = qUnqualified qcnv,
            lscrSubConv = subId
          }
  runFedClient
    @"leave-sub-conversation"
    client
    (tDomain rcid)
    lscr
    >>= liftIO . \case
      LeaveSubConversationResponseError e ->
        assertFailure $
          "error while leaving remote conversation: " <> show e
      LeaveSubConversationResponseProtocolError e ->
        assertFailure $
          "protocol error while leaving remote conversation: " <> T.unpack e
      LeaveSubConversationResponseOk -> pure ()

leaveCurrentConv ::
  (HasCallStack) =>
  ClientIdentity ->
  Qualified ConvOrSubConvId ->
  MLSTest ()
leaveCurrentConv cid qsub = case qUnqualified qsub of
  -- TODO: implement leaving main conversation as well
  Conv _ -> liftIO $ assertFailure "Leaving conversations is not supported"
  SubConv cnv subId -> do
    liftTest $ do
      loc <- qualifyLocal ()
      foldQualified
        loc
        ( \_ ->
            leaveSubConv (ciUser cid) (ciClient cid) (qsub $> cnv) subId
              !!! const 200 === statusCode
        )
        (\rcid -> remoteLeaveCurrentConv rcid (qsub $> cnv) subId)
        (cidQualifiedUser cid $> cid)
    State.modify $ \mls ->
      mls
        { mlsMembers = Set.difference (mlsMembers mls) (Set.singleton cid)
        }

getCurrentGroupId :: MLSTest GroupId
getCurrentGroupId = do
  State.gets mlsGroupId >>= \case
    Nothing -> liftIO $ assertFailure "Creating add proposal for non-existing group"
    Just g -> pure g

withTempKeyPackageFile :: ByteString -> ContT a MLSTest FilePath
withTempKeyPackageFile bs = do
  bd <- State.gets mlsBaseDir
  ContT $ \k ->
    bracket
      (liftIO (openBinaryTempFile bd "kp"))
      (\(fp, _) -> liftIO (removeFile fp))
      $ \(fp, h) -> do
        liftIO $ BS.hPut h bs `finally` hClose h
        k fp
