{-# 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.IndexedUserStore.ElasticSearch where

import Control.Error (lastMay)
import Control.Exception (throwIO)
import Data.Aeson
import Data.Aeson.Key qualified as Key
import Data.Aeson.Types (parseMaybe)
import Data.ByteString qualified as LBS
import Data.ByteString.Builder
import Data.ByteString.Conversion
import Data.Id
import Data.List.NonEmpty (NonEmpty (..))
import Data.Map.Strict qualified as M
import Data.Text qualified as Text
import Data.Text.Ascii
import Data.Text.Encoding qualified as Text
import Database.Bloodhound qualified as ES
import Imports
import Network.HTTP.Client
import Network.HTTP.Types
import Numeric.Natural (Natural)
import Polysemy
import Wire.API.Team.Role (roleName)
import Wire.API.Team.Size (TeamSize (TeamSize))
import Wire.API.User.Search
import Wire.IndexedUserStore
import Wire.Sem.Metrics (Metrics)
import Wire.Sem.Metrics qualified as Metrics
import Wire.UserSearch.Metrics
import Wire.UserSearch.Types
import Wire.UserStore.IndexUser

data ESConn = ESConn
  { env :: ES.BHEnv,
    indexName :: ES.IndexName
  }

data IndexedUserStoreConfig = IndexedUserStoreConfig
  { conn :: ESConn,
    additionalConn :: Maybe ESConn
  }

interpretIndexedUserStoreES ::
  ( Member (Embed IO) r,
    Member Metrics r
  ) =>
  IndexedUserStoreConfig ->
  InterpreterFor IndexedUserStore r
interpretIndexedUserStoreES cfg =
  interpret $ \case
    Upsert docId userDoc versioning -> upsertImpl cfg docId userDoc versioning
    UpdateTeamSearchVisibilityInbound tid vis ->
      updateTeamSearchVisibilityInboundImpl cfg tid vis
    BulkUpsert docs -> bulkUpsertImpl cfg docs
    DoesIndexExist -> doesIndexExistImpl cfg
    SearchUsers searcherId mSearcherTeam teamSearchInfo term maxResults mTypes ->
      searchUsersImpl cfg searcherId mSearcherTeam teamSearchInfo term maxResults mTypes
    PaginateTeamMembers filters maxResults mPagingState ->
      paginateTeamMembersImpl cfg filters maxResults mPagingState
    GetTeamSize tid -> getTeamSizeImpl cfg tid

getTeamSizeImpl ::
  (Member (Embed IO) r) =>
  IndexedUserStoreConfig ->
  TeamId ->
  Sem r TeamSize
getTeamSizeImpl cfg tid = do
  r <- embed $ ES.runBH cfg.conn.env $ do
    res <- ES.searchByType cfg.conn.indexName mappingName search
    liftIO $ ES.parseEsResponse res
  result <- either (embed . throwIO . IndexLookupError) pure (r :: Either ES.EsError (ES.SearchResult UserDoc))
  let aggs = fromMaybe mempty (ES.aggregations result)
      getCount name = maybe 0 (.filterDocCount) $ M.lookup name aggs >>= parseMaybe (parseJSON @FilterResult)
  pure $ TeamSize (getCount "regulars") (getCount "apps")
  where
    teamQ = termQ "team" (idToText tid)

    -- Regular users: type = "regular" or type field absent (legacy documents)
    regularQuery =
      ES.QueryBoolQuery
        boolQuery
          { ES.boolQueryMustMatch =
              [ teamQ,
                ES.QueryBoolQuery
                  boolQuery
                    { ES.boolQueryShouldMatch =
                        [ termQ "type" "regular",
                          ES.QueryBoolQuery
                            boolQuery
                              { ES.boolQueryMustNotMatch = [ES.QueryExistsQuery (ES.FieldName "type")]
                              }
                        ]
                    }
              ]
          }

    appQuery =
      ES.QueryBoolQuery
        boolQuery
          { ES.boolQueryMustMatch = [teamQ, termQ "type" "app"]
          }

    search =
      (ES.mkSearch Nothing Nothing)
        { ES.size = ES.Size 0,
          ES.aggBody =
            Just $
              ES.mkAggregations "regulars" (ES.FilterAgg (ES.FilterAggregation (ES.Filter regularQuery) Nothing))
                <> ES.mkAggregations "apps" (ES.FilterAgg (ES.FilterAggregation (ES.Filter appQuery) Nothing))
        }

upsertImpl ::
  forall r.
  ( Member (Embed IO) r,
    Member Metrics r
  ) =>
  IndexedUserStoreConfig ->
  ES.DocId ->
  UserDoc ->
  ES.VersionControl ->
  Sem r ()
upsertImpl cfg docId userDoc versioning = do
  void $ runInBothES cfg indexDoc
  where
    indexDoc :: ES.IndexName -> ES.BH (Sem r) ()
    indexDoc idx = do
      r <- ES.indexDocument idx mappingName settings userDoc docId
      unless (ES.isSuccess r || ES.isVersionConflict r) $ do
        lift $ Metrics.incCounter indexUpdateErrorCounter
        res <- liftIO $ ES.parseEsResponse r
        liftIO . throwIO . IndexUpdateError . either id id $ res
      lift $ Metrics.incCounter indexUpdateSuccessCounter

    settings = ES.defaultIndexDocumentSettings {ES.idsVersionControl = versioning}

updateTeamSearchVisibilityInboundImpl :: forall r. (Member (Embed IO) r) => IndexedUserStoreConfig -> TeamId -> SearchVisibilityInbound -> Sem r ()
updateTeamSearchVisibilityInboundImpl cfg tid vis =
  void $ runInBothES cfg updateAllDocs
  where
    updateAllDocs :: ES.IndexName -> ES.BH (Sem r) ()
    updateAllDocs idx = do
      r <- ES.updateByQuery idx query (Just script)
      unless (ES.isSuccess r || ES.isVersionConflict r) $ do
        res <- liftIO $ ES.parseEsResponse r
        liftIO . throwIO . IndexUpdateError . either id id $ res

    query :: ES.Query
    query = ES.TermQuery (ES.Term "team" $ idToText tid) Nothing

    script :: ES.Script
    script = ES.Script (Just (ES.ScriptLanguage "painless")) (Just (ES.ScriptInline scriptText)) Nothing Nothing

    -- Unfortunately ES disallows updating ctx._version with a "Update By Query"
    scriptText =
      "ctx._source."
        <> Key.toText searchVisibilityInboundFieldName
        <> " = '"
        <> Text.decodeUtf8 (toByteString' vis)
        <> "';"

bulkUpsertImpl :: (Member (Embed IO) r) => IndexedUserStoreConfig -> [(ES.DocId, UserDoc, ES.VersionControl)] -> Sem r ()
bulkUpsertImpl cfg docs = do
  let bhe = cfg.conn.env
      ES.IndexName idx = cfg.conn.indexName
      ES.MappingName mpp = mappingName
      (ES.Server base) = ES.bhServer bhe
  baseReq <- embed $ parseRequest (Text.unpack $ base <> "/" <> idx <> "/" <> mpp <> "/_bulk")
  let reqWithoutCreds =
        baseReq
          { method = "POST",
            requestHeaders = [(hContentType, "application/x-ndjson")],
            requestBody = RequestBodyLBS (toLazyByteString (foldMap encodeActionAndData docs))
          }
  req <- embed $ bhe.bhRequestHook reqWithoutCreds
  res <- embed $ httpLbs req (ES.bhManager bhe)
  unless (ES.isSuccess res) $ do
    parsedRes <- liftIO $ ES.parseEsResponse res
    liftIO . throwIO . IndexUpdateError . either id id $ parsedRes
  where
    encodeJSONToString :: (ToJSON a) => a -> Builder
    encodeJSONToString = fromEncoding . toEncoding

    encodeActionAndData :: (ES.DocId, UserDoc, ES.VersionControl) -> Builder
    encodeActionAndData (docId, userDoc, versionControl) =
      encodeJSONToString (bulkIndexAction docId versionControl)
        <> "\n"
        <> encodeJSONToString userDoc
        <> "\n"

    bulkIndexAction :: ES.DocId -> ES.VersionControl -> Value
    bulkIndexAction docId versionControl =
      let (versionType :: Maybe Text, version) = case versionControl of
            ES.NoVersionControl -> (Nothing, Nothing)
            ES.InternalVersion v -> (Nothing, Just v)
            ES.ExternalGT (ES.ExternalDocVersion v) -> (Just "external", Just v)
            ES.ExternalGTE (ES.ExternalDocVersion v) -> (Just "external_gte", Just v)
            ES.ForceVersion (ES.ExternalDocVersion v) -> (Just "force", Just v)
       in object
            [ "index"
                .= object
                  [ "_id" .= docId,
                    "version_type" .= versionType,
                    "version" .= version
                  ]
            ]

doesIndexExistImpl :: (Member (Embed IO) r) => IndexedUserStoreConfig -> Sem r Bool
doesIndexExistImpl cfg = do
  (mainExists, fromMaybe True -> additionalExists) <- runInBothES cfg ES.indexExists
  pure $ mainExists && additionalExists

searchUsersImpl ::
  (Member (Embed IO) r) =>
  IndexedUserStoreConfig ->
  UserId ->
  Maybe TeamId ->
  TeamSearchInfo ->
  Text ->
  Int ->
  Maybe [UserTypeFilter] ->
  Sem r (SearchResult UserDoc)
searchUsersImpl cfg searcherId mSearcherTeam teamSearchInfo term maxResults mTypes = do
  queryIndex cfg maxResults $
    defaultUserQuery searcherId mSearcherTeam teamSearchInfo mTypes term

-- | The default or canonical 'IndexQuery'.
--
-- The intention behind parameterising 'queryIndex' over the 'IndexQuery' is that
-- it allows to experiment with different queries (perhaps in an A/B context).
--
-- FUTUREWORK: Drop legacyPrefixMatch
defaultUserQuery :: UserId -> Maybe TeamId -> TeamSearchInfo -> Maybe [UserTypeFilter] -> Text -> IndexQuery Contact
defaultUserQuery searcher mSearcherTeamId teamSearchInfo mTypes (normalized -> term') =
  let matchPhraseOrPrefix =
        ES.QueryMultiMatchQuery $
          ( ES.mkMultiMatchQuery
              [ ES.FieldName "handle.prefix^2",
                ES.FieldName "normalized.prefix",
                ES.FieldName "normalized^3"
              ]
              (ES.QueryString term')
          )
            { ES.multiMatchQueryType = Just ES.MultiMatchMostFields,
              ES.multiMatchQueryOperator = ES.And
            }
      query =
        ES.QueryBoolQuery
          boolQuery
            { ES.boolQueryMustMatch =
                [ ES.QueryBoolQuery
                    boolQuery
                      { ES.boolQueryShouldMatch = [matchPhraseOrPrefix],
                        -- This removes exact handle matches, as they are fetched from cassandra
                        ES.boolQueryMustNotMatch = [termQ "handle" term']
                      }
                ],
              ES.boolQueryShouldMatch = [ES.QueryExistsQuery (ES.FieldName "handle")]
            }
      -- This reduces relevance on users not in team of search by 90% (no
      -- science behind that number). If the searcher is not part of a team the
      -- relevance is not reduced for any users.
      queryWithBoost =
        ES.QueryBoostingQuery
          ES.BoostingQuery
            { ES.positiveQuery = query,
              ES.negativeQuery = maybe ES.QueryMatchNoneQuery matchUsersNotInTeam mSearcherTeamId,
              ES.negativeBoost = ES.Boost 0.1
            }
   in mkUserQuery searcher mSearcherTeamId teamSearchInfo mTypes queryWithBoost

paginateTeamMembersImpl ::
  (Member (Embed IO) r) =>
  IndexedUserStoreConfig ->
  BrowseTeamFilters ->
  Int ->
  Maybe PagingState ->
  Sem r (SearchResult UserDoc)
paginateTeamMembersImpl cfg BrowseTeamFilters {..} maxResults mPagingState = do
  let (IndexQuery q f sortSpecs) =
        teamUserSearchQuery teamId mQuery mRoleFilter mSortBy mSortOrder mEmailVerificationFilter mSearchable
  let search =
        (ES.mkSearch (Just q) (Just f))
          { -- we are requesting one more result than the page size to determine if there is a next page
            ES.size = ES.Size (fromIntegral maxResults + 1),
            ES.sortBody = Just (fmap ES.DefaultSortSpec sortSpecs),
            ES.searchAfterKey = toSearchAfterKey =<< mPagingState
          }
  mkResult <$> searchInMainIndex cfg search
  where
    toSearchAfterKey ps = decode' . LBS.fromStrict =<< (decodeBase64Url . unPagingState) ps

    fromSearchAfterKey :: ES.SearchAfterKey -> PagingState
    fromSearchAfterKey = PagingState . encodeBase64Url . LBS.toStrict . encode

    mkResult es =
      let hitsPlusOne = ES.hits . ES.searchHits $ es
          hits = take (fromIntegral maxResults) hitsPlusOne
          mps = fromSearchAfterKey <$> lastMay (mapMaybe ES.hitSort hits)
          results = mapMaybe ES.hitSource hits
       in SearchResult
            { searchFound = ES.hitsTotalValue . ES.hitsTotal . ES.searchHits $ es,
              searchReturned = length results,
              searchTook = ES.took es,
              searchResults = results,
              searchPolicy = FullSearch,
              searchPagingState = mps,
              searchHasMore = Just $ length hitsPlusOne > length hits
            }

searchInMainIndex :: forall r. (Member (Embed IO) r) => IndexedUserStoreConfig -> ES.Search -> Sem r (ES.SearchResult UserDoc)
searchInMainIndex cfg search = do
  r <- ES.runBH cfg.conn.env $ do
    res <- ES.searchByType cfg.conn.indexName mappingName search
    liftIO $ ES.parseEsResponse res
  either (embed . throwIO . IndexLookupError) pure r

queryIndex ::
  (Member (Embed IO) r) =>
  IndexedUserStoreConfig ->
  Int ->
  IndexQuery x ->
  Sem r (SearchResult UserDoc)
queryIndex cfg s (IndexQuery q f _) = do
  let search = (ES.mkSearch (Just q) (Just f)) {ES.size = ES.Size (fromIntegral s)}
  mkResult <$> searchInMainIndex cfg search
  where
    mkResult es =
      let results = mapMaybe ES.hitSource . ES.hits . ES.searchHits $ es
       in SearchResult
            { searchFound = ES.hitsTotalValue . ES.hitsTotal . ES.searchHits $ es,
              searchReturned = length results,
              searchTook = ES.took es,
              searchResults = results,
              searchPolicy = FullSearch,
              searchPagingState = Nothing,
              searchHasMore = Nothing
            }

teamUserSearchQuery ::
  TeamId ->
  Maybe Text ->
  Maybe RoleFilter ->
  Maybe TeamUserSearchSortBy ->
  Maybe TeamUserSearchSortOrder ->
  Maybe EmailVerificationFilter ->
  Maybe Bool ->
  IndexQuery TeamContact
teamUserSearchQuery tid mbSearchText mRoleFilter mSortBy mSortOrder mEmailFilter mSearchable =
  IndexQuery
    ( maybe
        (ES.MatchAllQuery Nothing)
        matchPhraseOrPrefix
        mbQStr
    )
    teamFilter
    -- in combination with pagination a non-unique search specification can lead to missing results
    -- therefore we use the unique `_doc` value as a tie breaker
    -- - see https://www.elastic.co/guide/en/elasticsearch/reference/6.8/search-request-sort.html for details on `_doc`
    -- - see https://www.elastic.co/guide/en/elasticsearch/reference/6.8/search-request-search-after.html for details on pagination and tie breaker
    -- in the latter article it "is advised to duplicate (client side or [...]) the content of the _id field
    -- in another field that has doc value enabled and to use this new field as the tiebreaker for the sort"
    -- so alternatively we could use the user ID as a tie breaker, but this would require a change in the index mapping
    (sorting ++ sortingTieBreaker)
  where
    sorting :: [ES.DefaultSort]
    sorting =
      maybe
        [defaultSort SortByCreatedAt SortOrderDesc | isNothing mbQStr]
        (\tuSortBy -> [defaultSort tuSortBy (fromMaybe SortOrderAsc mSortOrder)])
        mSortBy
    sortingTieBreaker :: [ES.DefaultSort]
    sortingTieBreaker = [ES.DefaultSort (ES.FieldName "_doc") ES.Ascending Nothing Nothing Nothing Nothing]

    mbQStr :: Maybe Text
    mbQStr =
      case mbSearchText of
        Nothing -> Nothing
        Just q ->
          case normalized q of
            "" -> Nothing
            term' -> Just term'

    matchPhraseOrPrefix term' =
      ES.QueryMultiMatchQuery $
        ( ES.mkMultiMatchQuery
            [ ES.FieldName "email^4",
              ES.FieldName "handle^4",
              ES.FieldName "normalized^3",
              ES.FieldName "email.prefix^3",
              ES.FieldName "handle.prefix^2",
              ES.FieldName "normalized.prefix"
            ]
            (ES.QueryString term')
        )
          { ES.multiMatchQueryType = Just ES.MultiMatchMostFields,
            ES.multiMatchQueryOperator = ES.And
          }

    teamFilter :: ES.Filter
    teamFilter = ES.Filter $ ES.QueryBoolQuery boolQuery {ES.boolQueryMustMatch = mustMatch}
      where
        mustMatch :: [ES.Query]
        mustMatch = ES.TermQuery (ES.Term "team" $ idToText tid) Nothing : roleFilter <> emailFilter <> searchableFilter

        searchableFilter :: [ES.Query]
        searchableFilter = case mSearchable of
          Just False -> [ES.TermQuery (ES.Term "searchable" "false") Nothing]
          Just True -> [ES.QueryBoolQuery boolQuery {ES.boolQueryMustNotMatch = [ES.TermQuery (ES.Term "searchable" "false") Nothing]}]
          Nothing -> []

        roleFilter :: [ES.Query]
        roleFilter =
          case mRoleFilter of
            Nothing -> []
            Just (RoleFilter []) -> []
            Just (RoleFilter (r : rs)) -> [ES.TermsQuery "role" (roleName <$> r :| rs)]

        emailFilter :: [ES.Query]
        emailFilter =
          case mEmailFilter of
            Nothing -> []
            -- Verified: must have a verified email and must NOT have an unverified email
            Just EmailVerified ->
              [ ES.QueryBoolQuery
                  boolQuery
                    { ES.boolQueryMustMatch = [ES.QueryExistsQuery (ES.FieldName "email")],
                      ES.boolQueryMustNotMatch = [ES.QueryExistsQuery (ES.FieldName "email_unvalidated")]
                    }
              ]
            -- Unverified: must have an unverified email, regardless of verified
            Just EmailUnverified -> [ES.QueryExistsQuery (ES.FieldName "email_unvalidated")]

    defaultSort :: TeamUserSearchSortBy -> TeamUserSearchSortOrder -> ES.DefaultSort
    defaultSort tuSortBy sortOrder =
      ES.DefaultSort
        ( case tuSortBy of
            SortByName -> ES.FieldName "name"
            SortByHandle -> ES.FieldName "handle.keyword"
            SortByEmail -> ES.FieldName "email.keyword"
            SortBySAMLIdp -> ES.FieldName "saml_idp"
            SortByManagedBy -> ES.FieldName "managed_by"
            SortByRole -> ES.FieldName "role"
            SortByCreatedAt -> ES.FieldName "created_at"
        )
        ( case sortOrder of
            SortOrderAsc -> ES.Ascending
            SortOrderDesc -> ES.Descending
        )
        Nothing
        Nothing
        Nothing
        Nothing

mkUserQuery :: UserId -> Maybe TeamId -> TeamSearchInfo -> Maybe [UserTypeFilter] -> ES.Query -> IndexQuery Contact
mkUserQuery searcher mSearcherTeamId teamSearchInfo mTypes q =
  IndexQuery
    q
    ( ES.Filter
        . ES.QueryBoolQuery
        $ boolQuery
          { ES.boolQueryMustNotMatch =
              maybeToList (matchSelf searcher)
                <>
                -- The following matches both where searchable is true
                -- or where the field is missing. There didn't seem to
                -- be a more readable way to express
                -- "not(exists(searchable)) or searchable = true" in
                -- Elastic Search.
                [ES.TermQuery (ES.Term "searchable" "false") Nothing]
                <>
                -- Exclude apps from other teams
                maybeToList (matchAppsFromOtherTeams mSearcherTeamId),
            ES.boolQueryMustMatch =
              [ restrictSearchSpaceByTeam mSearcherTeamId teamSearchInfo,
                restrictSearchSpaceByUserType mTypes,
                ES.QueryBoolQuery
                  boolQuery
                    { ES.boolQueryShouldMatch =
                        [ termQ "account_status" "active",
                          -- Also match entries where the account_status field is not present.
                          -- These must have been inserted before we added the account_status
                          -- and at that time we only inserted active users in the first place.
                          -- This should be unnecessary after re-indexing, but let's be lenient
                          -- here for a while.
                          ES.QueryBoolQuery
                            boolQuery
                              { ES.boolQueryMustNotMatch =
                                  [ES.QueryExistsQuery (ES.FieldName "account_status")]
                              }
                        ]
                    }
              ]
          }
    )
    []

termQ :: Text -> Text -> ES.Query
termQ f v =
  ES.TermQuery
    ES.Term
      { ES.termField = f,
        ES.termValue = v
      }
    Nothing

matchSelf :: UserId -> Maybe ES.Query
matchSelf searcher = Just (termQ "_id" (idToText searcher))

-- | Exclude apps from other teams.
-- Apps should only be searchable within their own team.
matchAppsFromOtherTeams :: Maybe TeamId -> Maybe ES.Query
matchAppsFromOtherTeams mSearcherTeamId =
  Just $
    ES.QueryBoolQuery
      boolQuery
        { ES.boolQueryMustMatch =
            [ -- Match apps (type = "app")
              termQ "type" "app",
              -- That are from a different team than the searcher
              case mSearcherTeamId of
                -- If searcher has no team, exclude all apps
                Nothing ->
                  ES.QueryExistsQuery (ES.FieldName "team")
                -- If searcher has a team, exclude apps from other teams or with no team
                Just searcherTeam ->
                  ES.QueryBoolQuery
                    boolQuery
                      { ES.boolQueryShouldMatch =
                          [ -- Apps with no team
                            ES.QueryBoolQuery
                              boolQuery
                                { ES.boolQueryMustNotMatch =
                                    [ES.QueryExistsQuery (ES.FieldName "team")]
                                },
                            -- Apps from a different team
                            ES.QueryBoolQuery
                              boolQuery
                                { ES.boolQueryMustMatch =
                                    [ES.QueryExistsQuery (ES.FieldName "team")],
                                  ES.boolQueryMustNotMatch =
                                    [termQ "team" (idToText searcherTeam)]
                                }
                          ]
                      }
            ]
        }

-- | See 'TeamSearchInfo'
restrictSearchSpaceByTeam :: Maybe TeamId -> TeamSearchInfo -> ES.Query
--
-- FUTUREWORK(fisx): can this commented-out code be removed?  if not,
-- there should be a comment explaining why.
--
-- restrictSearchSpace (FederatedSearch Nothing) =
--   ES.QueryBoolQuery
--     boolQuery
--       { ES.boolQueryShouldMatch =
--           [ matchNonTeamMemberUsers,
--             matchTeamMembersSearchableByAllTeams
--           ]
--       }
-- restrictSearchSpace (FederatedSearch (Just [])) =
--   ES.QueryBoolQuery
--     boolQuery
--       { ES.boolQueryMustMatch =
--           [ -- if the list of allowed teams is empty, this is impossible to fulfill, and no results will be returned
--             -- this case should be handled earlier, so this is just a safety net
--             ES.TermQuery (ES.Term "team" "must not match any team") Nothing
--           ]
--       }
-- restrictSearchSpace (FederatedSearch (Just teams)) =
--   ES.QueryBoolQuery
--     boolQuery
--       { ES.boolQueryMustMatch =
--           [ matchTeamMembersSearchableByAllTeams,
--             onlyInTeams
--           ]
--       }
--   where
--     onlyInTeams = ES.QueryBoolQuery boolQuery {ES.boolQueryShouldMatch = map matchTeamMembersOf teams}
restrictSearchSpaceByTeam mteam searchInfo =
  case (mteam, searchInfo) of
    (Nothing, _) -> matchNonTeamMemberUsers
    (Just _, NoTeam) -> matchNonTeamMemberUsers
    (Just searcherTeam, TeamOnly team) ->
      if searcherTeam == team
        then matchTeamMembersOf team
        else ES.QueryMatchNoneQuery
    (Just searcherTeam, AllUsers) ->
      ES.QueryBoolQuery
        boolQuery
          { ES.boolQueryShouldMatch =
              [ matchNonTeamMemberUsers,
                matchTeamMembersSearchableByAllTeams,
                matchTeamMembersOf searcherTeam
              ]
          }

restrictSearchSpaceByUserType :: Maybe [UserTypeFilter] -> ES.Query
restrictSearchSpaceByUserType = \case
  -- Nothing (param omitted) and Just [] (param present but empty) both mean
  -- "no filter" to avoid surprising empty-result regressions when clients send
  -- `type=` without values.
  Nothing -> ES.MatchAllQuery Nothing
  Just [] -> ES.MatchAllQuery Nothing
  Just uts@(utsH : utsT) ->
    if UserTypeFilterRegular `elem` uts
      then
        ES.QueryBoolQuery
          boolQuery
            { ES.boolQueryShouldMatch =
                [ ES.TermsQuery "type" (userTypeFilterToText <$> (utsH :| utsT)),
                  -- Older index entries may lack the "type" field; treat those as regular.
                  ES.QueryBoolQuery
                    boolQuery
                      { ES.boolQueryMustNotMatch =
                          [ES.QueryExistsQuery (ES.FieldName "type")]
                      }
                ]
            }
      else ES.TermsQuery "type" (userTypeFilterToText <$> (utsH :| utsT))

matchTeamMembersOf :: TeamId -> ES.Query
matchTeamMembersOf team = ES.TermQuery (ES.Term "team" $ idToText team) Nothing

matchTeamMembersSearchableByAllTeams :: ES.Query
matchTeamMembersSearchableByAllTeams =
  ES.QueryBoolQuery
    boolQuery
      { ES.boolQueryMustMatch =
          [ ES.QueryExistsQuery $ ES.FieldName "team",
            ES.TermQuery (ES.Term (Key.toText searchVisibilityInboundFieldName) "searchable-by-all-teams") Nothing
          ]
      }

matchNonTeamMemberUsers :: ES.Query
matchNonTeamMemberUsers =
  ES.QueryBoolQuery
    boolQuery
      { ES.boolQueryMustNotMatch = [ES.QueryExistsQuery $ ES.FieldName "team"]
      }

matchUsersNotInTeam :: TeamId -> ES.Query
matchUsersNotInTeam tid =
  ES.QueryBoolQuery
    boolQuery
      { ES.boolQueryMustNotMatch = [ES.TermQuery (ES.Term "team" $ idToText tid) Nothing]
      }

--------------------------------------------
-- Utils

runInBothES :: (Monad m) => IndexedUserStoreConfig -> (ES.IndexName -> ES.BH m a) -> m (a, Maybe a)
runInBothES cfg f = do
  x <- ES.runBH cfg.conn.env $ f cfg.conn.indexName
  y <- forM cfg.additionalConn $ \additional ->
    ES.runBH additional.env $ f additional.indexName
  pure (x, y)

mappingName :: ES.MappingName
mappingName = ES.MappingName "user"

boolQuery :: ES.BoolQuery
boolQuery = ES.mkBoolQuery [] [] [] []

-- | (or can something like this be found in bloodhound?)
newtype FilterResult = FilterResult {filterDocCount :: Natural}

instance FromJSON FilterResult where
  parseJSON = withObject "FilterResult" $ \o -> FilterResult <$> o .: "doc_count"
