-- 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/>.
{-# 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.ConversationStore.MLS.Types where

import Data.Bifunctor
import Data.Domain
import Data.Id
import Data.IntMap (IntMap)
import Data.IntMap qualified as IntMap
import Data.IntSet qualified as IntSet
import Data.Map qualified as Map
import Data.Qualified
import GHC.Records (HasField (..))
import Imports
import Wire.API.Conversation
import Wire.API.Conversation.Protocol
import Wire.API.MLS.CipherSuite
import Wire.API.MLS.Credential
import Wire.API.MLS.Group.Serialisation qualified as Group
import Wire.API.MLS.LeafNode
import Wire.API.MLS.SubConversation
import Wire.StoredConversation

-- | A map of leaf index to members.
--
-- This is used to reconstruct client
-- identities from leaf indices in remove proposals, as well as to allocate new
-- indices for added clients.
--
-- Note that clients that are in the process of being removed from a group
-- (i.e. there is a pending remove proposals for them) are included in this
-- mapping.
newtype IndexMap = IndexMap {unIndexMap :: IntMap GroupMember}
  deriving (Eq, Show)
  deriving newtype (Semigroup, Monoid)

mkIndexMapFromParts ::
  [(Domain, UserId, ClientId, Int32, Bool)] ->
  [(HistoryClientId, Int32, Bool)] ->
  IndexMap
mkIndexMapFromParts regular history =
  IndexMap
    . flip (foldr addHistoryClient) history
    . flip (foldr addRegularClient) regular
    $ mempty
  where
    addHistoryClient (h, leafidx, _) =
      IntMap.insert (fromIntegral leafidx) (HistoryClient h)
    addRegularClient (dom, usr, c, leafidx, _) =
      IntMap.insert (fromIntegral leafidx) (RegularClient (ClientIdentity dom usr c))

imLookup :: IndexMap -> LeafIndex -> Maybe GroupMember
imLookup m i = IntMap.lookup (fromIntegral i) (unIndexMap m)

imFromList :: [(LeafIndex, GroupMember)] -> IndexMap
imFromList = IndexMap . IntMap.fromList . map (first fromIntegral)

imNextIndex :: IndexMap -> LeafIndex
imNextIndex im =
  fromIntegral . fromJust $
    find (\n -> not $ IntMap.member n (unIndexMap im)) [0 ..]

imAddClient :: IndexMap -> GroupMember -> (LeafIndex, IndexMap)
imAddClient im cid = let idx = imNextIndex im in (idx, IndexMap $ IntMap.insert (fromIntegral idx) cid $ unIndexMap im)

imRemoveClient :: IndexMap -> LeafIndex -> Maybe (GroupMember, IndexMap)
imRemoveClient im idx = do
  cid <- imLookup im idx
  pure (cid, IndexMap . IntMap.delete (fromIntegral idx) $ unIndexMap im)

imRemoveIndices :: [LeafIndex] -> IndexMap -> IndexMap
imRemoveIndices keys =
  IndexMap
    . flip IntMap.withoutKeys (IntSet.fromList (map fromIntegral keys))
    . unIndexMap

imAssocs :: IndexMap -> [(Int, GroupMember)]
imAssocs = IntMap.assocs . unIndexMap

-- | A two-level map of users to clients to leaf indices.
--
-- This is used to keep track of the state of an MLS group for e.g. propagating
-- a message to all the clients that are supposed to receive it.
--
-- Note that clients that are in the process of being removed from a group
-- (i.e. there is a pending remove proposals for them) are __not__ included in
-- this mapping.
newtype ClientMap a = ClientMap
  { unClientMap :: Map (Qualified UserId) (Map ClientId a)
  }
  deriving (Show, Eq, Functor)

instance Semigroup (ClientMap a) where
  ClientMap cm1 <> ClientMap cm2 = ClientMap $ Map.unionWith Map.union cm1 cm2

instance Monoid (ClientMap a) where
  mempty = ClientMap mempty

mkClientMap :: [(Domain, UserId, ClientId, Int32, Bool)] -> ClientMap LeafIndex
mkClientMap = foldr addEntry mempty
  where
    addEntry :: (Domain, UserId, ClientId, Int32, Bool) -> ClientMap LeafIndex -> ClientMap LeafIndex
    addEntry (dom, usr, c, leafidx, pending_removal)
      | pending_removal = id -- treat as removed, don't add to ClientMap
      | otherwise = ClientMap . Map.insertWith (<>) (Qualified usr dom) (Map.singleton c (fromIntegral leafidx)) . unClientMap

cmNull :: ClientMap a -> Bool
cmNull (ClientMap cm) = Map.null cm

cmToMap :: (Ord a) => ClientMap a -> Map a ClientIdentity
cmToMap = Map.fromList . map swap . cmAssocs

cmLookupIndex :: ClientIdentity -> ClientMap a -> Maybe a
cmLookupIndex cid (ClientMap cm) = do
  clients <- Map.lookup (cidQualifiedUser cid) cm
  Map.lookup (ciClient cid) clients

cmLookup :: Qualified UserId -> ClientMap a -> Maybe (Map ClientId a)
cmLookup quid (ClientMap cm) = Map.lookup quid cm

cmLookupClients :: Qualified UserId -> ClientMap a -> [ClientId]
cmLookupClients quid (ClientMap cm) = foldMap Map.keys (Map.lookup quid cm)

cmRemoveClient :: ClientIdentity -> ClientMap a -> ClientMap a
cmRemoveClient cid (ClientMap cm) = ClientMap $
  case Map.lookup (cidQualifiedUser cid) cm of
    Nothing -> cm
    Just clients ->
      let clients' = Map.delete (ciClient cid) clients
       in if Map.null clients'
            then Map.delete (cidQualifiedUser cid) cm
            else Map.insert (cidQualifiedUser cid) clients' cm

isClientMember :: ClientIdentity -> ClientMap LeafIndex -> Bool
isClientMember ci = isJust . cmLookupIndex ci

cmAssocs :: ClientMap a -> [(ClientIdentity, a)]
cmAssocs (ClientMap cm) = do
  (quid, clients) <- Map.assocs cm
  (clientId, idx) <- Map.assocs clients
  pure (mkClientIdentity quid clientId, idx)

cmIdentities :: ClientMap a -> [ClientIdentity]
cmIdentities = map fst . cmAssocs

cmSingleton :: ClientIdentity -> a -> ClientMap a
cmSingleton cid idx =
  ClientMap $
    Map.singleton
      (cidQualifiedUser cid)
      (Map.singleton (ciClient cid) idx)

-- | Construct the subset of the first client map consisting of users that are present in the second.
cmIntersect :: ClientMap a -> ClientMap b -> ClientMap a
cmIntersect (ClientMap m1) (ClientMap m2) = ClientMap (Map.intersection m1 m2)

-- | Inform a handler for 'POST /conversations/list-ids' if the MLS global team
-- conversation and the MLS self-conversation should be included in the
-- response.
data ListGlobalSelfConvs = ListGlobalSelf | DoNotListGlobalSelf
  deriving (Eq)

data MLSConversation = MLSConversation
  { mcId :: ConvId,
    mcMetadata :: ConversationMetadata GroupConvType,
    mcMLSData :: ConversationMLSData,
    mcLocalMembers :: [LocalMember],
    mcRemoteMembers :: [RemoteMember],
    mcMembers :: ClientMap LeafIndex,
    mcIndexMap :: IndexMap,
    mcMigrationState :: MLSMigrationState
  }
  deriving (Show)

data SubConversation = SubConversation
  { scParentConvId :: ConvId,
    scSubConvId :: SubConvId,
    scMLSData :: ConversationMLSData,
    scMembers :: ClientMap LeafIndex,
    scIndexMap :: IndexMap
  }
  deriving (Eq, Show)

newSubConversation :: ConvId -> SubConvId -> GroupId -> SubConversation
newSubConversation convId subConvId groupId =
  SubConversation
    { scParentConvId = convId,
      scSubConvId = subConvId,
      scMLSData =
        ConversationMLSData
          { cnvmlsGroupId = groupId,
            cnvmlsActiveData = Nothing
          },
      scMembers = mkClientMap [],
      scIndexMap = mempty
    }

newSubConversationFromParent ::
  Local ConvId ->
  SubConvId ->
  SubConversation
newSubConversationFromParent lconv sconv =
  let qcs = flip SubConv sconv <$> tUntagged lconv
      groupId = Group.newGroupId RegularConv qcs
   in newSubConversation (tUnqualified lconv) sconv groupId

toPublicSubConv :: Qualified SubConversation -> PublicSubConversation
toPublicSubConv (Qualified (SubConversation {..}) domain) =
  let members = map fst (cmAssocs scMembers)
   in PublicSubConversation
        { pscParentConvId = Qualified scParentConvId domain,
          pscSubConvId = scSubConvId,
          pscGroupId = cnvmlsGroupId scMLSData,
          pscActiveData = cnvmlsActiveData scMLSData,
          pscMembers = members
        }

type ConvOrSubConv = ConvOrSubChoice MLSConversation SubConversation

instance HasField "meta" ConvOrSubConv (ConversationMetadata GroupConvType) where
  getField x = x.conv.mcMetadata

instance HasField "mlsMeta" ConvOrSubConv ConversationMLSData where
  getField (Conv c) = mcMLSData c
  getField (SubConv _ s) = scMLSData s

instance HasField "members" ConvOrSubConv (ClientMap LeafIndex) where
  getField (Conv c) = mcMembers c
  getField (SubConv _ s) = scMembers s

instance HasField "indexMap" ConvOrSubConv IndexMap where
  getField (Conv c) = mcIndexMap c
  getField (SubConv _ s) = scIndexMap s

instance HasField "id" ConvOrSubConv ConvOrSubConvId where
  getField (Conv c) = Conv (mcId c)
  getField (SubConv c s) = SubConv (mcId c) (scSubConvId s)

instance HasField "migrationState" ConvOrSubConv MLSMigrationState where
  getField (Conv c) = c.mcMigrationState
  getField (SubConv _ _) = MLSMigrationMLS

instance HasField "ciphersuite" ConvOrSubConv (Maybe CipherSuiteTag) where
  getField x = x.mlsMeta.ciphersuite
