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

import Foundation
import GRDB
import XCTest
@testable import SignalServiceKit

class MediaGalleryAttachmentFinderTest: XCTestCase {

    private let attachmentStore = AttachmentStore()
    private var db: InMemoryDB!

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

    // MARK: - Queries

    func testQueryDateRange() throws {
        let (thread, messageRowId) = insertThreadAndInteraction()
        let threadRowId = thread.sqliteRowId!

        // Insert one matching content type before the date range
        insertAttachment(
            messageRowId: messageRowId,
            threadRowId: threadRowId,
            receivedAtTimestamp: 100,
            mimeType: "image/jpeg",
            orderInMessage: 0,
        )
        // ...and one matching content type after the date range
        insertAttachment(
            messageRowId: messageRowId,
            threadRowId: threadRowId,
            receivedAtTimestamp: 300,
            mimeType: "image/jpeg",
            orderInMessage: 1,
        )
        // ...and one non-matching content type within the date range
        insertAttachment(
            messageRowId: messageRowId,
            threadRowId: threadRowId,
            receivedAtTimestamp: 200,
            mimeType: "audio/mp3",
            orderInMessage: 2,
        )
        // ...and two matching content type within the date range
        insertAttachment(
            messageRowId: messageRowId,
            threadRowId: threadRowId,
            receivedAtTimestamp: 200,
            mimeType: "image/jpeg",
            orderInMessage: 3,
        )
        insertAttachment(
            messageRowId: messageRowId,
            threadRowId: threadRowId,
            receivedAtTimestamp: 200,
            mimeType: "image/jpeg",
            orderInMessage: 4,
        )
        // ...and one within the date range that we will exclude.
        insertAttachment(
            messageRowId: messageRowId,
            threadRowId: threadRowId,
            receivedAtTimestamp: 200,
            mimeType: "image/jpeg",
            orderInMessage: 5,
        )
        // ...and a view once attachment that will be excluded.
        insertAttachment(
            messageRowId: messageRowId,
            threadRowId: threadRowId,
            receivedAtTimestamp: 200,
            mimeType: "image/jpeg",
            orderInMessage: 6,
            isViewOnce: true,
        )
        let exclusionSet = Set<AttachmentReferenceId>([
            .init(ownerId: .messageBodyAttachment(messageRowId: messageRowId), orderInMessage: 5),
        ])

        let finder = MediaGalleryAttachmentFinder(threadId: thread.grdbId!.int64Value, filter: .allPhotoVideoCategory)

        // Should get two results with offset 0
        var query = finder.galleryItemQuery(
            in: .init(
                start: .init(millisecondsSince1970: 150),
                end: .init(millisecondsSince1970: 250),
            ),
            excluding: exclusionSet,
            offset: 0,
            ascending: true,
        )

        var results = try db.read { tx in
            return try query.fetchAll(tx.database)
        }

        XCTAssertEqual(results.count, 2)
        XCTAssertEqual(results[0].receivedAtTimestamp, 200)
        XCTAssertEqual(results[0].orderInMessage, 3)
        XCTAssertEqual(results[1].receivedAtTimestamp, 200)
        XCTAssertEqual(results[1].orderInMessage, 4)

        // Should get just the second result with offset 1
        query = finder.galleryItemQuery(
            in: .init(
                start: .init(millisecondsSince1970: 150),
                end: .init(millisecondsSince1970: 250),
            ),
            excluding: exclusionSet,
            offset: 1,
            ascending: true,
        )

        results = try db.read { tx in
            return try query.fetchAll(tx.database)
        }

        XCTAssertEqual(results.count, 1)
        XCTAssertEqual(results[0].receivedAtTimestamp, 200)
        XCTAssertEqual(results[0].orderInMessage, 4)
    }

    // MARK: - Index Usage

