content: {-# LANGUAGE RecordWildCards #-}
  
  -- 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 hiding (Set)
  import Cassandra qualified as Cas
  import Cassandra.Settings qualified as Cas
  import Conduit
  import Control.Lens
  import Control.Monad.Except
  import Data.Conduit.Combinators qualified as C
  import Data.Handle (Handle)
  import Data.Id
  import Data.Map.Strict qualified as Map
  import Data.Text qualified as T
  import Imports
  import Options
  import Options.Applicative hiding (action)
  import System.Logger qualified as Log
  import Types
  
  -- | The table user_handle grouped by user
  type HandleMap = Map UserId [Handle]
  
  readHandleMap :: Env -> IO HandleMap
  readHandleMap Env {..} =
    runConduit $
      transPipe (runClient envBrig) (paginateC selectUserHandle (paramsP LocalQuorum () envPageSize) x1)
        .| (C.foldM insertAndLog (Map.empty, 0) <&> fst)
    where
      selectUserHandle :: PrepQuery R () (Maybe UserId, Maybe Handle)
      selectUserHandle = "SELECT user, handle FROM user_handle"
  
      insertAndLog :: (HandleMap, Int) -> [(Maybe UserId, Maybe Handle)] -> IO (HandleMap, Int)
      insertAndLog (hmap, nTotal) pairs = do
        let n = length pairs
        Log.info envLogger $ Log.msg @Text $ "handles loaded: " <> (T.pack . show $ nTotal + n)
        pure (foldl' insert hmap pairs, nTotal + n)
  
      insert :: HandleMap -> (Maybe UserId, Maybe Handle) -> HandleMap
      insert hmap (Just uid, maybeToList -> handles) =
        Map.insertWith (++) uid handles hmap
      insert hmap (Nothing, _) = hmap
  
  data Action
    = -- | Set the handle for user (column "handle" in table user).
      -- Precondition:
      -- The handle is the sole handle owned by the user (according to table user_handle)
      SetHandle
        UserId
        Handle
    | -- | Change the handle for user (column "handle" in table user)
      -- and delete the other handle.
      -- Precondition:
      -- Both handles are already owned by the user (according to table user_handle)
      -- and the user doesnt own any more handles.
      -- syntax: ResetHandle uid handle handleToBeRemoved
      ResetHandle
        UserId
        Handle
        Handle
    | NoActionRequired UserId
    deriving (Show)
  
  data ActionError = ActionError UserId Text [Handle]
    deriving (Show)
  
  type ActionResult = Either ActionError Action
  
  decideAction ::
    UserId ->
    -- | "handle" column in table "brig.user"
    Maybe Handle ->
    -- | All handles owned by user in table "brig.user_handle"
    [Handle] ->
    ActionResult
  decideAction uid Nothing [handle] = pure $ SetHandle uid handle
  decideAction uid Nothing [] = pure $ NoActionRequired uid
  decideAction uid Nothing handles = throwError $ ActionError uid "No handle set, but multiple handles owned" handles
  decideAction uid (Just currentHandle) handles =
    case filter (/= currentHandle) handles of
      [] -> pure $ NoActionRequired uid
      [otherHandle] -> pure $ ResetHandle uid otherHandle currentHandle
      handles' -> throwError $ ActionError uid "Handle is set, but multiple handles owned" handles'
  
  sourceActions :: Env -> HandleMap -> ConduitM () ActionResult IO ()
  sourceActions Env {..} hmap =
    transPipe
      (runClient envGalley)
      ( paginateC selectTeam (paramsP LocalQuorum (pure envTeam) envPageSize) x5
          .| C.map (fmap runIdentity)
      )
      .| C.mapM readUsersPage
      .| C.concat
      .| C.map
        ( \(uid, mbHandle) ->
            decideAction uid mbHandle (fromMaybe [] $ Map.lookup uid hmap)
        )
    where
      selectTeam :: PrepQuery R (Identity TeamId) (Identity UserId)
      selectTeam = "SELECT user FROM team_member WHERE team = ?"
  
      readUsersPage :: [UserId] -> IO [(UserId, Maybe Handle)]
      readUsersPage uids =
        runClient envBrig $
          query selectUsers (params LocalQuorum (pure uids))
  
      selectUsers :: PrepQuery R (Identity [UserId]) (UserId, Maybe Handle)
      selectUsers = "SELECT id, handle FROM user WHERE id in ?"
  
  executeAction :: Env -> Action -> IO ()
  executeAction env = \case
    (NoActionRequired _) -> pure ()
    (SetHandle uid handle) ->
      setUserHandle env uid handle
    (ResetHandle uid handle handleRemove) -> do
      setUserHandle env uid handle
      removeHandle env handleRemove
    where
      setUserHandle :: Env -> UserId -> Handle -> IO ()
      setUserHandle Env {..} uid handle =
        runClient envBrig $
          Cas.write updateHandle $
            params LocalQuorum (handle, uid)
        where
          updateHandle :: PrepQuery W (Handle, UserId) ()
          updateHandle = {- `IF EXISTS`, but that requires benchmarking -} "UPDATE user SET handle = ? WHERE id = ?"
  
      removeHandle :: Env -> Handle -> IO ()
      removeHandle Env {..} handle =
        runClient envBrig $
          Cas.write deleteHandle $
            params LocalQuorum (pure handle)
        where
          deleteHandle :: PrepQuery W (Identity Handle) ()
          deleteHandle = "DELETE FROM user_handle WHERE handle = ?"
  
  runCommand :: Env -> IO ()
  runCommand env@Env {..} = do
    Log.info envLogger (Log.msg @Text "Loading all handles...")
    hmap <- readHandleMap env
    Log.info envLogger $ Log.msg @Text ("Processing team" <> (if envSettings ^. setDryRun then " (dry-run) " else " ") <> "...")
    runConduit $
      sourceActions env hmap
        .| C.iterM (handleAction env)
        .| chunkify 100
        .| C.mapM_ (logSummary env)
    Log.info envLogger $ Log.msg @Text "Done."
    where
      handleAction :: Env -> ActionResult -> IO ()
      handleAction _env (Left err) = Log.err envLogger (Log.msg . show $ err)
      handleAction _env (Right action) = do
        Log.debug envLogger $ Log.msg (show $ action)
        unless (envSettings ^. setDryRun) $
          executeAction env action
  
      logSummary :: Env -> [ActionResult] -> IO ()
      logSummary _env results = do
        Log.info envLogger $
          Log.msg $
            ("Action summary for batch" <> (if envSettings ^. setDryRun then " (dry-run) " else " ") <> ":\n")
              <> T.intercalate
                "\n"
                ( let (nErrs, nReset, nSet, nNoOp) = foldl' tally (0, 0, 0, 0) results
                   in catMaybes
                        [ mark "   error(!)" nErrs,
                          mark "   reset original handle" nReset,
                          mark "   set missing handle" nSet,
                          mark "   do nothing" nNoOp
                        ]
                )
        where
          mark :: Text -> Int -> Maybe Text
          mark msg n
            | n > 0 = Just $ msg <> ": " <> T.pack (show n)
            | otherwise = Nothing
  
          tally :: (Int, Int, Int, Int) -> ActionResult -> (Int, Int, Int, Int)
          tally (nErrs, nReset, nSet, nNoOp) (Right ResetHandle {}) = (nErrs, nReset + 1, nSet, nNoOp)
          tally (nErrs, nReset, nSet, nNoOp) (Right SetHandle {}) = (nErrs, nReset, nSet + 1, nNoOp)
          tally (nErrs, nReset, nSet, nNoOp) (Right NoActionRequired {}) = (nErrs, nReset, nSet, nNoOp + 1)
          tally (nErrs, nReset, nSet, nNoOp) (Left _) = (nErrs + 1, nReset, nSet, nNoOp)
  
      chunkify :: (Monad m) => Int -> ConduitT i [i] m ()
      chunkify n = void (C.map (: [])) .| C.chunksOfE n
  
  main :: IO ()
  main = do
    s <- execParser (info (helper <*> settingsParser) desc)
    lgr <- initLogger (if s ^. setDebug then Log.Debug else Log.Info)
  
    brig <- initCas (s ^. setCasBrig) lgr
    galley <- initCas (s ^. setCasGalley) lgr
  
    let env =
          Env
            { envBrig = brig,
              envGalley = galley,
              envPageSize = s ^. setPageSize,
              envTeam = s ^. setTeamId,
              envSettings = s,
              envLogger = lgr
            }
    runCommand env
    where
      desc =
        header "repair-handles"
          <> progDesc "Fix dangling and missing handles"
          <> fullDesc
      initLogger level =
        Log.new
          . Log.setOutput Log.StdOut
          . Log.setFormat Nothing
          . Log.setBufSize 0
          $ Log.setLogLevel level Log.defSettings
      initCas cas l =
        Cas.init
          . Cas.setLogger (Cas.mkLogger l)
          . Cas.setContacts (cas ^. cHosts) []
          . Cas.setPortNumber (fromIntegral $ cas ^. cPort)
          . Cas.setKeyspace (cas ^. cKeyspace)
          . Cas.setProtocolVersion Cas.V4
          $ Cas.defSettings
  