content: {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# OPTIONS_GHC -Wno-orphans -Wno-unused-imports #-}
  
  -- 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 Work where
  
  import Cassandra
  import Data.Conduit
  import Data.Conduit.Internal (zipSources)
  import Data.Conduit.List qualified as C
  import Data.Id
  import Data.Misc
  import Imports
  import System.Logger (Logger)
  import System.Logger qualified as Log
  import UnliftIO.Async (pooledMapConcurrentlyN)
  import Wire.API.Team.Feature
  import Wire.ConversationStore.Cassandra.Instances ()
  
  runCommand :: Logger -> ClientState -> ClientState -> IO ()
  runCommand l spar galley = do
    runConduit $
      zipSources
        (C.sourceList [(1 :: Int32) ..])
        (transPipe (runClient spar) getSsoTeams)
        .| C.mapM
          ( \(i, tids) -> do
              Log.info l (Log.field "number of idps processed: " (show (i * pageSize)))
              pure (runIdentity <$> tids)
          )
        .| C.mapM_ (\tids -> runClient galley (writeSsoFlags tids))
  
  pageSize :: Int32
  pageSize = 1000
  
  getSsoTeams :: ConduitM () [Identity TeamId] Client ()
  getSsoTeams = paginateC cql (paramsP LocalQuorum () pageSize) x5
    where
      cql :: PrepQuery R () (Identity TeamId)
      cql = "select team from idp"
  
  writeSsoFlags :: [TeamId] -> Client ()
  writeSsoFlags = mapM_ (`setSSOTeamConfig` FeatureStatusEnabled)
    where
      setSSOTeamConfig :: (MonadClient m) => TeamId -> FeatureStatus -> m ()
      setSSOTeamConfig tid ssoTeamConfigStatus = do
        retry x5 $ write updateSSOTeamConfig (params LocalQuorum (ssoTeamConfigStatus, tid))
  
      updateSSOTeamConfig :: PrepQuery W (FeatureStatus, TeamId) ()
      updateSSOTeamConfig = {- `IF EXISTS`, but that requires benchmarking -} "update team_features set sso_status = ? where team_id = ?"
  