-- 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.PostgresMigrations where

import Imports
import Prometheus
import System.Logger qualified as Log
import UnliftIO
import Wire.BackgroundWorker.Env
import Wire.BackgroundWorker.Util
import Wire.CodeStore.Migration
import Wire.ConversationStore.Migration
import Wire.DomainRegistrationStore.Migration
import Wire.Migration (MigrationOptions)
import Wire.TeamFeatureStore.Migration

conversations :: MigrationOptions -> AppT IO CleanupAction
conversations migOpts = do
  cassClient <- asks (.cassandraGalley)
  pgPool <- asks (.hasqlPool)
  logger <- asks (.logger)

  Log.info logger $ Log.msg (Log.val "starting conversation migration")
  convMigCounter <- register $ counter $ Prometheus.Info "wire_local_convs_migrated_to_pg" "Number of local conversations migrated to Postgresql"
  convMigFinished <- register $ counter $ Prometheus.Info "wire_local_convs_migration_finished" "Whether the conversation migration to Postgresql is finished successfully"
  convMigFailed <- register $ counter $ Prometheus.Info "wire_local_convs_migration_failed" "Whether the conversation migration to Postgresql has failed"
  convMigDuration <- register $ vector "outcome" $ histogram (Prometheus.Info "wire_local_convs_migration_duration_seconds" "Duration of local conversation migration attempts") defaultBuckets
  userMigCounter <- register $ counter $ Prometheus.Info "wire_user_remote_convs_migrated_to_pg" "Number of users whose remote conversation membership data is migrated to Postgresql"
  userMigFinished <- register $ counter $ Prometheus.Info "wire_user_remote_convs_migration_finished" "Whether the migration of remote conversation membership data to Postgresql is finished successfully"
  userMigFailed <- register $ counter $ Prometheus.Info "wire_user_remote_convs_migration_failed" "Whether the migration of remote conversation membership data to Postgresql has failed"
  userMigDuration <- register $ vector "outcome" $ histogram (Prometheus.Info "wire_user_remote_convs_migration_duration_seconds" "Duration of remote conversation membership migration attempts") defaultBuckets

  convLoop <- async . lift $ migrateConvsLoop migOpts cassClient pgPool logger convMigCounter convMigFinished convMigFailed convMigDuration
  userLoop <- async . lift $ migrateUsersLoop migOpts cassClient pgPool logger userMigCounter userMigFinished userMigFailed userMigDuration

  Log.info logger $ Log.msg (Log.val "started conversation migration")
  pure $ do
    Log.info logger $ Log.msg (Log.val "cancelling conversation migration")
    cancel convLoop
    cancel userLoop

conversationCodes :: MigrationOptions -> AppT IO CleanupAction
conversationCodes migOpts = do
  cassClient <- asks (.cassandraGalley)
  pgPool <- asks (.hasqlPool)
  logger <- asks (.logger)
  Log.info logger $ Log.msg (Log.val "starting conversation codes migration")
  count <- register $ counter $ Prometheus.Info "wire_conv_codes_migrated_to_pg" "Number of conversation codes migrated to Postgresql"
  finished <- register $ counter $ Prometheus.Info "wire_conv_codes_migration_finished" "Whether the conversation codes migration to Postgresql is finished successfully"
  failed <- register $ counter $ Prometheus.Info "wire_conv_codes_migration_failed" "Whether the conversation codes migration to Postgresql has failed"
  duration <- register $ vector "outcome" $ histogram (Prometheus.Info "wire_conv_codes_migration_duration_seconds" "Duration of conversation code migration attempts") defaultBuckets

  migrationLoop <- async . lift $ migrateCodesLoop migOpts cassClient pgPool logger count finished failed duration

  Log.info logger $ Log.msg (Log.val "started conversation codes migration")
  pure $ do
    Log.info logger $ Log.msg (Log.val "cancelling conversation codes migration")
    cancel migrationLoop

teamFeatures :: MigrationOptions -> AppT IO CleanupAction
teamFeatures migOpts = do
  cassClient <- asks (.cassandraGalley)
  pgPool <- asks (.hasqlPool)
  logger <- asks (.logger)
  Log.info logger $ Log.msg (Log.val "starting team features migration")
  count <- register $ counter $ Prometheus.Info "wire_team_features_migrated_to_pg" "Number of team features migrated to Postgresql"
  finished <- register $ counter $ Prometheus.Info "wire_team_features_migration_finished" "Whether the team features migration to Postgresql is finished successfully"
  failed <- register $ counter $ Prometheus.Info "wire_team_features_migration_failed" "Whether the team features migration to Postgresql has failed"
  duration <- register $ vector "outcome" $ histogram (Prometheus.Info "wire_team_features_migration_duration_seconds" "Duration of team feature migration attempts") defaultBuckets

  migrationLoop <- async . lift $ migrateTeamFeaturesLoop migOpts cassClient pgPool logger count finished failed duration

  Log.info logger $ Log.msg (Log.val "started team features migration")
  pure $ do
    Log.info logger $ Log.msg (Log.val "cancelling team features migration")
    cancel migrationLoop

domainRegistration :: MigrationOptions -> AppT IO CleanupAction
domainRegistration migOpts = do
  cassClient <- asks (.cassandraBrig)
  pgPool <- asks (.hasqlPool)
  logger <- asks (.logger)
  Log.info logger $ Log.msg (Log.val "starting domain registration migration")
  count <- register $ counter $ Prometheus.Info "wire_domain_registration_migrated_to_pg" "Number of domain registration rows migrated to Postgresql"
  finished <- register $ counter $ Prometheus.Info "wire_domain_registration_migration_finished" "Whether the domain registration migration to Postgresql is finished successfully"
  failed <- register $ counter $ Prometheus.Info "wire_domain_registration_migration_failed" "Whether the domain registration migration to Postgresql has failed"
  duration <- register $ vector "outcome" $ histogram (Prometheus.Info "wire_domain_registration_migration_duration_seconds" "Duration of domain registration migration attempts") defaultBuckets

  migrationLoop <- async . lift $ migrateDomainRegistrationsLoop migOpts cassClient pgPool logger count finished failed duration

  Log.info logger $ Log.msg (Log.val "started domain registration migration")
  pure $ do
    Log.info logger $ Log.msg (Log.val "cancelling domain registration migration")
    cancel migrationLoop
