//
// Copyright 2023 Signal Messenger, LLC
// SPDX-License-Identifier: AGPL-3.0-only
//

import Foundation

/**
 * Utility class to simplify the transition when we start handling a previously-unknown storage
 * service field for the first time.
 *
 * ---
 * All this is rather complicated. Lets start with an example:
 *
 * Say we have some setting for dark mode (isDarkMode) on iOS and Desktop.
 * Desktop has been syncing the setting to storage service. iOS has not. We want to start syncing it.
 * Locally we have three states: UNSET, FALSE, TRUE.
 *
 * If we just naively added the field to our local proto definition, we'd recognize it the next time
 * we read from storage service, but our naive merging logic would overwrite our local state
 * with the state we see in storage service. Uh oh! We should respect local state!
 *
 * Ok so maybe we prefer local state for the primary, so we ignore the value when we read and
 * keep our local state instead. We do that, write our local state, Desktop reads it, and then the user
 * updates the value on Desktop. Next time we read from storage service on iOS...we drop the
 * value the user set on Desktop in favor of our local value! Uh oh!
 *
 * In short, what we need is to ignore the value on reads UNTIL we get a chance to write at least once,
 * and then we can start respecting the value on reads.
 * This class solves this problem generally.
 *
 * ---
 *
 * Now, how does this class actually work? The scaffolding handles all the hooks for you so
 * that you can write a "migration" that does some subset of three things:
 * 1. ``MergeUnknownFields`` - do any special one-time handling the _first_ time we learn about
 *   some previously-unknown proto field. This lets you e.g. do any one-time merging of remote state
 *   into local state. In the isDarkMode example, we would take this opportunity to overwrite a local
 *   UNSET value with any value in storage service.
 *
 * Next are two methods that run on every record we read/write until we get the chance to successfully
 * write to storage service once. This lets us intercept records _before_ the standard merging code runs
 * and perform modifications while we are in this "in between" state when we first handle unknown fields.
 * 2. ``InterceptRemoteManifest`` - Take a record from a storage service manifest we fetched
 *   remotely and are about to merge with local state, and modify it _before_ we merge.
 *   In the isDarkMode example, we would overwrite the remote state with local state, so that the merge
 *   code always sees the remote state and local state as the same and changes nothing.
 * 3. ``InterceptLocalManifest`` - Take a record from a storage service manifest we generated
 *   locally and are about to _write_ to a remote, and modify it _before_ we write it.
 *   In the isDarkMode example we would do nothing in this step, but if we had some more complex merging
 *   logic that reconciles local and remote state, you could imagine this coming in handy.
 *
 * So the general flow is:
 * 1. Have some unknown field (isDarkMode) in storage service
 * 2. Update to a build that knows about this field for the first time
 * 3. Once, the _very first_ time on launch of this new build: run ``MergeUnknownFields`` providing
 *   all records with unknown fields. (May be empty. See ##Unknown Field Default Values## below).
 * 4. If we read from storage service: pass every read record to ``InterceptRemoteManifest`` before
 *   merging locally. (We pretend isDarkMode remote state is the same as local state so we never overwrite local state yet).
 * 5. If we write to storage service: pass every generated record to ``InterceptLocalManifest`` before
 *   writing remotely.
 * 6. Repeat 4 and 5 for every read/write until some write succeeds (step 3 is never repeated!)
 * 7. Succesfully write to storage service
 * 8. Done! Migration marked complete.
 *
 * To see this in pseudocode for the isDarkMode example, see the ``MigrationId.noOpExample`` migration below.
 *
 * ---
 *
 * There are some gotchas with ``MergeUnknownFields``.
 *
 * ##Unknown Field Default Values##
 *
 * One quirk of unknown fields in protos: if a field uses a primitive type (e.g. bool) and
 * is set to its default value (e.g. false), the serialized representation of that proto simply
 * omits the field entirely.
 * When a client that is unaware of the field parses such a proto, it doesn't know to
 * look for the field and doesn't see anything in the serialized bytes, so it has no idea
 * there is an unknown field at all.
 *
 * This means our "records with unknown fields" actually means "records with unknown
 * fields set to non-default values". So your ``MergeUnknownFields`` implementation
 * **SHOULD NOT** assume it is getting every record that had an unknown field; it only
 * gets those with non-default values. (And might actually be an empty array if all records
 * use default values!)
 *
 * What this means in practice is you should NOT write for loops over the passed in records;
 * instead you should loop over objects in our database (e.g. TSGroupThreads) and match those
 * against any passed in records (e.g. by groupId) and, if not present, assume the default value.
 *
 * ##Triggering Storage Service Writes##
 *
 * Most of the time, if you're merging remote state in ``MergeUnknownFields``, you will want
 * to write back the results of the merge to storage service. The way to do this is to mark any updated
 * record as needing an update via normal mechanisms, e.g. for the local account record you would call
 * `storageServiceManager.recordPendingLocalAccountUpdates()`
 */
