-- 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 Wire.FederationConfigStore.Cassandra
  ( interpretFederationDomainConfig,
    remotesMapFromCfgFile,
    AddFederationRemoteResult (..),
  )
where

import Cassandra
import Control.Exception (ErrorCall (ErrorCall))
import Control.Lens
import Control.Monad.Catch (throwM)
import Data.Domain
import Data.Id
import Data.Map qualified as Map
import Data.Qualified
import Database.CQL.Protocol (SerialConsistency (LocalSerialConsistency), serialConsistency)
import Imports
import Polysemy
import Polysemy.Embed
import Wire.API.Routes.FederationDomainConfig
import Wire.API.User.Search
import Wire.FederationConfigStore

-- | Interpreter for getting the federation config from the database and the config file.
-- The config file is injected into the interpreter and has precedence over the database.
-- The config file is static and can only be changed by restarting the service.
-- If a domain is configured in the config file, it is not allowed to add it to the database.
-- If a domain is configured in the config file, it is not allowed to update it in the database.
-- If a domain is configured in the config file, it is not allowed to add a team restriction to it in the database.
-- In the future the config file will be removed and the database will be the only source of truth.
interpretFederationDomainConfig ::
  forall r a.
  (Member (Embed IO) r) =>
  ClientState ->
  Maybe FederationStrategy ->
  Map Domain FederationDomainConfig ->
  Sem (FederationConfigStore ': r) a ->
  Sem r a
interpretFederationDomainConfig casClient mFedStrategy fedCfgs =
  interpret $
    runEmbedded (runClient casClient) . embed . \case
      GetFederationConfig d -> getFederationConfig' fedCfgs d
      GetFederationConfigs -> getFederationConfigs' mFedStrategy fedCfgs
      AddFederationConfig cnf -> addFederationConfig' fedCfgs cnf
      UpdateFederationConfig cnf -> updateFederationConfig' fedCfgs cnf
      AddFederationRemoteTeam d t -> addFederationRemoteTeam' fedCfgs d t
      RemoveFederationRemoteTeam d t -> removeFederationRemoteTeam' d t
      GetFederationRemoteTeams d -> getFederationRemoteTeams' d
      BackendFederatesWith mtid -> backendFederatesWithImpl mtid fedCfgs mFedStrategy

-- | Compile config file list into a map indexed by domains.  Use this to make sure the config
-- file is consistent (ie., no two entries for the same domain).
-- This is called during initialization of the interpreter and the service will fail if the config is not consistent.
remotesMapFromCfgFile :: [FederationDomainConfig] -> Map Domain FederationDomainConfig
remotesMapFromCfgFile cfg =
  let dict = [(cnf.domain, cnf) | cnf <- cfg]
      merge c c' =
        if c == c'
          then c
          else error $ "error in config file: conflicting parameters on domain: " <> show (c, c')
   in Map.fromListWith merge dict

getFederationConfigs' ::
  forall m.
  (MonadClient m) =>
  Maybe FederationStrategy ->
  Map Domain FederationDomainConfig ->
  m FederationDomainConfigs
getFederationConfigs' mFedStrategy cfgs = do
  -- FUTUREWORK: we should solely rely on `db` in the future for remote domains; merging
  -- remote domains from `cfg` is just for providing an easier, more robust migration path.
  -- See
  -- https://docs.wire.com/understand/federation/backend-communication.html#configuring-remote-connections,
  -- http://docs.wire.com/developer/developer/federation-design-aspects.html#configuring-remote-connections-dev-perspective
  -- (because the creation and update of a federation config is guarded, we can safely merge the two configs here)
  remotes <-
    (<>)
      <$> getFederationRemotesFromDb
      <*> pure (Map.elems cfgs)

  defFederationDomainConfigs
    & maybe id (\v cfg -> cfg {strategy = v}) mFedStrategy
    & (\cfg -> cfg {remotes = remotes})
    & pure

maxKnownNodes :: Int
maxKnownNodes = 10000

getFederationConfig' :: (MonadClient m) => Map Domain FederationDomainConfig -> Domain -> m (Maybe FederationDomainConfig)
getFederationConfig' cfgs rDomain = case find ((== rDomain) . domain) cfgs of
  Just cfg -> pure . Just $ cfg -- the configuration from the file has precedence (if exists there should not be a db entry at all)
  Nothing -> do
    mCnf <- retry x1 (query1 q (params LocalQuorum (Identity rDomain)))
    case mCnf of
      Just (p, r) -> Just . FederationDomainConfig rDomain p <$> toRestriction rDomain (fromMaybe 0 r)
      Nothing -> pure Nothing
  where
    q :: PrepQuery R (Identity Domain) (FederatedUserSearchPolicy, Maybe Int32)
    q = "SELECT search_policy, restriction FROM federation_remotes WHERE domain = ?"

getFederationRemotesFromDb :: forall m. (MonadClient m) => m [FederationDomainConfig]
getFederationRemotesFromDb = (\(d, p, r) -> FederationDomainConfig d p r) <$$> qry
  where
    qry :: m [(Domain, FederatedUserSearchPolicy, FederationRestriction)]
    qry = do
      res <- retry x1 . query get $ params LocalQuorum ()
      forM res $ \(d, p, rInt) -> do
        (d,p,) <$> toRestriction d (fromMaybe 0 rInt)

    get :: PrepQuery R () (Domain, FederatedUserSearchPolicy, Maybe Int32)
    get = fromString $ "SELECT domain, search_policy, restriction FROM federation_remotes LIMIT " <> show maxKnownNodes

addFederationConfig' :: (MonadClient m) => Map Domain FederationDomainConfig -> FederationDomainConfig -> m AddFederationRemoteResult
addFederationConfig' cfg (FederationDomainConfig rDomain searchPolicy restriction) = do
  -- if a domain already exists in a config, we do not allow to add it to the database
  conflict <- domainExistsInConfig (FederationDomainConfig rDomain searchPolicy restriction)
  if conflict
    then pure $ AddFederationRemoteDivergingConfig cfg
    else do
      l <- length <$> getFederationRemotesFromDb
      if l >= maxKnownNodes
        then pure AddFederationRemoteMaxRemotesReached
        else
          AddFederationRemoteSuccess <$ do
            retry x5 (write addConfig (params LocalQuorum (rDomain, searchPolicy, fromRestriction restriction)))
            case restriction of
              FederationRestrictionByTeam tids ->
                retry x5 . batch . forM_ tids $ addPrepQuery addTeams . (rDomain,)
              FederationRestrictionAllowAll -> pure ()
  where
    -- If remote domain is registered in config file, the version that can be added to the
    -- database must be the same.
    domainExistsInConfig :: (Monad m) => FederationDomainConfig -> m Bool
    domainExistsInConfig fedDomConf = do
      pure $ case Map.lookup (domain fedDomConf) cfg of
        Nothing -> False
        Just fedDomConf' -> fedDomConf' /= fedDomConf

    addConfig :: PrepQuery W (Domain, FederatedUserSearchPolicy, Int32) ()
    addConfig = "INSERT INTO federation_remotes (domain, search_policy, restriction) VALUES (?, ?, ?)"

    addTeams :: PrepQuery W (Domain, TeamId) ()
    addTeams = "INSERT INTO federation_remote_teams (domain, team) VALUES (?, ?)"

updateFederationConfig' :: (MonadClient m) => Map Domain FederationDomainConfig -> FederationDomainConfig -> m UpdateFederationResult
updateFederationConfig' cfgs (FederationDomainConfig rDomain searchPolicy restriction) = do
  -- if a domain already exists in a config, we do not allow update it
  if rDomain `elem` (domain <$> cfgs)
    then pure UpdateFederationRemoteDivergingConfig
    else do
      let configParams =
            ( params
                LocalQuorum
                (searchPolicy, fromRestriction restriction, rDomain)
            )
              { serialConsistency = Just LocalSerialConsistency
              }
      r <- retry x1 (trans updateConfig configParams)
      updateTeams
      case r of
        [] -> pure UpdateFederationRemoteNotFound
        [_] -> pure UpdateFederationSuccess
        _ -> throwM $ ErrorCall "Primary key violation detected federation_remotes"
  where
    updateConfig :: PrepQuery W (FederatedUserSearchPolicy, Int32, Domain) x
    updateConfig = "UPDATE federation_remotes SET search_policy = ?, restriction = ? WHERE domain = ? IF EXISTS"

    updateTeams :: (MonadClient m) => m ()
    updateTeams = retry x5 $ do
      write dropTeams (params LocalQuorum (Identity rDomain))
      case restriction of
        FederationRestrictionByTeam tids ->
          batch . forM_ tids $ addPrepQuery insertTeam . (rDomain,)
        FederationRestrictionAllowAll -> pure ()

    dropTeams :: PrepQuery W (Identity Domain) ()
    dropTeams = "DELETE FROM federation_remote_teams WHERE domain = ?"

    insertTeam :: PrepQuery W (Domain, TeamId) ()
    insertTeam = "INSERT INTO federation_remote_teams (domain, team) VALUES (?, ?)"

addFederationRemoteTeam' :: (MonadClient m) => Map Domain FederationDomainConfig -> Domain -> TeamId -> m AddFederationRemoteTeamResult
addFederationRemoteTeam' cfgs rDomain tid = do
  mDom <- getFederationConfig' cfgs rDomain
  case mDom of
    Nothing ->
      pure AddFederationRemoteTeamDomainNotFound
    Just (FederationDomainConfig _ _ FederationRestrictionAllowAll) ->
      pure AddFederationRemoteTeamRestrictionAllowAll
    Just _ -> do
      retry x1 $ write add (params LocalQuorum (rDomain, tid))
      pure AddFederationRemoteTeamSuccess
  where
    add :: PrepQuery W (Domain, TeamId) ()
    add = "INSERT INTO federation_remote_teams (domain, team) VALUES (?, ?)"

getFederationRemoteTeams' :: (MonadClient m) => Domain -> m [FederationRemoteTeam]
getFederationRemoteTeams' rDomain = do
  fmap (FederationRemoteTeam . runIdentity) <$> retry x1 (query get (params LocalQuorum (Identity rDomain)))
  where
    get :: PrepQuery R (Identity Domain) (Identity TeamId)
    get = "SELECT team FROM federation_remote_teams WHERE domain = ?"

removeFederationRemoteTeam' :: (MonadClient m) => Domain -> TeamId -> m ()
removeFederationRemoteTeam' rDomain rteam =
  retry x1 $ write delete (params LocalQuorum (rDomain, rteam))
  where
    delete :: PrepQuery W (Domain, TeamId) ()
    delete = "DELETE FROM federation_remote_teams WHERE domain = ? AND team = ?"

backendFederatesWithImpl ::
  (MonadClient m) =>
  Remote (Maybe TeamId) ->
  Map Domain FederationDomainConfig ->
  Maybe FederationStrategy ->
  m Bool
backendFederatesWithImpl (tUntagged -> Qualified Nothing rDomain) staticCfgs = \case
  Nothing -> pure False
  Just AllowNone -> pure False
  Just AllowAll -> pure True
  Just AllowDynamic -> do
    getFederationConfig' staticCfgs rDomain >>= \case
      Nothing -> pure False
      Just c -> pure $ restriction c == FederationRestrictionAllowAll
backendFederatesWithImpl (tUntagged -> Qualified (Just rTeam) rDomain) staticCfgs = \case
  Nothing -> pure False
  Just AllowNone -> pure False
  Just AllowAll -> pure True
  Just AllowDynamic ->
    getFederationConfig' staticCfgs rDomain >>= \case
      Nothing -> pure False
      Just (FederationDomainConfig _ _ FederationRestrictionAllowAll) ->
        pure True
      Just (FederationDomainConfig _ _ (FederationRestrictionByTeam ts)) -> pure $ rTeam `elem` ts

data RestrictionException = RestrictionException Int32

instance Show RestrictionException where
  show (RestrictionException v) =
    "Expected a RestrictionPolicy encoding, but found a value " <> show v

instance Exception RestrictionException

toRestriction :: (MonadClient m) => Domain -> Int32 -> m FederationRestriction
toRestriction _ 0 = pure FederationRestrictionAllowAll
toRestriction dom 1 =
  fmap FederationRestrictionByTeam $
    runIdentity <$$> retry x1 (query getTeams (params LocalQuorum (Identity dom)))
  where
    getTeams :: PrepQuery R (Identity Domain) (Identity TeamId)
    getTeams = fromString $ "SELECT team FROM federation_remote_teams WHERE domain = ?"
toRestriction _ v = throwM . RestrictionException $ v

fromRestriction :: FederationRestriction -> Int32
fromRestriction FederationRestrictionAllowAll = 0
fromRestriction (FederationRestrictionByTeam _) = 1
