Luma/Sources/Shared/Persistence/ChatArchive.swift

import Foundation
import SwiftData

/// Durable per-account chat history.
///
/// The in-memory domain model stays the value types `Conversation` and
/// `ChatMessage`; this type is only the persistence boundary. Storage is
/// SwiftData (`@Model` records in `ArchiveRecords.swift`) instead of the
/// former per-account JSON snapshot file. A one-shot migration reads the legacy
/// JSON file on first launch and seeds the SwiftData store.
final class ChatArchive {
    struct Snapshot: Codable, Sendable {
        static let currentSchemaVersion = 6

        var schemaVersion: Int
        var conversations: [Conversation]
        var messages: [ChatMessage]
        var locallyDeletedMessageIDs: Set<String>
        var rosterContactJIDs: Set<String>
        var lastSuccessfulMAMSync: Date?
        var lastSuccessfulMAMCursor: String?
        var mamCheckpoints: [MAMArchiveKey: MAMArchiveCheckpoint]

        init(
            conversations: [Conversation],
            messages: [ChatMessage],
            locallyDeletedMessageIDs: Set<String> = [],
            rosterContactJIDs: Set<String> = [],
            lastSuccessfulMAMSync: Date? = nil,
            lastSuccessfulMAMCursor: String? = nil,
            mamCheckpoints: [MAMArchiveKey: MAMArchiveCheckpoint] = [:]
        ) {
            self.schemaVersion = Self.currentSchemaVersion
            self.conversations = conversations
            self.messages = messages
            self.locallyDeletedMessageIDs = locallyDeletedMessageIDs
            self.rosterContactJIDs = rosterContactJIDs
            self.lastSuccessfulMAMSync = lastSuccessfulMAMSync
            self.lastSuccessfulMAMCursor = ArchiveSyncCheckpoint.normalizedCursor(
                lastSuccessfulMAMCursor
            )
            self.mamCheckpoints = mamCheckpoints
        }

        private enum CodingKeys: String, CodingKey {
            case schemaVersion
            case conversations
            case messages
            case locallyDeletedMessageIDs
            case rosterContactJIDs
            case lastSuccessfulMAMSync
            case lastSuccessfulMAMCursor
            case mamCheckpoints
        }

        init(from decoder: Decoder) throws {
            let values = try decoder.container(keyedBy: CodingKeys.self)
            schemaVersion = try values.decodeIfPresent(Int.self, forKey: .schemaVersion) ?? 1
            conversations = try values.decode([Conversation].self, forKey: .conversations)
            messages = try values.decode([ChatMessage].self, forKey: .messages)
            locallyDeletedMessageIDs = try values.decodeIfPresent(
                Set<String>.self,
                forKey: .locallyDeletedMessageIDs
            ) ?? []
            rosterContactJIDs = try values.decodeIfPresent(
                Set<String>.self,
                forKey: .rosterContactJIDs
            ) ?? []
            lastSuccessfulMAMSync = try values.decodeIfPresent(
                Date.self,
                forKey: .lastSuccessfulMAMSync
            )
            lastSuccessfulMAMCursor = ArchiveSyncCheckpoint.normalizedCursor(
                try values.decodeIfPresent(String.self, forKey: .lastSuccessfulMAMCursor)
            )
            mamCheckpoints = try values.decodeIfPresent(
                            [MAMArchiveKey: MAMArchiveCheckpoint].self, forKey: .mamCheckpoints
                        ) ?? [:]
        }

        static let empty = Snapshot(conversations: [], messages: [])
    }

    private let accountJID: String
    private let legacyFileURL: URL

    init(accountJID: String) {
        self.accountJID = accountJID.lowercased()
        let fileManager = FileManager.default
        let base = fileManager.urls(for: .applicationSupportDirectory, in: .userDomainMask).first
            ?? fileManager.temporaryDirectory
        let directory = base.appendingPathComponent("Luma/Accounts", isDirectory: true)
        legacyFileURL = directory.appendingPathComponent("\(Self.stableHash(accountJID)).json")
    }