public class StorageServiceUnknownFieldMigrator {

    fileprivate enum MigrationId: UInt, CaseIterable {
        case noOpExample = 0
        case dontNotifyForMentionsIfMuted = 1

        // MARK: - Migration Insertion Point

        // Never, ever, ever insert another migration before an existing one.
        // Increase the value only.

        static var highestKnownValue: UInt {
            return allCases.lazy.map(\.rawValue).max() ?? 0
        }
    }

    fileprivate enum Actions<RecordType: MigrateableStorageServiceRecordType> {
        typealias MergeUnknownFields = (_ records: [RecordType], _ isPrimaryDevice: Bool, _ tx: DBWriteTransaction) -> Void
        typealias InterceptRemoteManifest = (RecordType, inout RecordType.Builder, _ isPrimaryDevice: Bool, _ tx: DBReadTransaction) -> Void
        typealias InterceptLocalManifest = (RecordType, inout RecordType.Builder, _ isPrimaryDevice: Bool, _ tx: DBReadTransaction) -> Void
    }

    private static func registerMigrations(migrator: Migrator) {
        // This is simply an example migration (commented out) for how you _would_ migrate
        // if you had an isDarkMode setting that already existed both locally and in storage service
        // that needed to be merged such that the local value, if set, "wins" on the primary.
        // See doc comment on this class for context.
        migrator.registerMigration(
            .noOpExample,
            record: StorageServiceProtoAccountRecord.self,
            mergeUnknownFields: { accountRecords, isPrimaryDevice, tx in
                /**
                 let isRemoteDarkMode = accountRecords.first?.isDarkMode ?? false
                 guard isPrimaryDevice else {
                     // Just take what we have in storage service.
                     localDarkModeSetting.set(isRemoteDarkMode, tx: tx)
                 }
                 switch (localDarkModeSetting.get(tx: tx)) {
                 case .TRUE, .FALSE:
                     // Ignore if set locally, schedule an update so we overwrite
                     // storage service with our local value.
                     NSObject.storageServiceManager.recordPendingLocalAccountUpdates()
                 case .UNSET:
                     // Just take what we have in storage service.
                     localDarkModeSetting.set(isRemoteDarkMode, tx: tx)
                 }
                 */
            },
            interceptRemoteManifest: { accountRecord, accountRecordBuilder, isPrimaryDevice, tx in
                /**
                 if isPrimaryDevice {
                     // Until we get the chance to write to storage service, we want to
                     // always use our local value. Overwrite the value to our local value
                     // so that the merge logic uses it.
                     accountRecordBuilder.setIsDarkMode(localDarkModeSetting.get(tx: tx))
                 }
                 */
                return
            },
            interceptLocalManifest: { accountRecord, accountRecordBuilder, isPrimaryDevice, tx in
                /**
                 // Nothing to intercept for writes in this case
                  */
                return
            },
        )

        migrator.registerMigration(
            .dontNotifyForMentionsIfMuted,
            record: StorageServiceProtoGroupV2Record.self,
            mergeUnknownFields: { records, isPrimaryDevice, tx in
                // groupId -> dontNotifyForMentionsIfMuted
                var recordMap = [Data: Bool]()
                records.forEach {
                    if let groupId = (try? GroupV2ContextInfo.deriveFrom(masterKeyData: $0.masterKey))?.groupId.serialize() {
                        recordMap[groupId] = $0.dontNotifyForMentionsIfMuted
                    }
                }
                ThreadFinder().enumerateGroupThreads(tx: tx) { groupThread -> Bool in
                    let remoteValue: TSThreadMentionNotificationMode =
                        (recordMap[groupThread.groupId] ?? false) ? .never : .always
                    if isPrimaryDevice {
                        // On primaries, only set if previously unset.
                        if groupThread.mentionNotificationMode == .default {
                            groupThread.updateWithMentionNotificationMode(remoteValue, wasLocallyInitiated: false, transaction: tx)
                        } else {
                            // Schedule an update so we put up our local state onto storageService.
                            SSKEnvironment.shared.storageServiceManagerRef.recordPendingUpdates(groupModel: groupThread.groupModel)
                        }
                    } else {
                        // On secondaries, take the value from storage service.
                        groupThread.updateWithMentionNotificationMode(remoteValue, wasLocallyInitiated: false, transaction: tx)
                    }

                    return true
                }
            },
            interceptRemoteManifest: { record, recordBuilder, isPrimaryDevice, tx in
                // when we read a remote group on the primary, override with our local
                // value so we dont overwrite the local value with the remote's value
                // until we've had the chance to write ourselves.
                if isPrimaryDevice {
                    guard
                        let groupId = (try? GroupV2ContextInfo.deriveFrom(masterKeyData: record.masterKey))?.groupId.serialize(),
                        let groupThread = TSGroupThread.fetch(groupId: groupId, transaction: tx)
                    else {
                        return
                    }
                    switch groupThread.mentionNotificationMode {
                    case .default:
                        // Don't overwrite.
                        break
                    case .always:
                        recordBuilder.setDontNotifyForMentionsIfMuted(false)
                    case .never:
                        recordBuilder.setDontNotifyForMentionsIfMuted(true)
                    }
                }
            },
            interceptLocalManifest: { record, recordBuilder, isPrimaryDevice, tx in
                // Nothing to do
            },
        )

        // MARK: - Migration Insertion Point
    }

