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

import Foundation
import XCTest
@testable import SignalServiceKit

class AttachmentStoreTests: XCTestCase {

    private var db: InMemoryDB!

    private var attachmentStore: AttachmentStore!

    override func setUp() async throws {
        db = InMemoryDB()
        attachmentStore = AttachmentStore()
    }

    // MARK: - Inserts

    func testInsert() throws {
        var attachmentParams = Attachment.Record.mockPointer()
        let referenceParams = AttachmentReference.ConstructionParams.mock(
            owner: .thread(.globalThreadWallpaperImage(creationTimestamp: Date().ows_millisecondsSince1970)),
        )

        try db.write { tx in
            _ = try attachmentStore.insert(
                &attachmentParams,
                reference: referenceParams,
                tx: tx,
            )
        }

        let (references, attachments) = db.read { tx in
            let references = attachmentStore.fetchReferences(
                owners: [.globalThreadWallpaperImage],
                tx: tx,
            )
            let attachments = attachmentStore.fetch(
                ids: references.map(\.attachmentRowId),
                tx: tx,
            )
            return (references, attachments)
        }

        XCTAssertEqual(references.count, 1)
        XCTAssertEqual(attachments.count, 1)

        let reference = references[0]
        let attachment = attachments[0]

        XCTAssertEqual(reference.attachmentRowId, attachment.id)

        assertEqual(attachmentParams, attachment)
        assertEqual(referenceParams, reference)
    }

    func testMultipleInserts() throws {
        let (threadId1, messageId1) = insertThreadAndInteraction()
        let (threadId2, messageId2) = insertThreadAndInteraction()
        let (threadId3, messageId3) = insertThreadAndInteraction()

        let message1AttachmentIds: [UUID] = [.init()]
        let message2AttachmentIds: [UUID] = [.init(), .init()]
        let message3AttachmentIds: [UUID] = [.init(), .init(), .init()]

        var attachmentIdToAttachmentParams = [UUID: Attachment.Record]()
        var attachmentIdToAttachmentReferenceParams = [UUID: AttachmentReference.ConstructionParams]()

        try db.write { tx in
            for (messageId, threadId, attachmentIds) in [
                (messageId1, threadId1, message1AttachmentIds),
                (messageId2, threadId2, message2AttachmentIds),
                (messageId3, threadId3, message3AttachmentIds),
            ] {
                try attachmentIds.enumerated().forEach { index, id in
                    var attachmentParams = Attachment.Record.mockPointer()
                    let attachmentReferenceParams = AttachmentReference.ConstructionParams.mockMessageBodyAttachmentReference(
                        attachmentRecord: attachmentParams,
                        messageRowId: messageId,
                        threadRowId: threadId,
                        orderInMessage: UInt32(index),
                        idInOwner: id,
                    )
                    try attachmentStore.insert(
                        &attachmentParams,
                        reference: attachmentReferenceParams,
                        tx: tx,
                    )
                    attachmentIdToAttachmentParams[id] = attachmentParams
                    attachmentIdToAttachmentReferenceParams[id] = attachmentReferenceParams
                }
            }
        }

        for (messageId, attachmentIds) in [
            (messageId1, message1AttachmentIds),
            (messageId2, message2AttachmentIds),
            (messageId3, message3AttachmentIds),
        ] {
            let (references, attachments) = db.read { tx in
                let references = attachmentStore.fetchReferences(
                    owners: [.messageBodyAttachment(messageRowId: messageId)],
                    tx: tx,
                )
                let attachments = attachmentStore.fetch(
                    ids: references.map(\.attachmentRowId),
                    tx: tx,
                )
                return (references, attachments)
            }

            XCTAssertEqual(references.count, attachmentIds.count)
            XCTAssertEqual(attachments.count, attachmentIds.count)

            for (reference, attachment) in zip(references, attachments) {
                XCTAssertEqual(reference.attachmentRowId, attachment.id)

                let attachmentId: UUID
                switch reference.owner {
                case .message(.bodyAttachment(let metadata)):
                    attachmentId = metadata.idInOwner!
                default:
                    XCTFail("Unexpected owner type")
                    continue
                }

                guard
                    let attachmentParams = attachmentIdToAttachmentParams[attachmentId],
                    let referenceParams = attachmentIdToAttachmentReferenceParams[attachmentId]
                else {
                    XCTFail("Unexpected attachment id")
                    continue
                }

                assertEqual(attachmentParams, attachment)
                assertEqual(referenceParams, reference)
            }
        }
    }

