-- 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 Brig.Data.Connection
  ( -- * DB Operations
    insertConnection,
    updateConnection,
    updateConnectionStatus,
    lookupConnection,
    lookupLocalConnectionsPage,
    lookupRemoteConnectionsPage,
    lookupRelationWithHistory,
    lookupLocalConnections,
    lookupConnectionStatus,
    lookupConnectionStatus',
    lookupContactList,
    lookupContactListWithRelation,
    lookupLocalConnectionStatuses,
    lookupRemoteConnectionStatuses,
    lookupAllStatuses,
    lookupRemoteConnectedUsersPaginated,
    countConnections,
    deleteConnections,
    deleteRemoteConnections,
    remoteConnectionInsert,
    remoteConnectionSelect,
    remoteConnectionSelectFrom,
    remoteConnectionDelete,
    remoteConnectionClear,

    -- * Re-exports
    module T,
  )
where

import Brig.App
import Brig.Data.Types as T
import Cassandra
import Control.Monad.Morph
import Control.Monad.Trans.Maybe
import Data.Conduit (runConduit, (.|))
import Data.Conduit.List qualified as C
import Data.Domain (Domain)
import Data.Id
import Data.Json.Util (UTCTimeMillis, toUTCTimeMillis)
import Data.Qualified
import Data.Range
import Data.Time (getCurrentTime)
import Imports hiding (local)
import UnliftIO.Async (pooledForConcurrentlyN_, pooledMapConcurrentlyN, pooledMapConcurrentlyN_)
import Wire.API.Connection
import Wire.API.Routes.Internal.Brig.Connection

insertConnection ::
  (MonadClient m) =>
  Local UserId ->
  Qualified UserId ->
  RelationWithHistory ->
  Qualified ConvId ->
  m UserConnection
insertConnection self target rel qcnv@(Qualified cnv cdomain) = do
  now <- toUTCTimeMillis <$> liftIO getCurrentTime
  let local (tUnqualified -> ltarget) =
        write connectionInsert $
          params LocalQuorum (tUnqualified self, ltarget, rel, now, cnv)
  let remote (tUntagged -> Qualified rtarget domain) =
        write remoteConnectionInsert $
          params LocalQuorum (tUnqualified self, domain, rtarget, rel, now, cdomain, cnv)
  retry x5 $ foldQualified self local remote target
  pure $
    UserConnection
      { ucFrom = tUnqualified self,
        ucTo = target,
        ucStatus = relationDropHistory rel,
        ucLastUpdate = now,
        ucConvId = Just qcnv
      }

updateConnection :: (MonadClient m, MonadReader Env m) => UserConnection -> RelationWithHistory -> m UserConnection
updateConnection c status = do
  self <- qualifyLocal (ucFrom c)
  now <- updateConnectionStatus self (ucTo c) status
  pure $
    c
      { ucStatus = relationDropHistory status,
        ucLastUpdate = now
      }

updateConnectionStatus :: (MonadClient m) => Local UserId -> Qualified UserId -> RelationWithHistory -> m UTCTimeMillis
updateConnectionStatus self target status = do
  now <- toUTCTimeMillis <$> liftIO getCurrentTime
  let local (tUnqualified -> ltarget) =
        write connectionUpdate $
          params LocalQuorum (status, now, tUnqualified self, ltarget)
  let remote (tUntagged -> Qualified rtarget domain) =
        write remoteConnectionUpdate $
          params LocalQuorum (status, now, tUnqualified self, domain, rtarget)
  retry x5 $ foldQualified self local remote target
  pure now