    // MARK: - Public methods

    // If you are just writing a new migration, you don't need to worry about these.

    /// Check this before merging records with unknown fields; if true, call ``runMigrationsForRecordsWithUnknownFields``
    public static func needsAnyUnknownFieldsMigrations(tx: DBReadTransaction) -> Bool {
        return !necessaryMigrations(forKey: Keys.lastRunUnknownFieldsMerge, tx: tx).isEmpty
    }

    /// Check this before merging records from a remote manifest; if true, call ``interceptRemoteManifestBeforeMerging``
    public static func shouldInterceptRemoteManifestBeforeMerging(tx: DBReadTransaction) -> Bool {
        return !necessaryMigrations(forKey: Keys.lastSuccessfulStorageServiceWrite, tx: tx).isEmpty
    }

    /// Check this before uploading records generated locally; if true, call ``interceptLocalManifestBeforeUploading``
    public static func shouldInterceptLocalManifestBeforeUploading(tx: DBReadTransaction) -> Bool {
        return !necessaryMigrations(forKey: Keys.lastSuccessfulStorageServiceWrite, tx: tx).isEmpty
    }

    /// Call this after every succesful write of a manifest to Storage Service.
    public static func didWriteToStorageService(tx: DBWriteTransaction) {
        kvStore.setUInt(MigrationId.highestKnownValue, key: Keys.lastSuccessfulStorageServiceWrite, transaction: tx)
    }

    /// Given an array of every record from the latest synced manifest known to have unknown fields, runs any necessary migrations.
    public static func runMigrationsForRecordsWithUnknownFields(
        records: [any MigrateableStorageServiceRecordType],
        tx: DBWriteTransaction,
    ) {
        return _runMigrationsForRecordsWithUnknownFields(records: records, tx: tx)
    }

    public static func interceptRemoteManifestBeforeMerging<RecordType>(
        record: RecordType,
        tx: DBReadTransaction,
    ) -> RecordType {
        guard let recordTypecast = record as? (any MigrateableStorageServiceRecordType) else {
            // Not migrateable. Just no-op.
            return record
        }
        return _interceptRemoteManifestBeforeMerging(record: recordTypecast, tx: tx) as! RecordType
    }

    public static func interceptLocalManifestBeforeUploading<RecordType>(
        record: RecordType,
        tx: DBReadTransaction,
    ) -> RecordType {
        guard let recordTypecast = record as? (any MigrateableStorageServiceRecordType) else {
            // Not migrateable. Just no-op.
            return record
        }
        return _interceptLocalManifestBeforeUploading(record: recordTypecast, tx: tx) as! RecordType
    }

    // MARK: - Private Implementation

