-- 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.InvitationStore.Cassandra
  ( interpretInvitationStoreToCassandra,
  )
where

import Cassandra
import Control.Monad.Trans.Maybe
import Data.Conduit (runConduit, (.|))
import Data.Conduit.List qualified as Conduit
import Data.Id
import Data.Json.Util (UTCTimeMillis)
import Data.Range (Range, fromRange)
import Database.CQL.Protocol (Record (..), TupleType, asRecord)
import Imports
import Polysemy
import Polysemy.Embed
import UnliftIO.Async (pooledMapConcurrentlyN_)
import Util.Timeout
import Wire.API.Team.Role (Role)
import Wire.API.User
import Wire.InvitationStore

interpretInvitationStoreToCassandra :: (Member (Embed IO) r) => ClientState -> InterpreterFor InvitationStore r
interpretInvitationStoreToCassandra casClient =
  interpret $
    runEmbedded (runClient casClient) . \case
      InsertInvitation newInv timeout -> embed $ insertInvitationImpl newInv timeout
      LookupInvitation tid iid -> embed $ lookupInvitationImpl tid iid
      LookupInvitationsByEmail email -> embed $ lookupInvitationsByEmailImpl email
      InsertPendingScimUser tid email uid -> embed $ insertPendingScimUserImpl tid email uid
      LookupPendingScimUsers tid email -> embed $ lookupPendingScimUsersImpl tid email
      DeletePendingScimUser tid email uid -> embed $ deletePendingScimUserImpl tid email uid
      LookupInvitationByCode code -> embed $ lookupInvitationByCodeImpl code
      LookupInvitationsPaginated mSize tid miid -> embed $ lookupInvitationsPaginatedImpl mSize tid miid
      CountInvitations tid -> embed $ countInvitationsImpl tid
      DeleteInvitation tid invId -> embed $ deleteInvitationImpl tid invId
      DeleteAllTeamInvitations tid -> embed $ deleteInvitationsImpl tid

insertInvitationImpl ::
  InsertInvitation ->
  -- | The timeout for the invitation code.
  Timeout ->
  Client StoredInvitation
insertInvitationImpl (insertInvToStoredInv -> inv@(MkStoredInvitation teamId role invId now uid email name code)) timeout = do
  retry x5 . batch $ do
    setType BatchLogged
    setConsistency LocalQuorum
    addPrepQuery cqlInsert (teamId, role, invId, now, uid, email, name, code, round timeout)
    addPrepQuery cqlInsertInfo (code, teamId, invId, round timeout)
    addPrepQuery cqlInsertByEmail (email, teamId, invId, code, round timeout)
  pure inv
  where
    cqlInsert :: PrepQuery W (TeamId, Maybe Role, InvitationId, UTCTimeMillis, Maybe UserId, EmailAddress, Maybe Name, InvitationCode, Int32) ()
    cqlInsert =
      [sql|
        INSERT INTO team_invitation (team, role, id, created_at, created_by, email, name, code) VALUES (?, ?, ?, ?, ?, ?, ?, ?) USING TTL ?
      |]
    cqlInsertInfo :: PrepQuery W (InvitationCode, TeamId, InvitationId, Int32) ()
    cqlInsertInfo =
      [sql|
            INSERT INTO team_invitation_info (code, team, id) VALUES (?, ?, ?) USING TTL ?
          |]
    -- Note: the edge case of multiple invites to the same team by different admins from the
    -- same team results in last-invite-wins in the team_invitation_email table.
    cqlInsertByEmail :: PrepQuery W (EmailAddress, TeamId, InvitationId, InvitationCode, Int32) ()
    cqlInsertByEmail =
      [sql|
            INSERT INTO team_invitation_email (email, team, invitation, code) VALUES (?, ?, ?, ?) USING TTL ?
          |]

