{-# 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.UserGroupStore.Postgres where

import Control.Error (MaybeT (..))
import Data.Bifunctor (second)
import Data.Id
import Data.Json.Util
import Data.Map qualified as Map
import Data.Profunctor
import Data.Qualified (Local, QualifiedWithTag (tUntagged), inputQualifyLocal, qualifyAs)
import Data.Text qualified as T
import Data.Time
import Data.UUID as UUID
import Data.Vector (Vector)
import Data.Vector qualified as V
import Hasql.Decoders qualified as HD
import Hasql.Session
import Hasql.Statement
import Hasql.TH
import Hasql.Transaction qualified as Tx
import Hasql.Transaction.Sessions qualified as TxSessions
import Imports
import Polysemy
import Polysemy.Input
import Wire.API.Pagination
import Wire.API.PostgresMarshall
import Wire.API.User.Profile
import Wire.API.UserGroup hiding (UpdateUserGroupChannels)
import Wire.API.UserGroup.Pagination
import Wire.PaginationState
import Wire.Postgres
import Wire.UserGroupStore (UserGroupStore (..))

type UserGroupStorePostgresEffectConstraints r =
  PGConstraints r

interpretUserGroupStoreToPostgres ::
  forall r.
  (UserGroupStorePostgresEffectConstraints r, Member (Input (Local ())) r) =>
  InterpreterFor UserGroupStore r
interpretUserGroupStoreToPostgres =
  interpret $ \case
    CreateUserGroup team newUserGroup managedBy -> createUserGroup team newUserGroup managedBy
    GetUserGroup team userGroupId includeChannels -> getUserGroup team userGroupId includeChannels
    GetUserGroups tid req -> getUserGroups tid req
    GetUserGroupsWithMembers tid req -> getUserGroupsWithMembers tid req
    GetUserGroupsForConv convId -> getUserGroupsForConv convId
    UpdateUserGroup tid gid gup -> updateGroup tid gid gup
    DeleteUserGroup tid gid -> deleteGroup tid gid
    AddUser gid uid -> addUser gid uid
    UpdateUsers gid uids -> updateUsers gid uids
    RemoveUser gid uid -> removeUser gid uid
    AddUserGroupChannels gid convIds -> updateUserGroupChannels True gid convIds
    UpdateUserGroupChannels gid convIds -> updateUserGroupChannels False gid convIds
    GetUserGroupIdsForUsers uids -> getUserGroupIdsForUsers uids
    GetUserGroupChannels tid gid -> getUserGroupChannels tid gid

getUserGroupsForConv :: (UserGroupStorePostgresEffectConstraints r) => ConvId -> Sem r (Vector UserGroup)
getUserGroupsForConv convId = do
  runStatement convId stmt
  where
    decodeRow :: (UUID, Text, Int32, UTCTime, Vector UUID) -> Either Text UserGroup
    decodeRow (uuid, nameTxt, managedByInt, createdAtUtc, memberIds) = do
      name <- userGroupNameFromText nameTxt
      managedBy <- managedByFromInt32 managedByInt
      let id_ = Id uuid
          members = Identity (fmap Id memberIds)
          membersCount = Nothing
          channels = Nothing
          channelsCount = Nothing
          createdAt = toUTCTimeMillis createdAtUtc
      pure UserGroup_ {..}

    stmt :: Statement ConvId (Vector UserGroup)
    stmt =
      lmap (.toUUID)
        . refineResult (mapM decodeRow)
        $ [vectorStatement|
            select (ug.id :: uuid), (ug.name :: text), (ug.managed_by :: int), (ug.created_at :: timestamptz),
                   coalesce((select array_agg(ugm.user_id) from user_group_member ugm where ugm.user_group_id = ug.id), array[]::uuid[]) :: uuid[]
            from user_group ug
            inner join user_group_channel ugc on ugc.user_group_id = ug.id
            where ugc.conv_id = ($1 :: uuid)
          |]

getUserGroupIdsForUsers :: (UserGroupStorePostgresEffectConstraints r) => [UserId] -> Sem r (Map UserId [UserGroupId])
getUserGroupIdsForUsers uidsList = do
  rows <- runStatement (V.fromList (toUUID <$> uidsList)) stmt
  pure $ Map.fromListWith (<>) [(Id u, [Id g]) | (u, g) <- V.toList rows]
  where
    stmt :: Statement (Vector UUID) (Vector (UUID, UUID))
    stmt =
      [vectorStatement|
            select (ugm.user_id :: uuid), (ugm.user_group_id :: uuid)
              from user_group_member as ugm
              where ugm.user_id = ANY (($1 :: uuid[]))
          |]

updateUsers :: (UserGroupStorePostgresEffectConstraints r) => UserGroupId -> Vector UserId -> Sem r ()
updateUsers gid uids = do
  runTransactionWithRetry TxSessions.Serializable TxSessions.Write do
    Tx.statement gid deleteAllUsersStatement
    Tx.statement (toUUID gid, uids) insertGroupMembersStatement
  where
    deleteAllUsersStatement :: Statement UserGroupId ()
    deleteAllUsersStatement =
      dimap (.toUUID) (const ()) $
        [resultlessStatement|
          delete from user_group_member where user_group_id = ($1 :: uuid)
          |]

getUserGroup ::
  forall r.
  (UserGroupStorePostgresEffectConstraints r, Member (Input (Local ())) r) =>
  TeamId ->
  UserGroupId ->
  Bool ->
  Sem r (Maybe UserGroup)
getUserGroup team id_ includeChannels = do
  loc <- inputQualifyLocal ()
  runSessionWithRetry $
    if includeChannels
      then sessionWithChannels loc
      else session
  where
    session :: Session (Maybe UserGroup)
    session = runMaybeT do
      (name, managedBy, createdAt) <- MaybeT $ statement (id_, team) getGroupMetadataStatement
      members <- lift $ Identity <$> statement id_ getGroupMembersStatement
      let membersCount = Just . V.length $ runIdentity members
          channelsCount = Nothing
          channels = mempty
      pure $ UserGroup_ {..}

    sessionWithChannels :: Local a -> Session (Maybe UserGroup)
    sessionWithChannels loc = runMaybeT do
      (name, managedBy, createdAt, memberIds, channelIds) <- MaybeT $ statement (id_, team) getGroupWithMembersAndChannelsStatement
      let members = Identity (fmap Id memberIds)
          membersCount = Just $ V.length memberIds
          channels = Just (fmap (tUntagged . qualifyAs loc . Id) channelIds)
          channelsCount = Just $ V.length channelIds
      pure $ UserGroup_ {..}

    decodeMetadataRow :: (Text, Int32, UTCTime) -> Either Text (UserGroupName, ManagedBy, UTCTimeMillis)
    decodeMetadataRow (name, managedByInt, utcTime) =
      (,,toUTCTimeMillis utcTime)
        <$> userGroupNameFromText name
        <*> managedByFromInt32 managedByInt

    decodeWithArrays :: (Text, Int32, UTCTime, Vector UUID, Vector UUID) -> Either Text (UserGroupName, ManagedBy, UTCTimeMillis, Vector UUID, Vector UUID)
    decodeWithArrays (name, managedByInt, utcTime, membs, chans) = do
      n <- userGroupNameFromText name
      m <- managedByFromInt32 managedByInt
      pure (n, m, toUTCTimeMillis utcTime, membs, chans)

    getGroupMetadataStatement :: Statement (UserGroupId, TeamId) (Maybe (UserGroupName, ManagedBy, UTCTimeMillis))
    getGroupMetadataStatement =
      lmap (\(gid, tid) -> (gid.toUUID, tid.toUUID))
        . refineResult (mapM decodeMetadataRow)
        $ [maybeStatement|
         select (name :: text), (managed_by :: int), (created_at :: timestamptz)
           from user_group where id = ($1 :: uuid) AND team_id = ($2 :: uuid)
         |]

    getGroupMembersStatement :: Statement UserGroupId (Vector UserId)
    getGroupMembersStatement =
      dimap (.toUUID) (fmap Id) $
        [vectorStatement|
          select (user_id :: uuid) from user_group_member where user_group_id = ($1 :: uuid)
          |]

    getGroupWithMembersAndChannelsStatement :: Statement (UserGroupId, TeamId) (Maybe (UserGroupName, ManagedBy, UTCTimeMillis, Vector UUID, Vector UUID))
    getGroupWithMembersAndChannelsStatement =
      lmap (\(gid, tid) -> (gid.toUUID, tid.toUUID))
        . refineResult (mapM decodeWithArrays)
        $ [maybeStatement|
            select
              (name :: text),
              (managed_by :: int),
              (created_at :: timestamptz),
              coalesce((select array_agg(ugm.user_id) from user_group_member ugm where ugm.user_group_id = ug.id), array[]::uuid[]) :: uuid[],
              coalesce((select array_agg(ugc.conv_id) from user_group_channel ugc where ugc.user_group_id = ug.id), array[]::uuid[]) :: uuid[]
            from user_group ug
            where ug.id = ($1 :: uuid) and ug.team_id = ($2 :: uuid)
          |]

getUserGroupsWithMembers ::
  forall r.
  (UserGroupStorePostgresEffectConstraints r) =>
  TeamId ->
  UserGroupPageRequest ->
  Sem r UserGroupPageWithMembers
getUserGroupsWithMembers tid req =
  runTransactionWithRetry TxSessions.ReadCommitted TxSessions.Read $
    UserGroupPage
      <$> Tx.statement () (refineResult (mapM toUserGroup) $ buildStatement query rows)
      <*> getUserGroupCount tid req
  where
    rows :: HD.Result [(UUID, Text, Int32, UTCTime, Vector UUID, Int32)]
    rows =
      HD.rowList $
        (,,,,,)
          <$> HD.column (HD.nonNullable HD.uuid)
          <*> HD.column (HD.nonNullable HD.text)
          <*> HD.column (HD.nonNullable HD.int4)
          <*> HD.column (HD.nonNullable HD.timestamptz)
          <*> decodeUuidVector
          <*> HD.column (HD.nonNullable HD.int4)

    query :: QueryFragment
    query =
      mconcat $
        map
          literal
          [ "select",
            T.intercalate
              ", "
              [ "ug.id :: uuid",
                "ug.name :: text",
                "ug.managed_by :: int",
                "ug.created_at :: timestamptz",
                "coalesce(array_agg(gm.user_id) filter (WHERE gm.user_id IS NOT NULL), array[]::uuid[]) :: uuid[]",
                "count(gm.user_id) :: int"
              ],
            "from user_group ug",
            "left join user_group_member gm on ug.id = gm.user_group_id"
          ]
          <> [where_ (groupMatchIdName tid req <> groupPaginationWhereClause req)]
          <> [ literal "group by ug.team_id, ug.id"
             ]
          <> groupPaginationOrderBy req

    toUserGroup :: (UUID, Text, Int32, UTCTime, Vector UUID, Int32) -> Either Text UserGroup
    toUserGroup (Id -> id_, name', managedBy', createdAt', members', Just . fromIntegral -> membersCount) = do
      name <- userGroupNameFromText name'
      managedBy <- parseManagedBy managedBy'
      let createdAt = toUTCTimeMillis createdAt'
          channels = Nothing
          channelsCount = Nothing
          members = Identity (fmap Id members' :: Vector UserId)
      pure $ UserGroup_ {..}

groupMatchIdName :: TeamId -> UserGroupPageRequest -> [QueryFragment]
groupMatchIdName tid req =
  clause1 "ug.team_id" "=" tid
    : managedByClause
      <> nameClause
  where
    nameClause = case req.searchString of
      Just name -> [like "ug.name" name]
      Nothing -> []
    managedByClause = case req.managedByFilter of
      Just managedBy -> [clause1 "ug.managed_by" "=" (managedByToInt32 managedBy)]
      Nothing -> []

groupPaginationWhereClause :: UserGroupPageRequest -> [QueryFragment]
groupPaginationWhereClause req = case paginationClause req.paginationState of
  Just c -> [clause (sortOrderOperator req.sortOrder) c]
  Nothing -> []

groupPaginationOrderBy :: UserGroupPageRequest -> [QueryFragment]
groupPaginationOrderBy req =
  [ orderBy $
      case req.paginationState of
        PaginationSortByName _ -> [("ug.name", req.sortOrder), ("ug.id", req.sortOrder)]
        PaginationSortByCreatedAt _ -> [("ug.created_at", req.sortOrder), ("ug.id", req.sortOrder)]
        _ -> [("ug.id", req.sortOrder)],
    limit $
      fromIntegral @_ @Int32 (pageSizeToWord req.pageSize)
  ]
    <> case req.paginationState of
      PaginationOffset n -> [offset (fromIntegral n :: Int32)]
      _ -> []

getUserGroupCount :: TeamId -> UserGroupPageRequest -> Tx.Transaction Int
getUserGroupCount tid req = Tx.statement () $ refineResult parseCount $ buildStatement query decoder
  where
    query = literal "select count(*) from user_group ug" <> where_ (groupMatchIdName tid req)
    decoder = HD.singleRow (HD.column (HD.nonNullable HD.int8))

decodeUuidVector :: HD.Row (Vector UUID)
decodeUuidVector =
  HD.column $
    HD.nonNullable $
      HD.array $
        HD.dimension V.replicateM $
          HD.element $
            HD.nonNullable HD.uuid

parseManagedBy :: Int32 -> Either Text ManagedBy
parseManagedBy = \case
  0 -> pure ManagedByWire
  1 -> pure ManagedByScim
  bad -> Left $ "Could not parse managedBy value: " <> T.pack (show bad)

getUserGroups ::
  forall r.
  ( UserGroupStorePostgresEffectConstraints r,
    Member (Input (Local ())) r
  ) =>
  TeamId ->
  UserGroupPageRequest ->
  Sem r UserGroupPage
getUserGroups tid req@(UserGroupPageRequest {..}) = do
  loc <- inputQualifyLocal ()
  runTransactionWithRetry TxSessions.ReadCommitted TxSessions.Read $
    UserGroupPage <$> getUserGroupsSession loc <*> getUserGroupCount tid req
  where
    getUserGroupsSession :: Local () -> Tx.Transaction [UserGroupMeta]
    getUserGroupsSession loc =
      Tx.statement () $
        refineResult (mapM (parseRow loc)) $
          buildStatement
            ( mconcat $
                [ literal "select",
                  literal selectors,
                  literal "from user_group as ug",
                  where_ (groupMatchIdName tid req <> groupPaginationWhereClause req)
                ]
                  <> groupPaginationOrderBy req
            )
            decodeRow

    decodeRow :: HD.Result [(UUID, Text, Int32, UTCTime, Maybe Int64, Int64, Maybe (Vector UUID))]
    decodeRow =
      HD.rowList
        ( (,,,,,,)
            <$> HD.column (HD.nonNullable HD.uuid)
            <*> HD.column (HD.nonNullable HD.text)
            <*> HD.column (HD.nonNullable HD.int4)
            <*> HD.column (HD.nonNullable HD.timestamptz)
            <*> (if req.includeMemberCount then Just <$> HD.column (HD.nonNullable HD.int8) else pure Nothing)
            <*> HD.column (HD.nonNullable HD.int8)
            <*> ( if req.includeChannels
                    then Just <$> decodeUuidVector
                    else pure Nothing
                )
        )

    parseRow :: Local a -> (UUID, Text, Int32, UTCTime, Maybe Int64, Int64, Maybe (Vector UUID)) -> Either Text UserGroupMeta
    parseRow loc (Id -> id_, namePre, managedByPre, toUTCTimeMillis -> createdAt, membersCountRaw, channelsCountRaw, maybeChannels) = do
      managedBy <- parseManagedBy managedByPre
      name <- userGroupNameFromText namePre
      let members = Const ()
          membersCount = fromIntegral <$> membersCountRaw
          channelsCount = Just (fromIntegral channelsCountRaw)
          channels = fmap (fmap (tUntagged . qualifyAs loc . Id)) maybeChannels
      pure $ UserGroup_ {..}

    selectors :: Text
    selectors =
      T.intercalate ", " $
        ["id", "name", "managed_by", "created_at"]
          <> ["(select count(*) from user_group_member as ugm where ugm.user_group_id = ug.id) as members" | includeMemberCount]
          <> ["(select count(*) from user_group_channel as ugc where ugc.user_group_id = ug.id) as channels"]
          <> ["coalesce((select array_agg(ugc.conv_id) from user_group_channel as ugc where ugc.user_group_id = ug.id), array[]::uuid[]) as channel_ids" | includeChannels]

createUserGroup ::
  forall r.
  (UserGroupStorePostgresEffectConstraints r) =>
  TeamId ->
  NewUserGroup ->
  ManagedBy ->
  Sem r UserGroup
createUserGroup team newUserGroup managedBy = do
  runTransaction TxSessions.Serializable TxSessions.Write do
    (id_, name, managedBy_, createdAt) <- Tx.statement (newUserGroup.name, team, managedBy) insertGroupStatement
    Tx.statement (toUUID id_, newUserGroup.members) insertGroupMembersStatement
    pure
      UserGroup_
        { membersCount = Nothing,
          members = Identity newUserGroup.members,
          channels = mempty,
          managedBy = managedBy_,
          channelsCount = Nothing,
          id_,
          name,
          createdAt
        }
  where
    decodeMetadataRow :: (UUID, Text, Int32, UTCTime) -> Either Text (UserGroupId, UserGroupName, ManagedBy, UTCTimeMillis)
    decodeMetadataRow (groupId, name, managedByInt, utcTime) =
      (Id groupId,,,toUTCTimeMillis utcTime)
        <$> userGroupNameFromText name
        <*> managedByFromInt32 managedByInt

    insertGroupStatement :: Statement (UserGroupName, TeamId, ManagedBy) (UserGroupId, UserGroupName, ManagedBy, UTCTimeMillis)
    insertGroupStatement =
      lmap (\(n, t, m) -> (userGroupNameToText n, t.toUUID, managedByToInt32 m))
        . refineResult decodeMetadataRow
        $ [singletonStatement|
            insert into user_group (name, team_id, managed_by)
              values ($1 :: text, $2 :: uuid, $3 :: int)
              returning id :: uuid, name :: text, managed_by :: int, created_at :: timestamptz
            |]

-- This can perhaps be simplified using rel8
insertGroupMembersStatement :: Statement (UUID, Vector UserId) ()
insertGroupMembersStatement =
  lmap (second (fmap (.toUUID)) . uncurry toRelationTable) $
    [resultlessStatement|
      insert into user_group_member (user_group_id, user_id) select * from unnest ($1 :: uuid[], $2 :: uuid[])
      |]

updateGroup ::
  forall r.
  (UserGroupStorePostgresEffectConstraints r) =>
  TeamId ->
  UserGroupId ->
  UserGroupUpdate ->
  Sem r (Maybe ())
updateGroup tid gid gup = do
  found <- isJust <$> runStatement (tid, gid, gup.name) updateGroupStatement
  pure $ if found then Just () else Nothing
  where
    updateGroupStatement :: Statement (TeamId, UserGroupId, UserGroupName) (Maybe Bool)
    updateGroupStatement =
      lmap (\(t, g, n) -> (t.toUUID, g.toUUID, userGroupNameToText n)) $
        [maybeStatement|
          update user_group set name = ($3 :: text)
            where team_id = ($1 :: uuid) and id = ($2 :: uuid)
            returning (true :: bool)
          |]

deleteGroup ::
  forall r.
  (UserGroupStorePostgresEffectConstraints r) =>
  TeamId ->
  UserGroupId ->
  Sem r (Maybe ())
deleteGroup tid gid = do
  found <- isJust <$> runStatement (tid, gid) deleteGroupStatement
  pure $ if found then Just () else Nothing
  where
    deleteGroupStatement :: Statement (TeamId, UserGroupId) (Maybe Bool)
    deleteGroupStatement =
      lmap (\(t, g) -> (t.toUUID, g.toUUID)) $
        [maybeStatement|
          delete from user_group
            where team_id = ($1 :: uuid) and id = ($2 :: uuid)
            returning (true :: bool)
          |]

addUser ::
  forall r.
  (UserGroupStorePostgresEffectConstraints r) =>
  UserGroupId ->
  UserId ->
  Sem r ()
addUser ugid uid =
  runStatement (ugid, uid) insertStatement
  where
    insertStatement :: Statement (UserGroupId, UserId) ()
    insertStatement =
      lmapPG
        [resultlessStatement|
        insert into user_group_member (user_group_id, user_id) values (($1 :: uuid), ($2 :: uuid))
        |]

removeUser ::
  forall r.
  (UserGroupStorePostgresEffectConstraints r) =>
  UserGroupId ->
  UserId ->
  Sem r ()
removeUser ugid uid =
  runStatement (ugid, uid) deleteStatement
  where
    deleteStatement :: Statement (UserGroupId, UserId) ()
    deleteStatement =
      lmapPG
        [resultlessStatement|
        delete from user_group_member where user_group_id = ($1 :: uuid) and user_id = ($2 :: uuid)
        |]

updateUserGroupChannels ::
  forall r.
  (UserGroupStorePostgresEffectConstraints r) =>
  Bool ->
  UserGroupId ->
  Vector ConvId ->
  Sem r ()
updateUserGroupChannels appendOnly gid convIds = do
  runTransaction TxSessions.Serializable TxSessions.Write $ do
    unless appendOnly $
      Tx.statement (gid, convIds) deleteStatement
    Tx.statement (gid, convIds) insertStatement
  where
    deleteStatement :: Statement (UserGroupId, Vector ConvId) ()
    deleteStatement =
      lmap
        (bimap toUUID (fmap toUUID))
        $ [resultlessStatement|
          delete from user_group_channel where user_group_id = ($1 :: uuid) and conv_id not in (SELECT unnest($2 :: uuid[]))
          |]

    insertStatement :: Statement (UserGroupId, Vector ConvId) ()
    insertStatement =
      lmap (bimap (fmap (.toUUID)) (fmap (.toUUID)) . uncurry toRelationTable) $
        [resultlessStatement|
          insert into user_group_channel (user_group_id, conv_id)  select * from unnest ($1 :: uuid[], $2 :: uuid[])
          on conflict (user_group_id, conv_id) do nothing
          |]

getUserGroupChannels ::
  forall r.
  (UserGroupStorePostgresEffectConstraints r) =>
  TeamId ->
  UserGroupId ->
  Sem r (Maybe (Vector ConvId))
getUserGroupChannels tid gid = do
  mbUuids <- runStatement (gid, tid) getChannelsStatement
  pure (fmap (fmap Id) mbUuids)
  where
    getChannelsStatement :: Statement (UserGroupId, TeamId) (Maybe (Vector UUID))
    getChannelsStatement =
      lmap (\(g, t) -> (g.toUUID, t.toUUID)) $
        [maybeStatement|
          select
            coalesce(
              (select array_agg(ugc.conv_id) from user_group_channel ugc where ugc.user_group_id = ug.id),
              array[]::uuid[]
            ) :: uuid[]
          from user_group ug
          where ug.id = ($1 :: uuid) and ug.team_id = ($2 :: uuid)
        |]

toRelationTable :: a -> Vector b -> (Vector a, Vector b)
toRelationTable a bs = (a <$ bs, bs)