    private enum Keys {
        // The value at this key is a UInt representing the highest known MigrationId
        // at the time this device last merged unknown fields from a previously-stored
        // storage service manifest. If this number is lower than the current highest known
        // MigrationId, we should run any unknown field merging operation migrations higher
        // than the stored value.
        static let lastRunUnknownFieldsMerge = "lastRunUnknownFieldsMerge"

        // The value at this key is a UInt representing the highest known MigrationId
        // at the time that this device last succesfully updated the manifest in storage service.
        // If this number is lower than the current highest known MigrationId, we should run any
        // higher value migrations' Record mutation operations on every record we read and write
        // locally.
        static let lastSuccessfulStorageServiceWrite = "lastSuccessfulStorageServiceWrite"
    }

    private static var kvStore = KeyValueStore(collection: "StorageServiceUnknownFieldMigrator")

    private class Migrator {
        var migrations: [MigrationId: any StorageServiceUnknownFieldMigration] = [:]

        func registerMigration<RecordType: MigrateableStorageServiceRecordType>(
            _ migrationId: MigrationId,
            record: RecordType.Type,
            mergeUnknownFields: @escaping Actions<RecordType>.MergeUnknownFields,
            interceptRemoteManifest: @escaping Actions<RecordType>.InterceptRemoteManifest,
            interceptLocalManifest: @escaping Actions<RecordType>.InterceptLocalManifest,
        ) {
            let migration = StorageServiceUnknownFieldMigrationImpl(
                id: migrationId,
                mergeUnknownFields: mergeUnknownFields,
                interceptRemoteManifest: interceptRemoteManifest,
                interceptLocalManifest: interceptLocalManifest,
            )
            migrations[migrationId] = migration
        }
    }

    private static var _migrator: Migrator?

    private static var migrator: Migrator {
        if let _migrator {
            return _migrator
        }
        let migrator = Migrator()
        registerMigrations(migrator: migrator)
        _migrator = migrator
        return migrator
    }

    private static func necessaryMigrations(
        forKey key: String,
        tx: DBReadTransaction,
    ) -> LazyFilterSequence<[MigrationId]> {
        guard let latestMigrationId = kvStore.getUInt(key, transaction: tx) else {
            // We've never run any migrations!
            return MigrationId.allCases.lazy.filter { _ in true }
        }
        return MigrationId.allCases.lazy.filter { $0.rawValue > latestMigrationId }
    }

    private static func _runMigrationsForRecordsWithUnknownFields(
        records: [any MigrateableStorageServiceRecordType],
        tx: DBWriteTransaction,
    ) {
        let necessaryMigrations = Self.necessaryMigrations(forKey: Keys.lastRunUnknownFieldsMerge, tx: tx)
        if necessaryMigrations.isEmpty {
            return
        }

        let tsAccountManager = DependenciesBridge.shared.tsAccountManager
        guard let registeredState = try? tsAccountManager.registeredState(tx: tx) else {
            return
        }

        func doMergeUnknownFields<Migration: StorageServiceUnknownFieldMigration>(
            records: [any MigrateableStorageServiceRecordType],
            migration: Migration,
        ) {
            let filteredRecords = records.compactMap { $0 as? Migration.RecordType }
            migration.mergeUnknownFields(filteredRecords, registeredState.isPrimary, tx)
        }

        necessaryMigrations.forEach { migrationId in
            guard let migration = migrator.migrations[migrationId] else {
                return
            }
            doMergeUnknownFields(records: records, migration: migration)
        }

        kvStore.setUInt(MigrationId.highestKnownValue, key: Keys.lastRunUnknownFieldsMerge, transaction: tx)
    }

    private static func _interceptRemoteManifestBeforeMerging<RecordType: MigrateableStorageServiceRecordType>(
        record: RecordType,
        tx: DBReadTransaction,
    ) -> RecordType {
        let necessaryMigrations = Self.necessaryMigrations(forKey: Keys.lastSuccessfulStorageServiceWrite, tx: tx)
        if necessaryMigrations.isEmpty {
            return record
        }

        let tsAccountManager = DependenciesBridge.shared.tsAccountManager
        guard let registeredState = try? tsAccountManager.registeredState(tx: tx) else {
            return record
        }

        var _builder: RecordType.Builder?

        func doModifyRemoteManifest<Migration: StorageServiceUnknownFieldMigration>(
            migration: Migration,
        ) {
            guard RecordType.self == Migration.RecordType.self else {
                return
            }
            let builder = _builder ?? record.asBuilder()
            var typecastBuilder = builder as! Migration.RecordType.Builder
            migration.interceptRemoteManifest(record as! Migration.RecordType, &typecastBuilder, registeredState.isPrimary, tx)
            _builder = (typecastBuilder as! RecordType.Builder)
        }

        necessaryMigrations.forEach { migrationId in
            guard let migration = migrator.migrations[migrationId] else {
                return
            }
            doModifyRemoteManifest(migration: migration)
        }

        if let builder = _builder {
            return builder.buildInfallibly()
        } else {
            // It was unmodified.
            return record
        }
    }

