content: {-# OPTIONS -Wno-ambiguous-fields -Wno-orphans #-}
  
  -- 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 Work where
  
  import Barbies
  import Barbies.Bare
  import Cassandra
  import Cassandra qualified as C
  import Conduit
  import Control.Monad.Catch
  import Data.Conduit.List qualified as C
  import Data.Default
  import Data.Id
  import Data.Misc
  import Data.Schema
  import Data.Time.Clock
  import Data.Time.Clock.POSIX (posixSecondsToUTCTime, utcTimeToPOSIXSeconds)
  import Imports
  import System.Logger qualified as Log
  import Wire.API.Conversation.Protocol
  import Wire.API.MLS.CipherSuite
  import Wire.API.Team.Feature
  
  data MigrationOpts = MigrationOpts
    { granularity :: Int,
      logger :: Log.Logger,
      clientState :: ClientState
    }
  
  runCommand :: MigrationOpts -> IO ()
  runCommand opts = do
    countRef <- newIORef (0, 0)
    runClient opts.clientState $
      runConduit $
        getFeatures
          .| C.concat
          .| C.mapM_ (writeFeatures opts countRef)
  
    (migrated, skipped) <- readIORef countRef
    Log.info opts.logger $
      Log.msg ("migration complete" :: ByteString)
        . Log.field "migrated" migrated
        . Log.field "skipped" skipped
  
  pageSize :: Int32
  pageSize = 1000
  
  ----------------------------------------------------------------------------
  -- Combined Row Type
  
  type FeatureRow =
    ( TeamId,
      Maybe EnforceAppLock, -- app_lock_enforce
      Maybe Int32, -- app_lock_inactivity_timeout_secs
      Maybe FeatureStatus, -- app_lock_status
      Maybe LockStatus, -- conference_calling
      Maybe One2OneCalls, -- conference_calling_one_to_one
      Maybe FeatureStatus, -- conference_calling_status
      Maybe FeatureStatus, -- digital_signatures
      Maybe LockStatus, -- domain_registration_lock_status
      Maybe FeatureStatus, -- domain_registration_status
      Maybe Text, -- enforce_file_download_location
      Maybe LockStatus, -- enforce_file_download_location_lock_status
      Maybe FeatureStatus, -- enforce_file_download_location_status
      Maybe FeatureStatus, -- expose_invitation_urls_to_team_admin
      Maybe FeatureStatus, -- file_sharing
      Maybe LockStatus, -- file_sharing_lock_status
      Maybe LockStatus, -- guest_links_lock_status
      Maybe FeatureStatus, -- guest_links_status
      Maybe FeatureStatus, -- legalhold_status
      Maybe FeatureStatus, -- limited_event_fanout_status
      Maybe (Cassandra.Set CipherSuiteTag), -- mls_allowed_ciphersuites
      Maybe CipherSuiteTag, -- mls_default_ciphersuite
      Maybe ProtocolTag, -- mls_default_protocol
      Maybe HttpsUrl, -- mls_e2eid_acme_discovery_url
      Maybe HttpsUrl, -- mls_e2eid_crl_proxy
      Maybe Int32, -- mls_e2eid_grace_period
      Maybe LockStatus, -- mls_e2eid_lock_status
      Maybe FeatureStatus, -- mls_e2eid_status
      Maybe Bool, -- mls_e2eid_use_proxy_on_mobile
      Maybe LockStatus, -- mls_lock_status
      Maybe OptionalUTCTime, -- mls_migration_finalise_regardless_after
      Maybe LockStatus, -- mls_migration_lock_status
      Maybe OptionalUTCTime, -- mls_migration_start_time
      Maybe FeatureStatus, -- mls_migration_status
      Maybe (Cassandra.Set UserId), -- mls_protocol_toggle_users
      Maybe FeatureStatus, -- mls_status
      Maybe (Cassandra.Set ProtocolTag), -- mls_supported_protocols
      Maybe LockStatus, -- outlook_cal_integration_lock_status
      Maybe FeatureStatus, -- outlook_cal_integration_status
      Maybe FeatureStatus, -- search_visibility_inbound_status
      Maybe FeatureStatus, -- search_visibility_status
      Maybe LockStatus, -- self_deleting_messages_lock_status
      Maybe FeatureStatus, -- self_deleting_messages_status
      Maybe Int32, -- self_deleting_messages_ttl
      Maybe LockStatus, -- snd_factor_password_challenge_lock_status
      Maybe FeatureStatus, -- snd_factor_password_challenge_status
      Maybe FeatureStatus, -- sso_status
      Maybe FeatureStatus -- validate_saml_emails
    )
  
  ----------------------------------------------------------------------------
  -- Query
  
  rowQuery :: PrepQuery R () FeatureRow
  rowQuery =
    "select \
    \ team_id, \
    \ app_lock_enforce, \
    \ app_lock_inactivity_timeout_secs, \
    \ app_lock_status, \
    \ conference_calling, \
    \ conference_calling_one_to_one, \
    \ conference_calling_status, \
    \ digital_signatures, \
    \ domain_registration_lock_status, \
    \ domain_registration_status, \
    \ enforce_file_download_location, \
    \ enforce_file_download_location_lock_status, \
    \ enforce_file_download_location_status, \
    \ expose_invitation_urls_to_team_admin, \
    \ file_sharing, \
    \ file_sharing_lock_status, \
    \ guest_links_lock_status, \
    \ guest_links_status, \
    \ legalhold_status, \
    \ limited_event_fanout_status, \
    \ mls_allowed_ciphersuites, \
    \ mls_default_ciphersuite, \
    \ mls_default_protocol, \
    \ mls_e2eid_acme_discovery_url, \
    \ mls_e2eid_crl_proxy, \
    \ mls_e2eid_grace_period, \
    \ mls_e2eid_lock_status, \
    \ mls_e2eid_status, \
    \ mls_e2eid_use_proxy_on_mobile, \
    \ mls_lock_status, \
    \ mls_migration_finalise_regardless_after, \
    \ mls_migration_lock_status, \
    \ mls_migration_start_time, \
    \ mls_migration_status, \
    \ mls_protocol_toggle_users, \
    \ mls_status, \
    \ mls_supported_protocols, \
    \ outlook_cal_integration_lock_status, \
    \ outlook_cal_integration_status, \
    \ search_visibility_inbound_status, \
    \ search_visibility_status, \
    \ self_deleting_messages_lock_status, \
    \ self_deleting_messages_status, \
    \ self_deleting_messages_ttl, \
    \ snd_factor_password_challenge_lock_status, \
    \ snd_factor_password_challenge_status, \
    \ sso_status, \
    \ validate_saml_emails \
    \ from team_features"
  
  writeFeatures :: (MonadClient m, MonadCatch m) => MigrationOpts -> IORef (Int, Int) -> FeatureRow -> m ()
  writeFeatures
    opts
    countRef
    ( team_id,
      app_lock_enforce,
      app_lock_inactivity_timeout_secs,
      app_lock_status,
      conference_calling,
      conference_calling_one_to_one,
      conference_calling_status,
      digital_signatures,
      domain_registration_lock_status,
      domain_registration_status,
      enforce_file_download_location,
      enforce_file_download_location_lock_status,
      enforce_file_download_location_status,
      expose_invitation_urls_to_team_admin,
      file_sharing,
      file_sharing_lock_status,
      guest_links_lock_status,
      guest_links_status,
      legalhold_status,
      limited_event_fanout_status,
      mls_allowed_ciphersuites,
      mls_default_ciphersuite,
      mls_default_protocol,
      mls_e2eid_acme_discovery_url,
      mls_e2eid_crl_proxy,
      mls_e2eid_grace_period,
      mls_e2eid_lock_status,
      mls_e2eid_status,
      mls_e2eid_use_proxy_on_mobile,
      mls_lock_status,
      mls_migration_finalise_regardless_after,
      mls_migration_lock_status,
      mls_migration_start_time,
      mls_migration_status,
      mls_protocol_toggle_users,
      mls_status,
      mls_supported_protocols,
      outlook_cal_integration_lock_status,
      outlook_cal_integration_status,
      search_visibility_inbound_status,
      search_visibility_status,
      self_deleting_messages_lock_status,
      self_deleting_messages_status,
      self_deleting_messages_ttl,
      snd_factor_password_challenge_lock_status,
      snd_factor_password_challenge_status,
      sso_status,
      validate_saml_emails
      ) = do
      state <- getMigrationState team_id
      when (state /= MigrationCompleted) $ do
        onException
          ( do
              -- set team features to read-only
              setMigrationState team_id MigrationInProgress
  
              retry x5 . batch $ do
                setConsistency LocalQuorum
                setType BatchLogged
  
                writeFeatureB team_id $
                  (def :: LockableFeaturePatch AppLockConfig)
                    { status = app_lock_status,
                      config = Just $ AppLockConfig @Covered app_lock_enforce app_lock_inactivity_timeout_secs
                    }
                writeFeatureB team_id $
                  LockableFeaturePatch
                    { status = conference_calling_status,
                      lockStatus = conference_calling,
                      config =
                        Just $ ConferenceCallingConfig @Covered conference_calling_one_to_one
                    }
  
                writeFeature team_id $
                  (def :: LockableFeaturePatch DigitalSignaturesConfig) {status = digital_signatures}
  
                writeFeature team_id $
                  (def :: LockableFeaturePatch DomainRegistrationConfig)
                    { status = domain_registration_status,
                      lockStatus = domain_registration_lock_status
                    }
  
                writeFeatureB team_id $
                  LockableFeaturePatch
                    { status = enforce_file_download_location_status,
                      lockStatus = enforce_file_download_location_lock_status,
                      config =
                        Just $
                          EnforceFileDownloadLocationConfig @Covered $
                            case enforce_file_download_location of
                              Nothing -> Nothing
                              Just "" -> Just Nothing
                              Just loc -> Just (Just loc)
                    }
  
                writeFeature team_id $
                  (def :: LockableFeaturePatch ExposeInvitationURLsToTeamAdminConfig)
                    { status = expose_invitation_urls_to_team_admin
                    }
  
                writeFeature team_id $
                  (def :: LockableFeaturePatch FileSharingConfig)
                    { status = file_sharing,
                      lockStatus = file_sharing_lock_status
                    }
  
                writeFeature team_id $
                  (def :: LockableFeaturePatch GuestLinksConfig)
                    { status = guest_links_status,
                      lockStatus = guest_links_lock_status
                    }
  
                writeFeature team_id $
                  (def :: LockableFeaturePatch LegalholdConfig)
                    { status = legalhold_status
                    }
  
                writeFeature team_id $
                  (def :: LockableFeaturePatch LimitedEventFanoutConfig)
                    { status = limited_event_fanout_status
                    }
  
                writeFeatureB team_id $
                  LockableFeaturePatch
                    { status = mls_status,
                      lockStatus = mls_lock_status,
                      config =
                        Just $
                          ( MLSConfig @Covered
                              (fmap C.fromSet mls_protocol_toggle_users)
                              mls_default_protocol
                              (fmap C.fromSet mls_allowed_ciphersuites)
                              mls_default_ciphersuite
                              (fmap C.fromSet mls_supported_protocols)
                          )
                    }
  
                writeFeatureB team_id $
                  LockableFeaturePatch
                    { status = mls_e2eid_status,
                      lockStatus = mls_e2eid_lock_status,
                      config =
                        Just $
                          ( MlsE2EIdConfig @Covered
                              (fmap fromIntegral mls_e2eid_grace_period)
                              (Alt mls_e2eid_acme_discovery_url)
                              (Alt mls_e2eid_crl_proxy)
                              (maybe def UseProxyOnMobile mls_e2eid_use_proxy_on_mobile)
                          )
                    }
  
                writeFeatureB team_id $
                  LockableFeaturePatch
                    { status = mls_migration_status,
                      lockStatus = mls_migration_lock_status,
                      config =
                        Just $
                          ( MlsMigrationConfig @Covered
                              (fmap unOptionalUTCTime mls_migration_start_time)
                              (fmap unOptionalUTCTime mls_migration_finalise_regardless_after)
                          )
                    }
  
                writeFeature team_id $
                  (def :: LockableFeaturePatch OutlookCalIntegrationConfig)
                    { status = outlook_cal_integration_status,
                      lockStatus = outlook_cal_integration_lock_status
                    }
  
                writeFeature team_id $
                  (def :: LockableFeaturePatch SearchVisibilityInboundConfig)
                    { status = search_visibility_inbound_status
                    }
  
                writeFeature team_id $
                  (def :: LockableFeaturePatch SearchVisibilityAvailableConfig)
                    { status = search_visibility_status
                    }
  
                writeFeatureB team_id $
                  LockableFeaturePatch
                    { status = self_deleting_messages_status,
                      lockStatus = self_deleting_messages_lock_status,
                      config =
                        Just $
                          SelfDeletingMessagesConfig @Covered self_deleting_messages_ttl
                    }
  
                writeFeature team_id $
                  (def :: LockableFeaturePatch SndFactorPasswordChallengeConfig)
                    { status = snd_factor_password_challenge_status,
                      lockStatus = snd_factor_password_challenge_lock_status
                    }
  
                writeFeature team_id $
                  (def :: LockableFeaturePatch SSOConfig) {status = sso_status}
  
                writeFeature team_id $
                  (def :: LockableFeaturePatch ValidateSAMLEmailsConfig)
                    { status = validate_saml_emails
                    }
  
              -- set migration state to completed
              setMigrationState team_id MigrationCompleted
          )
          (setMigrationState team_id MigrationNotStarted)
  
      modifyIORef countRef $ \(migrated, skipped) ->
        if state == MigrationNotStarted
          then (migrated + 1, skipped)
          else (migrated, skipped + 1)
  
      (migrated, skipped) <- readIORef countRef
      when ((migrated + skipped) `mod` opts.granularity == 0) $ do
        Log.info opts.logger $
          Log.msg ("migration progress" :: ByteString)
            . Log.field "migrated" migrated
            . Log.field "skipped" skipped
  
  ----------------------------------------------------------------------------
  
  -- Pagination
  
  getFeatures :: (MonadClient m) => ConduitM () [FeatureRow] m ()
  getFeatures = paginateC rowQuery (paramsP LocalQuorum () pageSize) x5
  
  ----------------------------------------------------------------------------
  -- Instances (unchanged)
  
  instance Cql EnforceAppLock where
    ctype = Tagged IntColumn
    toCql (EnforceAppLock False) = CqlInt 0
    toCql (EnforceAppLock True) = CqlInt 1
    fromCql (CqlInt n) = case n of
      0 -> pure (EnforceAppLock False)
      1 -> pure (EnforceAppLock True)
      _ -> Left "fromCql EnforceAppLock: int out of range"
    fromCql _ = Left "fromCql EnforceAppLock: int expected"
  
  -- Optional time stamp. A 'Nothing' value is represented as 0.
  newtype OptionalUTCTime = OptionalUTCTime {unOptionalUTCTime :: Maybe UTCTime}
  
  instance Cql OptionalUTCTime where
    ctype = Tagged (untag (ctype @UTCTime))
    toCql = toCql . fromMaybe (posixSecondsToUTCTime 0) . unOptionalUTCTime
    fromCql x = do
      t <- fromCql x
      pure . OptionalUTCTime $ guard (utcTimeToPOSIXSeconds t /= 0) $> t
  
  writeFeature ::
    forall cfg.
    (IsFeatureConfig cfg) =>
    TeamId ->
    LockableFeaturePatch cfg ->
    BatchM ()
  writeFeature tid feat =
    writeDbFeature @cfg
      tid
      feat {config = fmap (DbConfig . schemaToJSON) feat.config}
  
  writeFeatureB ::
    forall b cfg.
    ( cfg ~ b Bare Identity,
      ToSchema (b Covered Maybe),
      TraversableB (b Covered),
      IsFeatureConfig cfg
    ) =>
    TeamId ->
    LockableFeaturePatch (b Covered Maybe) ->
    BatchM ()
  writeFeatureB tid feat = do
    let dbConfig = feat.config >>= serialiseBarbieConfig
    writeDbFeature @cfg tid feat {config = dbConfig}
  
  writeDbFeature ::
    forall cfg.
    (IsFeatureConfig cfg) =>
    TeamId ->
    LockableFeaturePatch DbConfig ->
    BatchM ()
  writeDbFeature tid feat =
    let q :: PrepQuery W (Maybe FeatureStatus, Maybe LockStatus, Maybe DbConfig, TeamId, Text) ()
        q = "update team_features_dyn set status = ?, lock_status = ?, config = ? where team = ? and feature = ?"
     in addPrepQuery q (feat.status, feat.lockStatus, feat.config, tid, featureName @cfg)
  
  serialiseBarbieConfig ::
    ( TraversableB (b Covered),
      ToSchema (b Covered Maybe)
    ) =>
    b Covered Maybe ->
    Maybe DbConfig
  serialiseBarbieConfig cfg = do
    -- ensure at least one field is set
    void $ getAlt (bfoldMap (Alt . void) cfg)
    pure . DbConfig $ schemaToJSON cfg
  
  setMigrationState :: (MonadClient m) => TeamId -> TeamFeatureMigrationState -> m ()
  setMigrationState tid state = do
    let q :: PrepQuery W (TeamFeatureMigrationState, TeamId) ()
        q = "update team_features set migration_state = ? where team_id = ?"
    retry x5 $ write q (params LocalQuorum (state, tid))
  
  getMigrationState :: (MonadClient m) => TeamId -> m TeamFeatureMigrationState
  getMigrationState tid = do
    let q :: PrepQuery R (Identity TeamId) (Identity (Maybe TeamFeatureMigrationState))
        q = "select migration_state from team_features where team_id = ?"
    fmap (fromMaybe def . (runIdentity =<<)) $
      retry x1 $
        query1 q (params LocalQuorum (Identity tid))
  