    // MARK: -

    func testInsertSameMediaName() {
        let mediaName = Data(repeating: 27, count: 10).hexadecimalString

        switch testAttachmentInsertError(
            attachmentParams1: Attachment.Record.mockStream(streamInfo: .mock(mediaName: mediaName)),
            attachmentParams2: Attachment.Record.mockStream(streamInfo: .mock(mediaName: mediaName)),
        ) {
        case .duplicateMediaName:
            break
        case nil, .duplicatePlaintextHash:
            XCTFail()
        }
    }

    func testInsertSamePlaintextHash() throws {
        let plaintextHash = UUID().data

        switch testAttachmentInsertError(
            attachmentParams1: Attachment.Record.mockStream(streamInfo: .mock(plaintextHash: plaintextHash)),
            attachmentParams2: Attachment.Record.mockStream(streamInfo: .mock(plaintextHash: plaintextHash)),
        ) {
        case .duplicatePlaintextHash:
            break
        case nil, .duplicateMediaName:
            XCTFail()
        }
    }

    private func testAttachmentInsertError(
        attachmentParams1: Attachment.Record,
        attachmentParams2: Attachment.Record,
    ) -> AttachmentInsertError? {
        var attachmentParams1 = attachmentParams1
        var attachmentParams2 = attachmentParams2
        var attachmentInsertError: AttachmentInsertError?

        let (threadId1, messageId1) = insertThreadAndInteraction()
        let (threadId2, messageId2) = insertThreadAndInteraction()
        db.write { tx in
            let attachmentReferenceParams1 = AttachmentReference.ConstructionParams.mockMessageBodyAttachmentReference(
                attachmentRecord: attachmentParams1,
                messageRowId: messageId1,
                threadRowId: threadId1,
            )
            try! attachmentStore.insert(
                &attachmentParams1,
                reference: attachmentReferenceParams1,
                tx: tx,
            )

            let message1References = attachmentStore.fetchReferences(
                owners: [.messageBodyAttachment(messageRowId: messageId1)],
                tx: tx,
            )
            let message1Attachment = attachmentStore.fetch(
                ids: message1References.map(\.attachmentRowId),
                tx: tx,
            ).first!

            let attachmentReferenceParams2 = AttachmentReference.ConstructionParams.mockMessageBodyAttachmentReference(
                attachmentRecord: attachmentParams2,
                messageRowId: messageId2,
                threadRowId: threadId2,
            )
            do throws(AttachmentInsertError) {
                try attachmentStore.insert(
                    &attachmentParams2,
                    reference: attachmentReferenceParams2,
                    tx: tx,
                )
                XCTFail("Should have thrown error!")
                attachmentInsertError = nil
            } catch {
                attachmentInsertError = error
            }

            // Try again but insert using explicit owner adding.
            attachmentStore.addReference(
                attachmentReferenceParams2,
                attachmentRowId: message1Attachment.id,
                tx: tx,
            )
        }

        let (
            message1References,
            message1Attachments,
            message2References,
            message2Attachments,
        ) = db.read { tx in
            let message1References = attachmentStore.fetchReferences(
                owners: [.messageBodyAttachment(messageRowId: messageId1)],
                tx: tx,
            )
            let message1Attachments = attachmentStore.fetch(
                ids: message1References.map(\.attachmentRowId),
                tx: tx,
            )
            let message2References = attachmentStore.fetchReferences(
                owners: [.messageBodyAttachment(messageRowId: messageId2)],
                tx: tx,
            )
            let message2Attachments = attachmentStore.fetch(
                ids: message1References.map(\.attachmentRowId),
                tx: tx,
            )
            return (message1References, message1Attachments, message2References, message2Attachments)
        }

        // Both messages should have one reference, to one attachment.
        XCTAssertEqual(message1References.count, 1)
        XCTAssertEqual(message1Attachments.count, 1)
        XCTAssertEqual(message2References.count, 1)
        XCTAssertEqual(message2Attachments.count, 1)

        // But the attachments should be the same!
        XCTAssertEqual(message1Attachments[0].id, message2Attachments[0].id)

        // And it should have used the first attachment inserted.
        XCTAssertEqual(message2Attachments[0].encryptionKey, attachmentParams1.encryptionKey)

        return attachmentInsertError
    }

    // MARK: -