    private static func _interceptLocalManifestBeforeUploading<RecordType: MigrateableStorageServiceRecordType>(
        record: RecordType,
        tx: DBReadTransaction,
    ) -> RecordType {
        let necessaryMigrations = Self.necessaryMigrations(forKey: Keys.lastSuccessfulStorageServiceWrite, tx: tx)
        if necessaryMigrations.isEmpty {
            return record
        }

        let tsAccountManager = DependenciesBridge.shared.tsAccountManager
        guard let registeredState = try? tsAccountManager.registeredState(tx: tx) else {
            return record
        }

        var _builder: RecordType.Builder?

        func doModifyLocalManifest<Migration: StorageServiceUnknownFieldMigration>(
            migration: Migration,
        ) {
            guard RecordType.self == Migration.RecordType.self else {
                return
            }
            let builder = _builder ?? record.asBuilder()
            var typecastBuilder = builder as! Migration.RecordType.Builder
            migration.interceptLocalManifest(record as! Migration.RecordType, &typecastBuilder, registeredState.isPrimary, tx)
            _builder = (typecastBuilder as! RecordType.Builder)
        }

        necessaryMigrations.forEach { migrationId in
            guard let migration = migrator.migrations[migrationId] else {
                return
            }
            doModifyLocalManifest(migration: migration)
        }

        if let builder = _builder {
            return builder.buildInfallibly()
        } else {
            // It was unmodified.
            return record
        }
    }
}

public protocol MigrateableStorageServiceRecordType {
    associatedtype Builder: MigrateableStorageServiceRecordTypeBuilder where Builder.RecordType == Self

    func asBuilder() -> Builder
}

public protocol MigrateableStorageServiceRecordTypeBuilder {
    associatedtype RecordType: MigrateableStorageServiceRecordType

    func buildInfallibly() -> RecordType
}

private protocol StorageServiceUnknownFieldMigration {

    associatedtype RecordType: MigrateableStorageServiceRecordType

    var mergeUnknownFields: StorageServiceUnknownFieldMigrator.Actions<RecordType>.MergeUnknownFields { get }
    var interceptRemoteManifest: StorageServiceUnknownFieldMigrator.Actions<RecordType>.InterceptRemoteManifest { get }
    var interceptLocalManifest: StorageServiceUnknownFieldMigrator.Actions<RecordType>.InterceptLocalManifest { get }
}

extension StorageServiceUnknownFieldMigrator {

    private struct StorageServiceUnknownFieldMigrationImpl<RecordType: MigrateableStorageServiceRecordType>: StorageServiceUnknownFieldMigration {
        let id: StorageServiceUnknownFieldMigrator.MigrationId
        let mergeUnknownFields: Actions<RecordType>.MergeUnknownFields
        let interceptRemoteManifest: Actions<RecordType>.InterceptRemoteManifest
        let interceptLocalManifest: Actions<RecordType>.InterceptLocalManifest
    }
}

extension StorageServiceProtoAccountRecord: MigrateableStorageServiceRecordType {}
extension StorageServiceProtoContactRecord: MigrateableStorageServiceRecordType {}
extension StorageServiceProtoGroupV1Record: MigrateableStorageServiceRecordType {}
extension StorageServiceProtoGroupV2Record: MigrateableStorageServiceRecordType {}

extension StorageServiceProtoAccountRecordBuilder: MigrateableStorageServiceRecordTypeBuilder {}
extension StorageServiceProtoContactRecordBuilder: MigrateableStorageServiceRecordTypeBuilder {}
extension StorageServiceProtoGroupV1RecordBuilder: MigrateableStorageServiceRecordTypeBuilder {}
extension StorageServiceProtoGroupV2RecordBuilder: MigrateableStorageServiceRecordTypeBuilder {}