    func testAllQueriesUseIndex() throws {
        DebugFlags.mediaGalleryOnlyAllowDownloaded.set(false)

        let (thread, _) = insertThreadAndInteraction()

        // Set up some parametrized values for tests.
        // Specific values for many things don't matter, just presence
        // and combinations thereof.
        let dateIntervals: [DateInterval?] = [
            nil,
            .init(
                start: .init(millisecondsSince1970: 100),
                end: .init(millisecondsSince1970: 200),
            ),
        ]
        let exclusionSets: [Set<AttachmentReferenceId>] = [
            Set(),
            Set([.init(ownerId: .messageBodyAttachment(messageRowId: 100), orderInMessage: nil)]),
            Set([.init(ownerId: .messageBodyAttachment(messageRowId: 200), orderInMessage: 5)]),
            Set([
                .init(ownerId: .messageBodyAttachment(messageRowId: 100), orderInMessage: nil),
                .init(ownerId: .messageBodyAttachment(messageRowId: 200), orderInMessage: nil),
                .init(ownerId: .messageBodyAttachment(messageRowId: 300), orderInMessage: 5),
            ]),
        ]
        let offsets: [Int] = [0, 5]
        let limits: [Int] = [5, 100]
        let ascendings: [Bool] = [true, false]

        for filter in AllMediaFilter.allCases {
            let finder = MediaGalleryAttachmentFinder(threadId: thread.grdbId!.int64Value, filter: filter)
            var queries = [QueryInterfaceRequest<RecordType>]()

            for dateInterval in dateIntervals {
                for exclusionSet in exclusionSets {
                    for offset in offsets {
                        for ascending in ascendings {
                            queries.append(finder.galleryItemQuery(
                                in: dateInterval,
                                excluding: exclusionSet,
                                offset: offset,
                                ascending: ascending,
                            ))
                        }
                    }
                }
            }
            for dateInterval in dateIntervals.compacted() {
                for exclusionSet in exclusionSets {
                    for offset in offsets {
                        for limit in limits {
                            queries.append(finder.enumerateMediaAttachmentsQuery(
                                in: dateInterval,
                                excluding: exclusionSet,
                                range: .init(location: offset, length: limit),
                            ))
                        }
                    }
                }
            }
            for dateInterval in dateIntervals {
                for exclusionSet in exclusionSets {
                    for limit in limits {
                        queries.append(finder.enumerateTimestampsQuery(
                            beforeDate: dateInterval?.end,
                            afterDate: nil,
                            excluding: exclusionSet,
                            count: limit,
                            ascending: false,
                        ))
                        queries.append(finder.enumerateTimestampsQuery(
                            beforeDate: nil,
                            afterDate: dateInterval?.start,
                            excluding: exclusionSet,
                            count: limit,
                            ascending: true,
                        ))
                    }
                }
            }
            for limit in limits {
                queries.append(finder.recentMediaAttachmentsQuery(limit: limit))
            }

            try db.read { tx in
                for query in queries {
                    let preparedStatement = try query.makePreparedRequest(tx.database).statement
                    let queryPlan: [String] = try Row.fetchAll(
                        tx.database,
                        sql: "EXPLAIN QUERY PLAN \(preparedStatement.sql);",
                        arguments: preparedStatement.arguments,
                    ).map { $0["detail"] }

                    // Ensure we use the relevant indexes and...
                    // * we use all the columns up to the ordering columns
                    // * we DONT use expensive B trees for ordering
                    let allowedQueryPlans: [String] = [
                        "SEARCH MessageAttachmentReference USING INDEX message_attachment_reference_media_gallery_single_content_type_index (threadRowId=? AND ownerType=? AND contentType=?",
                        "SEARCH MessageAttachmentReference USING INDEX message_attachment_reference_media_gallery_visualMedia_content_type_index (threadRowId=? AND ownerType=? AND isVisualMediaContentType=?",
                        "SEARCH MessageAttachmentReference USING INDEX message_attachment_reference_media_gallery_fileOrInvalid_content_type_index (threadRowId=? AND ownerType=? AND isInvalidOrFileContentType=?",
                    ]
                    XCTAssert(queryPlan.allSatisfy { queryPlan in
                        for allowedQueryPlan in allowedQueryPlans {
                            if queryPlan.hasPrefix(allowedQueryPlan) {
                                return true
                            }
                        }
                        return false
                    })
                    // There should NOT be expensive B-TREE usage.
                    XCTAssert(queryPlan.allSatisfy { $0.contains("USE TEMP B-TREE").negated })
                }
            }
        }
    }

    // MARK: - Helpers

    typealias RecordType = MediaGalleryAttachmentFinder.RecordType

    private func insertThreadAndInteraction() -> (thread: TSThread, interactionRowId: Int64) {
        let thread = TSThread(uniqueId: UUID().uuidString)
        let interaction = TSInteraction(timestamp: 0, receivedAtTimestamp: 0, thread: thread)

        db.write { tx in
            try! thread.insert(tx.database)
            try! interaction.asRecord().insert(tx.database)
        }

        return (thread, interaction.sqliteRowId!)
    }

    @discardableResult
    private func insertAttachment(
        messageRowId: Int64,
        threadRowId: Int64,
        receivedAtTimestamp: UInt64,
        mimeType: String,
        orderInMessage: UInt32,
        isViewOnce: Bool = false,
    ) -> Attachment.IDType {
        db.write { tx in
            var attachmentRecord = Attachment.Record.mockStream(
                mimeType: mimeType,
            )
            try! attachmentRecord.insert(tx.database)

            let attachment = Attachment(record: attachmentRecord)

            let referenceParams = AttachmentReference.ConstructionParams.mock(
                owner: .message(.bodyAttachment(.init(
                    messageRowId: messageRowId,
                    receivedAtTimestamp: receivedAtTimestamp,
                    threadRowId: threadRowId,
                    contentType: attachment.contentType,
                    isPastEditRevision: false,
                    caption: nil,
                    renderingFlag: .default,
                    orderInMessage: orderInMessage,
                    idInOwner: nil,
                    isViewOnce: isViewOnce,
                ))),
            )

            attachmentStore.addReference(
                referenceParams,
                attachmentRowId: attachment.id,
                tx: tx,
            )

            return attachment.id
        }
    }
}