    func testReinsertGlobalThreadAttachment() throws {
        var attachmentParams1 = Attachment.Record.mockPointer()
        let date1 = Date()
        let referenceParams1 = AttachmentReference.ConstructionParams.mock(
            owner: .thread(.globalThreadWallpaperImage(creationTimestamp: date1.ows_millisecondsSince1970)),
        )
        var attachmentParams2 = Attachment.Record.mockPointer()
        let date2 = date1.addingTimeInterval(100)
        let referenceParams2 = AttachmentReference.ConstructionParams.mock(
            owner: .thread(.globalThreadWallpaperImage(creationTimestamp: date2.ows_millisecondsSince1970)),
        )

        try db.write { tx in
            try attachmentStore.insert(
                &attachmentParams1,
                reference: referenceParams1,
                tx: tx,
            )
            // Insert which should overwrite the existing row.
            try attachmentStore.insert(
                &attachmentParams2,
                reference: referenceParams2,
                tx: tx,
            )
        }

        let (references, attachments) = db.read { tx in
            let references = attachmentStore.fetchReferences(
                owners: [.globalThreadWallpaperImage],
                tx: tx,
            )
            let attachments = attachmentStore.fetch(
                ids: references.map(\.attachmentRowId),
                tx: tx,
            )
            return (references, attachments)
        }

        XCTAssertEqual(references.count, 1)
        XCTAssertEqual(attachments.count, 1)

        let reference = references[0]
        let attachment = attachments[0]

        XCTAssertEqual(reference.attachmentRowId, attachment.id)

        assertEqual(attachmentParams2, attachment)
        assertEqual(referenceParams2, reference)
    }

    // MARK: - Enumerate

    func testEnumerateAttachmentReferences() throws {
        let threadIdAndMessageIds = (0..<5).map { _ in
            insertThreadAndInteraction()
        }

        // Insert many references to the same Params over and over.
        var attachmentParams = Attachment.Record.mockStream()

        let attachmentIdsInOwner: [UUID] = try db.write { tx in
            var attachmentRowId: Attachment.IDType?
            return try threadIdAndMessageIds.flatMap { threadId, messageId in
                return try (0..<5).map { index in
                    let attachmentIdInOwner = UUID()
                    let attachmentReferenceParams = AttachmentReference.ConstructionParams.mockMessageBodyAttachmentReference(
                        attachmentRecord: attachmentParams,
                        messageRowId: messageId,
                        threadRowId: threadId,
                        orderInMessage: UInt32(index),
                        idInOwner: attachmentIdInOwner,
                    )
                    if let attachmentRowId {
                        attachmentStore.addReference(
                            attachmentReferenceParams,
                            attachmentRowId: attachmentRowId,
                            tx: tx,
                        )
                    } else {
                        try attachmentStore.insert(
                            &attachmentParams,
                            reference: attachmentReferenceParams,
                            tx: tx,
                        )
                        attachmentRowId = tx.database.lastInsertedRowID
                    }
                    return attachmentIdInOwner
                }
            }
        }

        // Get the attachment id we just inserted.
        let attachmentId = db.read { tx in
            let references = attachmentStore.fetchReferences(
                owners: [.messageBodyAttachment(messageRowId: threadIdAndMessageIds[0].interactionRowId)],
                tx: tx,
            )
            let attachment = attachmentStore.fetch(
                ids: references.map(\.attachmentRowId),
                tx: tx,
            )
            return attachment.first!.id
        }

        // Insert some other references to other arbitrary attachments
        try db.write { tx in
            try threadIdAndMessageIds.forEach { threadId, messageId in
                try (0..<5).forEach { index in
                    var record = Attachment.Record.mockPointer()
                    let attachmentReferenceParams = AttachmentReference.ConstructionParams.mockMessageBodyAttachmentReference(
                        attachmentRecord: record,
                        messageRowId: messageId,
                        threadRowId: threadId,
                        orderInMessage: UInt32(index),
                    )
                    try attachmentStore.insert(
                        &record,
                        reference: attachmentReferenceParams,
                        tx: tx,
                    )
                }
            }
        }

        // Check that we enumerate all the ids we created for the original attachment's id.
        var enumeratedCount = 0
        db.read { tx in
            attachmentStore.enumerateAllReferences(
                toAttachmentId: attachmentId,
                tx: tx,
                block: { reference, _ in
                    enumeratedCount += 1

                    XCTAssertEqual(reference.attachmentRowId, attachmentId)

                    switch reference.owner {
                    case .message(.bodyAttachment(let metadata)):
                        XCTAssertTrue(attachmentIdsInOwner.contains(metadata.idInOwner!))
                    default:
                        XCTFail("Unexpected attachment type!")
                    }
                },
            )
        }

        XCTAssertEqual(enumeratedCount, attachmentIdsInOwner.count)
    }