-- | Lookup the connection from a user 'A' to a user 'B' (A -> B).
lookupConnection :: (MonadClient m) => Local UserId -> Qualified UserId -> m (Maybe UserConnection)
lookupConnection self target = runMaybeT $ do
  let local (tUnqualified -> ltarget) = do
        (_, _, rel, time, mcnv) <-
          MaybeT . query1 connectionSelect $
            params LocalQuorum (tUnqualified self, ltarget)
        pure (rel, time, fmap (tUntagged . qualifyAs self) mcnv)
  let remote (tUntagged -> Qualified rtarget domain) = do
        (rel, time, cdomain, cnv) <-
          MaybeT . query1 remoteConnectionSelectFrom $
            params LocalQuorum (tUnqualified self, domain, rtarget)
        pure (rel, time, Just (Qualified cnv cdomain))
  (rel, time, mqcnv) <- hoist (retry x1) $ foldQualified self local remote target
  pure $
    UserConnection
      { ucFrom = tUnqualified self,
        ucTo = target,
        ucStatus = relationDropHistory rel,
        ucLastUpdate = time,
        ucConvId = mqcnv
      }

-- | 'lookupConnection' with more 'Relation' info.
lookupRelationWithHistory ::
  (MonadClient m) =>
  -- | User 'A'
  Local UserId ->
  -- | User 'B'
  Qualified UserId ->
  m (Maybe RelationWithHistory)
lookupRelationWithHistory self target = do
  let local (tUnqualified -> ltarget) =
        query1 relationSelect (params LocalQuorum (tUnqualified self, ltarget))
  let remote (tUntagged -> Qualified rtarget domain) =
        query1 remoteRelationSelect (params LocalQuorum (tUnqualified self, domain, rtarget))
  runIdentity <$$> retry x1 (foldQualified self local remote target)

-- | For a given user 'A', lookup their outgoing connections (A -> X) to other users.
lookupLocalConnections ::
  (MonadClient m) =>
  Local UserId ->
  Maybe UserId ->
  Range 1 500 Int32 ->
  m (ResultPage UserConnection)
lookupLocalConnections lfrom start (fromRange -> size) =
  toResult <$> case start of
    Just u ->
      retry x1 $
        paginate connectionsSelectFrom (paramsP LocalQuorum (tUnqualified lfrom, u) (size + 1))
    Nothing ->
      retry x1 $
        paginate connectionsSelect (paramsP LocalQuorum (Identity (tUnqualified lfrom)) (size + 1))
  where
    toResult = cassandraResultPage . fmap (toLocalUserConnection lfrom) . trim
    trim p = p {result = take (fromIntegral size) (result p)}

-- | For a given user 'A', lookup their outgoing connections (A -> X) to other users.
-- Similar to lookupLocalConnections
lookupLocalConnectionsPage ::
  (MonadClient m) =>
  Local UserId ->
  Maybe PagingState ->
  Range 1 1000 Int32 ->
  m (PageWithState Void UserConnection)
lookupLocalConnectionsPage self pagingState (fromRange -> size) =
  fmap (toLocalUserConnection self) <$> paginateWithState connectionsSelect (paramsPagingState LocalQuorum (Identity (tUnqualified self)) size pagingState) x1

-- | For a given user 'A', lookup their outgoing connections (A -> X) to remote users.
lookupRemoteConnectionsPage ::
  (MonadClient m) =>
  Local UserId ->
  Maybe PagingState ->
  Int32 ->
  m (PageWithState Void UserConnection)
lookupRemoteConnectionsPage self pagingState size =
  fmap (toRemoteUserConnection self)
    <$> paginateWithState
      remoteConnectionSelect
      (paramsPagingState LocalQuorum (Identity (tUnqualified self)) size pagingState)
      x1

-- | Lookup all relations between two sets of users (cartesian product).
lookupConnectionStatus :: (MonadClient m) => [UserId] -> [UserId] -> m [ConnectionStatus]
lookupConnectionStatus from to =
  map toConnectionStatus
    <$> retry x1 (query connectionStatusSelect (params LocalQuorum (from, to)))