lookupInvitationsPaginatedImpl :: Maybe (Range 1 500 Int32) -> TeamId -> Maybe InvitationId -> Client (PaginatedResult [StoredInvitation])
lookupInvitationsPaginatedImpl mSize tid miid = do
  page <- retry x1 case miid of
    Just ref -> paginate cqlSelectFrom (paramsP LocalQuorum (tid, ref) (pageSize + 1))
    Nothing -> paginate cqlSelect (paramsP LocalQuorum (Identity tid) (pageSize + 1))
  pure $ mkPage (hasMore page) $ map asRecord $ trim page
  where
    pageSize :: Int32
    pageSize = maybe 100 fromRange mSize

    trim :: Page a -> [a]
    trim p = take (fromIntegral pageSize) (result p)

    mkPage more invs = if more then PaginatedResultHasMore invs else PaginatedResult invs

    cqlSelect :: PrepQuery R (Identity TeamId) (TeamId, Maybe Role, InvitationId, UTCTimeMillis, Maybe UserId, EmailAddress, Maybe Name, InvitationCode)
    cqlSelect =
      [sql|
      SELECT team, role, id, created_at, created_by, email, name, code FROM team_invitation WHERE team = ? ORDER BY id ASC
      |]
    cqlSelectFrom :: PrepQuery R (TeamId, InvitationId) (TeamId, Maybe Role, InvitationId, UTCTimeMillis, Maybe UserId, EmailAddress, Maybe Name, InvitationCode)
    cqlSelectFrom =
      [sql|
      SELECT team, role, id, created_at, created_by, email, name, code FROM team_invitation WHERE team = ? AND id > ? ORDER BY id ASC
      |]

countInvitationsImpl :: TeamId -> Client (Int64)
countInvitationsImpl t =
  maybe 0 runIdentity
    <$> retry x1 (query1 cql (params LocalQuorum (Identity t)))
  where
    cql :: PrepQuery R (Identity TeamId) (Identity Int64)
    cql = [sql| SELECT count(*) FROM team_invitation WHERE team = ?|]

lookupInvitationByCodeImpl :: InvitationCode -> Client (Maybe StoredInvitation)
lookupInvitationByCodeImpl code = runMaybeT do
  (teamId, invId, _) <-
    MaybeT $
      retry x1 (query1 cqlInfo (params LocalQuorum (Identity code)))
  MaybeT $ fmap asRecord <$> retry x1 (query1 cqlMain (params LocalQuorum (teamId, invId)))
  where
    cqlInfo :: PrepQuery R (Identity InvitationCode) (TeamId, InvitationId, InvitationCode)
    cqlInfo =
      [sql|
      SELECT team, id, code FROM team_invitation_info WHERE code = ?
      |]
    cqlMain :: PrepQuery R (TeamId, InvitationId) (TupleType StoredInvitation)
    cqlMain =
      [sql|
      SELECT team, role, id, created_at, created_by, email, name, code FROM team_invitation WHERE team = ? AND id = ?
      |]

lookupInvitationsByEmailImpl :: EmailAddress -> Client [StoredInvitation]
lookupInvitationsByEmailImpl email = do
  infoList <-
    retry x1 (query cqlInfo (params LocalQuorum (Identity email)))
  fmap catMaybes $ forM infoList $ \(tid, invId, _invCode) ->
    fmap asRecord <$> retry x1 (query1 cqlMain (params LocalQuorum (tid, invId)))
  where
    cqlInfo :: PrepQuery R (Identity EmailAddress) (TeamId, InvitationId, InvitationCode)
    cqlInfo =
      [sql|
      SELECT team, invitation, code FROM team_invitation_email WHERE email = ?
      |]
    cqlMain :: PrepQuery R (TeamId, InvitationId) (TupleType StoredInvitation)
    cqlMain =
      [sql|
      SELECT team, role, id, created_at, created_by, email, name, code FROM team_invitation WHERE team = ? AND id = ?
      |]