    func testOldestStickerPackReferences() throws {
        // Sticker pack id to oldest row id.
        var stickerPackIds = [Data: Int64]()

        try db.write { tx in
            let thread = insertThread(tx: tx)
            let threadRowId = thread.sqliteRowId!

            // Add some non sticker attachments.
            for _ in 0..<10 {
                let messageRowId = insertInteraction(thread: thread, tx: tx)

                var record = Attachment.Record.mockStream()
                try attachmentStore.insert(
                    &record,
                    reference: .mock(
                        owner: .message(.bodyAttachment(.init(
                            messageRowId: messageRowId,
                            receivedAtTimestamp: .random(in: 0..<9999999),
                            threadRowId: threadRowId,
                            contentType: .image,
                            isPastEditRevision: false,
                            caption: nil,
                            renderingFlag: .default,
                            orderInMessage: 0,
                            idInOwner: nil,
                            isViewOnce: false,
                        ))),
                    ),
                    tx: tx,
                )
            }

            // Add some sticker attachments
            for _ in 0..<10 {
                let packId = UUID().data
                // Each sticker pack gets multiple stickers, but we should
                // dedupe down to the pack IDs.
                let numStickersInPack = Int.random(in: 1...10)
                for _ in 0..<numStickersInPack {
                    let stickerId = UInt32.random(in: 0..<UInt32.max)

                    let messageRowId = insertInteraction(thread: thread, tx: tx)

                    var record = Attachment.Record.mockStream()
                    try attachmentStore.insert(
                        &record,
                        reference: .mock(
                            owner: .message(.sticker(.init(
                                messageRowId: messageRowId,
                                receivedAtTimestamp: .random(in: 0..<9999999),
                                threadRowId: threadRowId,
                                contentType: .image,
                                isPastEditRevision: false,
                                stickerPackId: packId,
                                stickerId: stickerId,
                            ))),
                        ),
                        tx: tx,
                    )

                    stickerPackIds[packId] = min(messageRowId, stickerPackIds[packId] ?? .max)
                }
            }
        }

        let readPackReferences = db.read { tx in
            attachmentStore.oldestStickerPackReferences(tx: tx)
        }

        XCTAssertEqual(readPackReferences.count, stickerPackIds.count)
        readPackReferences.forEach { packRef in
            XCTAssertEqual(packRef.messageRowId, stickerPackIds[packRef.stickerPackId])
        }
    }

    func testallAttachmentIdsForSticker() throws {
        // Sticker info to count of attachments with that info.
        var stickerInfos = [StickerInfo: Int]()

        try db.write { tx in
            let thread = insertThread(tx: tx)
            let threadRowId = thread.sqliteRowId!

            // Add some non sticker attachments.
            for _ in 0..<10 {
                let messageRowId = insertInteraction(thread: thread, tx: tx)

                var record = Attachment.Record.mockStream()
                try attachmentStore.insert(
                    &record,
                    reference: .mock(
                        owner: .message(.bodyAttachment(.init(
                            messageRowId: messageRowId,
                            receivedAtTimestamp: .random(in: 0..<9999999),
                            threadRowId: threadRowId,
                            contentType: .image,
                            isPastEditRevision: false,
                            caption: nil,
                            renderingFlag: .default,
                            orderInMessage: 0,
                            idInOwner: nil,
                            isViewOnce: false,
                        ))),
                    ),
                    tx: tx,
                )
            }

            // Add some sticker attachments
            for _ in 0..<10 {
                let packId = UUID().data
                // Each sticker pack gets multiple stickers, but we should
                // dedupe down to the pack IDs.
                let numStickersInPack = Int.random(in: 1...10)
                for _ in 0..<numStickersInPack {
                    let stickerId = UInt32.random(in: 0..<UInt32.max)

                    let attachmentCount = Int.random(in: 1...10)
                    stickerInfos[StickerInfo(
                        packId: packId,
                        packKey: Data(repeating: 8, count: Int(StickerManager.packKeyLength)),
                        stickerId: stickerId,
                    )] = attachmentCount
                    for _ in 0..<attachmentCount {

                        let messageRowId = insertInteraction(thread: thread, tx: tx)

                        var record = Attachment.Record.mockStream()
                        try attachmentStore.insert(
                            &record,
                            reference: .mock(
                                owner: .message(.sticker(.init(
                                    messageRowId: messageRowId,
                                    receivedAtTimestamp: .random(in: 0..<9999999),
                                    threadRowId: threadRowId,
                                    contentType: .image,
                                    isPastEditRevision: false,
                                    stickerPackId: packId,
                                    stickerId: stickerId,
                                ))),
                            ),
                            tx: tx,
                        )
                    }

                }
            }
        }

        db.read { tx in
            for (stickerInfo, expectedCount) in stickerInfos {
                XCTAssertEqual(
                    expectedCount,
                    attachmentStore
                        .allAttachmentIdsForSticker(stickerInfo, tx: tx)
                        .count,
                )
            }
        }

    }