    func load() async -> Snapshot {
        let jid = accountJID
        let legacyURL = legacyFileURL
        return await MainActor.run {
            let context = LumaPersistence.archiveContainer.mainContext

            guard
                let conversations = try? context.fetch(FetchDescriptor<ConversationRecord>()),
                let messages = try? context.fetch(FetchDescriptor<MessageRecord>()),
                let metadata = try? context.fetch(FetchDescriptor<ArchiveMetadataRecord>())
            else {
                return Snapshot.empty
            }

            let accountConversations = conversations.filter { $0.accountJID == jid }
            let accountMessages = messages.filter { $0.accountJID == jid }
            let accountMetadata = metadata.filter { $0.accountJID == jid }

            if !accountConversations.isEmpty || !accountMessages.isEmpty || !accountMetadata.isEmpty {
                return Self.snapshot(
                    conversations: accountConversations,
                    messages: accountMessages,
                    metadata: accountMetadata.first
                )
            }

            // The SwiftData store is empty for this account: migrate the legacy
            // JSON snapshot once, if it still exists.
            if let legacy = Self.decodeLegacySnapshot(from: legacyURL) {
                try? Self.replaceAll(with: legacy, accountJID: jid, context: context)
                try? context.save()
                try? FileManager.default.removeItem(at: legacyURL)
                return legacy
            }

            return .empty
        }
    }

    func save(_ snapshot: Snapshot) async throws {
        let jid = accountJID
        let legacyURL = legacyFileURL
        try await MainActor.run {
            let context = LumaPersistence.archiveContainer.mainContext
            try Self.replaceAll(with: snapshot, accountJID: jid, context: context)
            try context.save()
            // The legacy file is stale once SwiftData owns the history.
            try? FileManager.default.removeItem(at: legacyURL)
        }
    }

    func erase() async throws {
        let jid = accountJID
        let legacyURL = legacyFileURL
        try await MainActor.run {
            let context = LumaPersistence.archiveContainer.mainContext

            for record in try context.fetch(FetchDescriptor<ConversationRecord>())
                where record.accountJID == jid {
                context.delete(record)
            }
            for record in try context.fetch(FetchDescriptor<MessageRecord>())
                where record.accountJID == jid {
                context.delete(record)
            }
            for record in try context.fetch(FetchDescriptor<ArchiveMetadataRecord>())
                where record.accountJID == jid {
                context.delete(record)
            }
            try context.save()
            try? FileManager.default.removeItem(at: legacyURL)
        }
    }

    // MARK: - Helpers

    private static func snapshot(
        conversations: [ConversationRecord],
        messages: [MessageRecord],
        metadata: ArchiveMetadataRecord?
    ) -> Snapshot {
        Snapshot(
            conversations: conversations.map { $0.value },
            messages: messages.map { $0.value },
            locallyDeletedMessageIDs: metadata.flatMap {
                ArchiveBlobCodec.decode([String].self, from: $0.locallyDeletedMessageIDsData).map { Set($0) }
            } ?? [],
            rosterContactJIDs: metadata.flatMap {
                ArchiveBlobCodec.decode([String].self, from: $0.rosterContactJIDsData).map { Set($0) }
            } ?? [],
            lastSuccessfulMAMSync: metadata?.lastSuccessfulMAMSync,
            lastSuccessfulMAMCursor: metadata?.lastSuccessfulMAMCursor,
            mamCheckpoints: metadata.flatMap {
                ArchiveBlobCodec.decode([MAMArchiveKey: MAMArchiveCheckpoint].self, from: $0.mamCheckpointsData)
            } ?? [:]
        )
    }

    private static func replaceAll(
        with snapshot: Snapshot,
        accountJID: String,
        context: ModelContext
    ) throws {
        for record in try context.fetch(FetchDescriptor<ConversationRecord>())
            where record.accountJID == accountJID {
            context.delete(record)
        }
        for record in try context.fetch(FetchDescriptor<MessageRecord>())
            where record.accountJID == accountJID {
            context.delete(record)
        }
        for record in try context.fetch(FetchDescriptor<ArchiveMetadataRecord>())
            where record.accountJID == accountJID {
            context.delete(record)
        }

        for conversation in snapshot.conversations {
            context.insert(ConversationRecord(accountJID: accountJID, conversation: conversation))
        }
        for message in snapshot.messages {
            context.insert(MessageRecord(accountJID: accountJID, message: message))
        }
        context.insert(ArchiveMetadataRecord(accountJID: accountJID, snapshot: snapshot))
    }

    private static func decodeLegacySnapshot(from url: URL) -> Snapshot? {
        guard let data = try? Data(contentsOf: url) else { return nil }
        let decoder = JSONDecoder()
        decoder.dateDecodingStrategy = .millisecondsSince1970
        return try? decoder.decode(Snapshot.self, from: data)
    }

    private static func stableHash(_ value: String) -> String {
        let hash = value.lowercased().utf8.reduce(UInt64(14_695_981_039_346_656_037)) { partial, byte in
            (partial ^ UInt64(byte)) &* 1_099_511_628_211
        }
        return String(hash, radix: 16, uppercase: false)
    }
}