-- | Lookup all relations between two sets of users (cartesian product).
lookupConnectionStatus' :: (MonadClient m) => [UserId] -> m [ConnectionStatus]
lookupConnectionStatus' from =
  map toConnectionStatus
    <$> retry x1 (query connectionStatusSelect' (params LocalQuorum (Identity from)))

lookupLocalConnectionStatuses ::
  ( MonadClient m,
    MonadUnliftIO m
  ) =>
  [UserId] ->
  Local [UserId] ->
  m [ConnectionStatusV2]
lookupLocalConnectionStatuses froms tos = do
  concat <$> pooledMapConcurrentlyN 16 lookupStatuses froms
  where
    lookupStatuses :: (MonadClient m) => UserId -> m [ConnectionStatusV2]
    lookupStatuses from =
      map (uncurry $ toConnectionStatusV2 from (tDomain tos))
        <$> retry x1 (query relationsSelect (params LocalQuorum (from, tUnqualified tos)))

lookupRemoteConnectionStatuses ::
  ( MonadClient m,
    MonadUnliftIO m
  ) =>
  [UserId] ->
  Remote [UserId] ->
  m [ConnectionStatusV2]
lookupRemoteConnectionStatuses froms tos = do
  concat <$> pooledMapConcurrentlyN 16 lookupStatuses froms
  where
    lookupStatuses :: (MonadClient m) => UserId -> m [ConnectionStatusV2]
    lookupStatuses from =
      map (uncurry $ toConnectionStatusV2 from (tDomain tos))
        <$> retry x1 (query remoteRelationsSelect (params LocalQuorum (from, tDomain tos, tUnqualified tos)))

lookupAllStatuses ::
  ( MonadClient m,
    MonadUnliftIO m
  ) =>
  Local [UserId] ->
  m [ConnectionStatusV2]
lookupAllStatuses lfroms = do
  let froms = tUnqualified lfroms
  concat <$> pooledMapConcurrentlyN 16 lookupAndCombine froms
  where
    lookupAndCombine :: (MonadClient m) => UserId -> m [ConnectionStatusV2]
    lookupAndCombine u = (<>) <$> lookupLocalStatuses u <*> lookupRemoteStatuses u

    lookupLocalStatuses :: (MonadClient m) => UserId -> m [ConnectionStatusV2]
    lookupLocalStatuses from =
      map (uncurry $ toConnectionStatusV2 from (tDomain lfroms))
        <$> retry x1 (query relationsSelectAll (params LocalQuorum (Identity from)))
    lookupRemoteStatuses :: (MonadClient m) => UserId -> m [ConnectionStatusV2]
    lookupRemoteStatuses from =
      map (\(d, u, r) -> toConnectionStatusV2 from d u r)
        <$> retry x1 (query remoteRelationsSelectAll (params LocalQuorum (Identity from)))

lookupRemoteConnectedUsersPaginated :: (MonadClient m) => Local UserId -> Int32 -> m (Page (Remote UserConnection))
lookupRemoteConnectedUsersPaginated u maxResults = do
  (\x@(d, _, _, _, _, _) -> toRemoteUnsafe d (toRemoteUserConnection u x)) <$$> retry x1 (paginate remoteConnectionSelect (paramsP LocalQuorum (Identity (tUnqualified u)) maxResults))

-- | See 'lookupContactListWithRelation'.
lookupContactList :: (MonadClient m) => UserId -> m [UserId]
lookupContactList u =
  fst <$$> (filter ((== AcceptedWithHistory) . snd) <$> lookupContactListWithRelation u)

-- | For a given user 'A', lookup the list of users that form his contact list,
-- i.e. the users to whom 'A' has an outgoing 'Accepted' relation (A -> B).
lookupContactListWithRelation :: (MonadClient m) => UserId -> m [(UserId, RelationWithHistory)]
lookupContactListWithRelation u =
  retry x1 (query contactsSelect (params LocalQuorum (Identity u)))

-- | Count the number of connections a user has in a specific relation status.
-- (If you want to distinguish 'RelationWithHistory', write a new function.)
-- Note: The count is eventually consistent.
countConnections :: (MonadClient m) => Local UserId -> [Relation] -> m Int64
countConnections u r = do
  rels <- retry x1 . query selectStatus $ params One (Identity (tUnqualified u))
  relsRemote <- retry x1 . query selectStatusRemote $ params One (Identity (tUnqualified u))

  pure $ foldl' count 0 rels + foldl' count 0 relsRemote
  where
    selectStatus :: QueryString R (Identity UserId) (Identity RelationWithHistory)
    selectStatus = "SELECT status FROM connection WHERE left = ?"

    selectStatusRemote :: QueryString R (Identity UserId) (Identity RelationWithHistory)
    selectStatusRemote = "SELECT status FROM connection_remote WHERE left = ?"

    count n (Identity s) | relationDropHistory s `elem` r = n + 1
    count n _ = n

deleteConnections :: (MonadClient m, MonadUnliftIO m) => UserId -> m ()
deleteConnections u = do
  runConduit $
    paginateC contactsSelect (paramsP LocalQuorum (Identity u) 100) x1
      .| C.mapM_
        (pooledMapConcurrentlyN_ 16 delete)
  do
    retry x1 . write connectionClear $ params LocalQuorum (Identity u)
    retry x1 . write remoteConnectionClear $ params LocalQuorum (Identity u)
  where
    delete (other, _status) = write connectionDelete $ params LocalQuorum (other, u)

deleteRemoteConnections ::
  ( MonadClient m,
    MonadUnliftIO m
  ) =>
  Remote UserId ->
  Range 1 1000 [UserId] ->
  m ()
deleteRemoteConnections (tUntagged -> Qualified remoteUser remoteDomain) (fromRange -> locals) =
  pooledForConcurrentlyN_ 16 locals $ \u ->
    write remoteConnectionDelete $ params LocalQuorum (u, remoteDomain, remoteUser)

-- Queries

connectionInsert :: PrepQuery W (UserId, UserId, RelationWithHistory, UTCTimeMillis, ConvId) ()
connectionInsert = "INSERT INTO connection (left, right, status, last_update, conv) VALUES (?, ?, ?, ?, ?)"

connectionUpdate :: PrepQuery W (RelationWithHistory, UTCTimeMillis, UserId, UserId) ()
connectionUpdate = {- `IF EXISTS`, but that requires benchmarking -} "UPDATE connection SET status = ?, last_update = ? WHERE left = ? AND right = ?"

connectionSelect :: PrepQuery R (UserId, UserId) (UserId, UserId, RelationWithHistory, UTCTimeMillis, Maybe ConvId)
connectionSelect = "SELECT left, right, status, last_update, conv FROM connection WHERE left = ? AND right = ?"

relationSelect :: PrepQuery R (UserId, UserId) (Identity RelationWithHistory)
relationSelect = "SELECT status FROM connection WHERE left = ? AND right = ?"

relationsSelect :: PrepQuery R (UserId, [UserId]) (UserId, RelationWithHistory)
relationsSelect = "SELECT right, status FROM connection where left = ? AND right IN ?"

relationsSelectAll :: PrepQuery R (Identity UserId) (UserId, RelationWithHistory)
relationsSelectAll = "SELECT right, status FROM connection where left = ?"

-- FUTUREWORK: Delete this query, we shouldn't use `IN` with the primary key of
-- the table.
connectionStatusSelect :: PrepQuery R ([UserId], [UserId]) (UserId, UserId, RelationWithHistory)
connectionStatusSelect = "SELECT left, right, status FROM connection WHERE left IN ? AND right IN ?"

-- FUTUREWORK: Delete this query, we shouldn't use `IN` with the primary key of
-- the table.
connectionStatusSelect' :: PrepQuery R (Identity [UserId]) (UserId, UserId, RelationWithHistory)
connectionStatusSelect' = "SELECT left, right, status FROM connection WHERE left IN ?"

contactsSelect :: PrepQuery R (Identity UserId) (UserId, RelationWithHistory)
contactsSelect = "SELECT right, status FROM connection WHERE left = ?"

connectionsSelect :: PrepQuery R (Identity UserId) (UserId, UserId, RelationWithHistory, UTCTimeMillis, Maybe ConvId)
connectionsSelect = "SELECT left, right, status, last_update, conv FROM connection WHERE left = ? ORDER BY right ASC"

connectionsSelectFrom :: PrepQuery R (UserId, UserId) (UserId, UserId, RelationWithHistory, UTCTimeMillis, Maybe ConvId)
connectionsSelectFrom = "SELECT left, right, status, last_update, conv FROM connection WHERE left = ? AND right > ? ORDER BY right ASC"

connectionDelete :: PrepQuery W (UserId, UserId) ()
connectionDelete = "DELETE FROM connection WHERE left = ? AND right = ?"

connectionClear :: PrepQuery W (Identity UserId) ()
connectionClear = "DELETE FROM connection WHERE left = ?"

-- Remote connections

remoteConnectionInsert :: PrepQuery W (UserId, Domain, UserId, RelationWithHistory, UTCTimeMillis, Domain, ConvId) ()
remoteConnectionInsert = "INSERT INTO connection_remote (left, right_domain, right_user, status, last_update, conv_domain, conv_id) VALUES (?, ?, ?, ?, ?, ?, ?)"

remoteConnectionSelect :: PrepQuery R (Identity UserId) (Domain, UserId, RelationWithHistory, UTCTimeMillis, Domain, ConvId)
remoteConnectionSelect = "SELECT right_domain, right_user, status, last_update, conv_domain, conv_id FROM connection_remote where left = ?"

remoteConnectionSelectFrom :: PrepQuery R (UserId, Domain, UserId) (RelationWithHistory, UTCTimeMillis, Domain, ConvId)
remoteConnectionSelectFrom = "SELECT status, last_update, conv_domain, conv_id FROM connection_remote where left = ? AND right_domain = ? AND right_user = ?"

remoteConnectionUpdate :: PrepQuery W (RelationWithHistory, UTCTimeMillis, UserId, Domain, UserId) ()
remoteConnectionUpdate = {- `IF EXISTS`, but that requires benchmarking -} "UPDATE connection_remote set status = ?, last_update = ? WHERE left = ? and right_domain = ? and right_user = ?"

remoteConnectionDelete :: PrepQuery W (UserId, Domain, UserId) ()
remoteConnectionDelete = "DELETE FROM connection_remote where left = ? AND right_domain = ? AND right_user = ?"

remoteConnectionClear :: PrepQuery W (Identity UserId) ()
remoteConnectionClear = "DELETE FROM connection_remote where left = ?"

remoteRelationSelect :: PrepQuery R (UserId, Domain, UserId) (Identity RelationWithHistory)
remoteRelationSelect = "SELECT status FROM connection_remote WHERE left = ? AND right_domain = ? AND right_user = ?"

remoteRelationsSelect :: PrepQuery R (UserId, Domain, [UserId]) (UserId, RelationWithHistory)
remoteRelationsSelect = "SELECT right_user, status FROM connection_remote WHERE left = ? AND right_domain = ? AND right_user IN ?"

remoteRelationsSelectAll :: PrepQuery R (Identity UserId) (Domain, UserId, RelationWithHistory)
remoteRelationsSelectAll = "SELECT right_domain, right_user, status FROM connection_remote WHERE left = ?"

-- Conversions

toLocalUserConnection ::
  Local x ->
  (UserId, UserId, RelationWithHistory, UTCTimeMillis, Maybe ConvId) ->
  UserConnection
toLocalUserConnection loc (l, r, relationDropHistory -> rel, time, cid) =
  UserConnection l (tUntagged (qualifyAs loc r)) rel time (fmap (tUntagged . qualifyAs loc) cid)

toRemoteUserConnection ::
  Local UserId ->
  (Domain, UserId, RelationWithHistory, UTCTimeMillis, Domain, ConvId) ->
  UserConnection
toRemoteUserConnection l (rDomain, r, relationDropHistory -> rel, time, cDomain, cid) =
  UserConnection (tUnqualified l) (Qualified r rDomain) rel time (Just $ Qualified cid cDomain)

toConnectionStatus :: (UserId, UserId, RelationWithHistory) -> ConnectionStatus
toConnectionStatus (l, r, relationDropHistory -> rel) = ConnectionStatus l r rel

toConnectionStatusV2 :: UserId -> Domain -> UserId -> RelationWithHistory -> ConnectionStatusV2
toConnectionStatusV2 from toDomain toUser relWithHistory =
  ConnectionStatusV2 from (Qualified toUser toDomain) (relationDropHistory relWithHistory)