    // MARK: - Update

    func testMarkAsUploaded() throws {
        let (threadId, messageId) = insertThreadAndInteraction()

        try db.write { tx in
            var attachmentParams = Attachment.Record.mockStream()
            let attachmentReferenceParams = AttachmentReference.ConstructionParams.mockMessageBodyAttachmentReference(
                attachmentRecord: attachmentParams,
                messageRowId: messageId,
                threadRowId: threadId,
            )
            try attachmentStore.insert(
                &attachmentParams,
                reference: attachmentReferenceParams,
                tx: tx,
            )
        }

        func fetchAttachment() -> Attachment {
            return db.read { tx in
                let references = attachmentStore.fetchReferences(
                    owners: [.messageBodyAttachment(messageRowId: messageId)],
                    tx: tx,
                )
                return attachmentStore.fetch(
                    ids: references.map(\.attachmentRowId),
                    tx: tx,
                ).first!
            }
        }

        var attachment = fetchAttachment()

        guard let stream = attachment.asStream() else {
            XCTFail("Expected attachment stream!")
            return
        }
        XCTAssertNil(attachment.latestTransitTierInfo)

        let transitTierInfo = Attachment.TransitTierInfo(
            cdnNumber: 3,
            cdnKey: UUID().uuidString,
            uploadTimestamp: Date().ows_millisecondsSince1970,
            encryptionKey: UUID().data,
            unencryptedByteCount: 100,
            integrityCheck: .ciphertextDigest(UUID().data),
            incrementalMacInfo: nil,
            lastDownloadAttemptTimestamp: nil,
        )

        // Mark it as uploaded.
        db.write { tx in
            attachmentStore.saveLatestTransitTierInfo(
                attachmentStream: stream,
                transitTierInfo: transitTierInfo,
                tx: tx,
            )
        }

        // Refetch and check that it is appropriately marked.
        attachment = fetchAttachment()

        XCTAssertEqual(attachment.latestTransitTierInfo, transitTierInfo)
    }

    // MARK: - Remove Owner

    func testRemoveOwner() throws {
        let (threadId1, messageId1) = insertThreadAndInteraction()
        let (threadId2, messageId2) = insertThreadAndInteraction()

        // Create two references to the same attachment.
        var attachmentParams = Attachment.Record.mockStream()

        try db.write { tx in
            try attachmentStore.insert(
                &attachmentParams,
                reference: AttachmentReference.ConstructionParams.mockMessageBodyAttachmentReference(
                    attachmentRecord: attachmentParams,
                    messageRowId: messageId1,
                    threadRowId: threadId1,
                ),
                tx: tx,
            )
            let attachmentId = tx.database.lastInsertedRowID
            attachmentStore.addReference(
                AttachmentReference.ConstructionParams.mockMessageBodyAttachmentReference(
                    attachmentRecord: attachmentParams,
                    messageRowId: messageId2,
                    threadRowId: threadId2,
                ),
                attachmentRowId: attachmentId,
                tx: tx,
            )
        }

        let (reference1, reference2) = db.read { tx in
            let reference1 = attachmentStore.fetchReferences(
                owners: [.messageBodyAttachment(messageRowId: messageId1)],
                tx: tx,
            ).first!
            let reference2 = attachmentStore.fetchReferences(
                owners: [.messageBodyAttachment(messageRowId: messageId2)],
                tx: tx,
            ).first!
            return (reference1, reference2)
        }

        XCTAssertEqual(reference1.attachmentRowId, reference2.attachmentRowId)

        db.write { tx in
            // Remove the first reference.
            attachmentStore.removeReference(
                reference: reference1,
                tx: tx,
            )

            // The attachment should still exist.
            XCTAssertNotNil(attachmentStore.fetch(
                ids: [reference1.attachmentRowId],
                tx: tx,
            ).first)

            // Remove the second reference.
            attachmentStore.removeReference(
                reference: reference2,
                tx: tx,
            )

            // The attachment should no longer exist.
            XCTAssertNil(attachmentStore.fetch(
                ids: [reference1.attachmentRowId],
                tx: tx,
            ).first)
        }
    }