insertPendingScimUserImpl :: TeamId -> EmailAddress -> UserId -> Client ()
insertPendingScimUserImpl team email uid =
  retry x5 $ write cql (params LocalQuorum (team, email, uid))
  where
    cql :: PrepQuery W (TeamId, EmailAddress, UserId) ()
    cql =
      [sql|
        INSERT INTO team_scim_pending_user_email (team, email, user) VALUES (?, ?, ?)
      |]

lookupPendingScimUsersImpl :: TeamId -> EmailAddress -> Client [UserId]
lookupPendingScimUsersImpl team email =
  map runIdentity <$> retry x1 (query cql (params LocalQuorum (team, email)))
  where
    cql :: PrepQuery R (TeamId, EmailAddress) (Identity UserId)
    cql =
      [sql|
        SELECT user FROM team_scim_pending_user_email WHERE team = ? AND email = ?
      |]

deletePendingScimUserImpl :: TeamId -> EmailAddress -> UserId -> Client ()
deletePendingScimUserImpl team email uid =
  retry x5 $ write cql (params LocalQuorum (team, email, uid))
  where
    cql :: PrepQuery W (TeamId, EmailAddress, UserId) ()
    cql =
      [sql|
        DELETE FROM team_scim_pending_user_email WHERE team = ? AND email = ? AND user = ?
      |]

lookupInvitationImpl :: TeamId -> InvitationId -> Client (Maybe StoredInvitation)
lookupInvitationImpl tid iid =
  fmap asRecord
    <$> retry x1 (query1 cql (params LocalQuorum (tid, iid)))
  where
    cql :: PrepQuery R (TeamId, InvitationId) (TupleType StoredInvitation)
    cql =
      [sql|
      SELECT team, role, id, created_at, created_by, email, name, code FROM team_invitation WHERE team = ? AND id = ?
      |]

deleteInvitationImpl :: TeamId -> InvitationId -> Client ()
deleteInvitationImpl teamId invId = do
  codeEmail <- lookupInvitationCodeEmail
  case codeEmail of
    Just (invCode, invEmail) -> retry x5 . batch $ do
      setType BatchLogged
      setConsistency LocalQuorum
      addPrepQuery cqlInvitation (teamId, invId)
      addPrepQuery cqlInvitationInfo (Identity invCode)
      addPrepQuery cqlInvitationEmail (invEmail, teamId)
    Nothing ->
      retry x5 $ write cqlInvitation (params LocalQuorum (teamId, invId))
  where
    lookupInvitationCodeEmail :: Client (Maybe (InvitationCode, EmailAddress))
    lookupInvitationCodeEmail = retry x1 (query1 cqlInvitationCodeEmail (params LocalQuorum (teamId, invId)))

    cqlInvitation :: PrepQuery W (TeamId, InvitationId) ()
    cqlInvitation =
      [sql|
        DELETE FROM team_invitation where team = ? AND id = ?
      |]

    cqlInvitationInfo :: PrepQuery W (Identity InvitationCode) ()
    cqlInvitationInfo =
      [sql|
        DELETE FROM team_invitation_info WHERE code = ?
      |]

    cqlInvitationEmail :: PrepQuery W (EmailAddress, TeamId) ()
    cqlInvitationEmail =
      [sql|
        DELETE FROM team_invitation_email WHERE email = ? AND team = ?
      |]

    cqlInvitationCodeEmail :: PrepQuery R (TeamId, InvitationId) (InvitationCode, EmailAddress)
    cqlInvitationCodeEmail =
      [sql|
        SELECT code, email FROM team_invitation WHERE team = ? AND id = ?
      |]

deleteInvitationsImpl :: TeamId -> Client ()
deleteInvitationsImpl teamId =
  runConduit $
    paginateC cqlSelect (paramsP LocalQuorum (Identity teamId) 100) x1
      .| Conduit.mapM_ (pooledMapConcurrentlyN_ 16 (deleteInvitationImpl teamId . runIdentity))
  where
    cqlSelect :: PrepQuery R (Identity TeamId) (Identity InvitationId)
    cqlSelect = "SELECT id FROM team_invitation WHERE team = ? ORDER BY id ASC"