    func testRemoveOwner_sameOwnerMultipleReferences() throws {
        let (threadId1, messageId1) = insertThreadAndInteraction()

        let referenceUUID1 = UUID()
        let referenceUUID2 = UUID()

        // Create two references to the same attachment on the same message.
        var attachmentParams = Attachment.Record.mockStream()

        try db.write { tx in
            try attachmentStore.insert(
                &attachmentParams,
                reference: AttachmentReference.ConstructionParams.mockMessageBodyAttachmentReference(
                    attachmentRecord: attachmentParams,
                    messageRowId: messageId1,
                    threadRowId: threadId1,
                    orderInMessage: 0,
                    idInOwner: referenceUUID1,
                ),
                tx: tx,
            )
            let attachmentId = tx.database.lastInsertedRowID
            attachmentStore.addReference(
                AttachmentReference.ConstructionParams.mockMessageBodyAttachmentReference(
                    attachmentRecord: attachmentParams,
                    messageRowId: messageId1,
                    threadRowId: threadId1,
                    orderInMessage: 1,
                    idInOwner: referenceUUID2,
                ),
                attachmentRowId: attachmentId,
                tx: tx,
            )
        }

        let (reference1, reference2) = db.read { tx in
            let references = attachmentStore.fetchReferences(
                owners: [.messageBodyAttachment(messageRowId: messageId1)],
                tx: tx,
            ).sorted(by: {
                $0.orderInOwningMessage! < $1.orderInOwningMessage!
            })
            return (references[0], references[1])
        }

        func checkIdInOwner(of reference: AttachmentReference, matches uuid: UUID) {
            switch reference.owner {
            case .message(let messageSource):
                switch messageSource {
                case .bodyAttachment(let metadata):
                    XCTAssertEqual(metadata.idInOwner, uuid)
                default:
                    XCTFail("Unexpected reference type")
                }
            default:
                XCTFail("Unexpected reference type")
            }
        }

        checkIdInOwner(of: reference1, matches: referenceUUID1)
        checkIdInOwner(of: reference2, matches: referenceUUID2)

        db.write { tx in
            // Remove the first reference.
            attachmentStore.removeReference(
                reference: reference1,
                tx: tx,
            )

            // The attachment should still exist.
            XCTAssertNotNil(attachmentStore.fetch(
                ids: [reference1.attachmentRowId],
                tx: tx,
            ).first)

            // Refetch references
            let references = attachmentStore.fetchReferences(
                owners: [.messageBodyAttachment(messageRowId: messageId1)],
                tx: tx,
            )
            XCTAssertEqual(references.count, 1)
            checkIdInOwner(of: references[0], matches: referenceUUID2)
        }
    }

    // MARK: Remove all thread owners

    func testRemoveAllThreadOwners() throws {
        var threadRowIds: [Int64?] = [nil]
        db.write { tx in
            for _ in 0..<5 {
                threadRowIds.append(self.insertThread(tx: tx).sqliteRowId!)
            }
        }

        try db.write { tx in
            try threadRowIds.forEach { threadRowId in
                var attachmentParams = Attachment.Record.mockPointer()
                let timestamp = Date().ows_millisecondsSince1970
                let referenceParams = AttachmentReference.ConstructionParams.mock(
                    owner: .thread(threadRowId.map {
                        .threadWallpaperImage(.init(threadRowId: $0, creationTimestamp: timestamp))
                    } ?? .globalThreadWallpaperImage(creationTimestamp: timestamp)),
                )
                try attachmentStore.insert(
                    &attachmentParams,
                    reference: referenceParams,
                    tx: tx,
                )
            }
        }

        func assertAttachmentCount(_ expectedCount: Int) {
            db.read { tx in
                let references = attachmentStore.fetchReferences(
                    owners: threadRowIds.map { threadRowId in
                        threadRowId.map {
                            .threadWallpaperImage(threadRowId: $0)
                        } ?? .globalThreadWallpaperImage
                    },
                    tx: tx,
                )
                let attachments = attachmentStore.fetch(
                    ids: references.map(\.attachmentRowId),
                    tx: tx,
                )
                XCTAssertEqual(references.count, expectedCount)
                XCTAssertEqual(attachments.count, expectedCount)
            }
        }

        assertAttachmentCount(threadRowIds.count)

        // Remove all and count should be 0.

        db.write { tx in
            attachmentStore.removeAllThreadOwners(
                tx: tx,
            )
        }

        assertAttachmentCount(0)
    }

    // MARK: - Thread merging

    func testThreadMerging() throws {
        var threadId1: Int64!
        var threadId2: Int64!
        var threadId3: Int64!
        var messageId1: Int64!
        var messageId2: Int64!
        var messageId3: Int64!
        db.write { tx in
            let thread1 = insertThread(tx: tx)
            threadId1 = thread1.sqliteRowId!
            messageId1 = insertInteraction(thread: thread1, tx: tx)
            messageId2 = insertInteraction(thread: thread1, tx: tx)
            let thread2 = insertThread(tx: tx)
            threadId2 = thread2.sqliteRowId!
            messageId3 = insertInteraction(thread: thread2, tx: tx)
            let thread3 = insertThread(tx: tx)
            threadId3 = thread3.sqliteRowId!
        }

        try db.write { tx in
            for (threadRowId, messageRowId) in [
                (threadId1, messageId1),
                (threadId1, messageId2),
                (threadId2, messageId3),
            ] {
                // Create an attachment and reference for each message.
                var attachmentParams = Attachment.Record.mockStream()
                try attachmentStore.insert(
                    &attachmentParams,
                    reference: AttachmentReference.ConstructionParams.mockMessageBodyAttachmentReference(
                        attachmentRecord: attachmentParams,
                        messageRowId: messageRowId!,
                        threadRowId: threadRowId!,
                    ),
                    tx: tx,
                )
            }
        }

        var (reference1, reference2, reference3) = db.read { tx in
            let reference1 = attachmentStore.fetchReferences(
                owners: [.messageBodyAttachment(messageRowId: messageId1)],
                tx: tx,
            ).first!
            let reference2 = attachmentStore.fetchReferences(
                owners: [.messageBodyAttachment(messageRowId: messageId2)],
                tx: tx,
            ).first!
            let reference3 = attachmentStore.fetchReferences(
                owners: [.messageBodyAttachment(messageRowId: messageId3)],
                tx: tx,
            ).first!
            return (reference1, reference2, reference3)
        }

        func checkThreadRowId(of reference: AttachmentReference, matches threadId: Int64) {
            switch reference.owner {
            case .message(let messageSource):
                switch messageSource {
                case .bodyAttachment(let metadata):
                    XCTAssertEqual(metadata.threadRowId, threadId)
                case .oversizeText(let metadata):
                    XCTAssertEqual(metadata.threadRowId, threadId)
                case .linkPreview(let metadata):
                    XCTAssertEqual(metadata.threadRowId, threadId)
                case .quotedReply(let metadata):
                    XCTAssertEqual(metadata.threadRowId, threadId)
                case .sticker(let metadata):
                    XCTAssertEqual(metadata.threadRowId, threadId)
                case .contactAvatar(let metadata):
                    XCTAssertEqual(metadata.threadRowId, threadId)
                }
            default:
                XCTFail("Unexpected reference type")
            }
        }

        checkThreadRowId(of: reference1, matches: threadId1)
        checkThreadRowId(of: reference2, matches: threadId1)
        checkThreadRowId(of: reference3, matches: threadId2)

        db.write { tx in
            // Merge threads
            attachmentStore.updateMessageAttachmentThreadRowIdsForThreadMerge(
                fromThreadRowId: threadId1,
                intoThreadRowId: threadId3,
                tx: tx,
            )
        }

        (reference1, reference2, reference3) = db.read { tx in
            let reference1 = attachmentStore.fetchReferences(
                owners: [.messageBodyAttachment(messageRowId: messageId1)],
                tx: tx,
            ).first!
            let reference2 = attachmentStore.fetchReferences(
                owners: [.messageBodyAttachment(messageRowId: messageId2)],
                tx: tx,
            ).first!
            let reference3 = attachmentStore.fetchReferences(
                owners: [.messageBodyAttachment(messageRowId: messageId3)],
                tx: tx,
            ).first!
            return (reference1, reference2, reference3)
        }

        // The ones that were thread 1 are now thread 3.
        checkThreadRowId(of: reference1, matches: threadId3)
        checkThreadRowId(of: reference2, matches: threadId3)
        checkThreadRowId(of: reference3, matches: threadId2)
    }

    // MARK: - UInt64 max values

    /// Added as part of the introduction of `DBUInt64` and its first usage in
    /// `Attachment.Record`. We don't need to implement this for every
    /// subsequent type that uses `DBUInt64`.
    func testUInt64MaxValues() throws {
        try InMemoryDB().write { tx in
            let attachment = Attachment.mock(
                transitTierInfo: .mock(
                    uploadTimestamp: .max,
                    lastDownloadAttemptTimestamp: .max,
                ),
                mediaTierInfo: .mock(
                    lastDownloadAttemptTimestamp: .max,
                ),
                thumbnailInfo: .mock(
                    lastDownloadAttemptTimestamp: .max,
                ),
                lastFullscreenViewTimestamp: .max,
            )

            var attachmentRecord = Attachment.Record(attachment: attachment)
            try attachmentRecord.insert(tx.database)
            let persisted = try Attachment.Record.fetchAll(tx.database)
            XCTAssertEqual(persisted.count, 1)
            XCTAssertEqual(persisted[0], attachmentRecord)
        }
    }

    // MARK: - Helpers

    private func insertThreadAndInteraction() -> (threadRowId: Int64, interactionRowId: Int64) {
        return db.write { tx in
            let thread = insertThread(tx: tx)
            let interactionRowId = insertInteraction(thread: thread, tx: tx)
            return (thread.sqliteRowId!, interactionRowId)
        }
    }

    private func insertThread(tx: DBWriteTransaction) -> TSThread {
        let thread = TSThread(uniqueId: UUID().uuidString)
        try! thread.insert(tx.database)
        return thread
    }

    private func insertInteraction(thread: TSThread, tx: DBWriteTransaction) -> Int64 {
        let interaction = TSInteraction(timestamp: 0, receivedAtTimestamp: 0, thread: thread)
        try! interaction.asRecord().insert(tx.database)
        return interaction.sqliteRowId!
    }

    private func assertEqual(_ record: Attachment.Record, _ attachment: Attachment) {
        var record = record
        record.sqliteId = attachment.id
        XCTAssertEqual(record, .init(attachment: attachment))
    }

    private func assertEqual(_ params: AttachmentReference.ConstructionParams, _ reference: AttachmentReference) {
        switch (params.owner, reference.owner) {
        case (.message, .message(let messageSource)):
            let paramsRecord = AttachmentReference.MessageAttachmentReferenceRecord(
                attachmentRowId: reference.attachmentRowId,
                sourceFilename: params.sourceFilename,
                sourceUnencryptedByteCount: params.sourceUnencryptedByteCount,
                sourceMediaSizePixels: params.sourceMediaSizePixels,
                messageSource: messageSource,
            )
            let referenceRecord = AttachmentReference.MessageAttachmentReferenceRecord(
                attachmentReference: reference,
                messageSource: messageSource,
            )
            XCTAssertEqual(paramsRecord, referenceRecord)
        case (.storyMessage, .storyMessage(let storyMessageSource)):
            let paramsRecord = AttachmentReference.StoryMessageAttachmentReferenceRecord(
                attachmentRowId: reference.attachmentRowId,
                sourceFilename: params.sourceFilename,
                sourceUnencryptedByteCount: params.sourceUnencryptedByteCount,
                sourceMediaSizePixels: params.sourceMediaSizePixels,
                storyMessageSource: storyMessageSource,
            )
            let referenceRecord = AttachmentReference.StoryMessageAttachmentReferenceRecord(
                attachmentReference: reference,
                storyMessageSource: storyMessageSource,
            )
            XCTAssertEqual(paramsRecord, referenceRecord)
        case (.thread, .thread(let threadSource)):
            let paramsRecord = AttachmentReference.ThreadAttachmentReferenceRecord(
                attachmentRowId: reference.attachmentRowId,
                threadSource: threadSource,
            )
            let referenceRecord = AttachmentReference.ThreadAttachmentReferenceRecord(
                attachmentRowId: reference.attachmentRowId,
                threadSource: threadSource,
            )
            XCTAssertEqual(paramsRecord, referenceRecord)
        case (.message, _), (.storyMessage, _), (.thread, _):
            XCTFail("Non matching owner types")
        }
    }
}
