//
// Copyright Amazon.com Inc. or its affiliates.
// All Rights Reserved.
//
// SPDX-License-Identifier: Apache-2.0
//

import Combine
import Foundation
import XCTest

@testable import Amplify
@testable import AmplifyTestCommon
@testable import AWSDataStorePlugin
@testable import AWSPluginsCore

class ObserveQueryTaskRunnerTests: XCTestCase {
    var storageEngine: MockStorageEngineBehavior!
    var dataStorePublisher: ModelSubcriptionBehavior!
    var dataStoreStateSubject = PassthroughSubject<DataStoreState, DataStoreError>()

    override func setUp() {
        ModelRegistry.register(modelType: Post.self)
        storageEngine = MockStorageEngineBehavior()
        dataStorePublisher = DataStorePublisher()
    }

    ///  An item changed observed will be returned in a single snapshot
    ///
    /// - Given:  The operation has started and the first query has completed.
    /// - When:
    ///    -  Observe an item change.
    /// - Then:
    ///    - The item observed will be returned in the second snapshot
    ///
    func testItemChangedWillGenerateSnapshot() async throws {
        let firstSnapshot = expectation(description: "first query snapshots")
        let secondSnapshot = expectation(description: "second query snapshots")
        let taskRunner = ObserveQueryTaskRunner(
            modelType: Post.self,
            modelSchema: Post.schema,
            predicate: nil,
            sortInput: nil,
            storageEngine: storageEngine,
            dataStorePublisher: dataStorePublisher,
            dataStoreConfiguration: .testDefault(),
            dispatchedModelSyncedEvent: AtomicValue(initialValue: false),
            dataStoreStatePublisher: dataStoreStateSubject.eraseToAnyPublisher()
        )
        let snapshots = taskRunner.sequence
        Task {
            var querySnapshots = [DataStoreQuerySnapshot<Post>]()
            do {
                for try await querySnapshot in snapshots {
                    querySnapshots.append(querySnapshot)
                    if querySnapshots.count == 1 {
                        XCTAssertEqual(querySnapshot.items.count, 0)
                        firstSnapshot.fulfill()
                    } else if querySnapshots.count == 2 {
                        XCTAssertEqual(querySnapshot.items.count, 1)
                        secondSnapshot.fulfill()
                    }
                }
            } catch {
                XCTFail("Failed with error \(error)")
            }
        }

        await fulfillment(of: [firstSnapshot], timeout: 1)

        let post = try createPost(id: "1")
        dataStorePublisher.send(input: post)
        await fulfillment(of: [secondSnapshot], timeout: 10)
    }

    ///  ObserveQuery will send a single snapshot when the sync state toggles
    ///  from false to true. The operation internally listens to `.modelSynced` event from
    ///  the Hub.
    ///
    /// - Given: ObserveQuery has started and the first snapshot has been received.
    /// - When:
    ///    - modelSyncedEvent is sent to the Hub
    /// - Then:
    ///    - ObserveQuery will send a second snapshot
    ///
    func testGenerateSnapshotOnObserveQueryWhenModelSynced() async throws {
        let firstSnapshot = expectation(description: "first query snapshots")
        let secondSnapshot = expectation(description: "second query snapshots")
        let thirdSnapshot = expectation(description: "third query snapshot")
        thirdSnapshot.isInverted = true

        let dispatchedModelSyncedEvent = AtomicValue(initialValue: false)
        let taskRunner = ObserveQueryTaskRunner(
            modelType: Post.self,
            modelSchema: Post.schema,
            predicate: nil,
            sortInput: nil,
            storageEngine: storageEngine,
            dataStorePublisher: dataStorePublisher,
            dataStoreConfiguration: .testDefault(),
            dispatchedModelSyncedEvent: dispatchedModelSyncedEvent,
            dataStoreStatePublisher: dataStoreStateSubject.eraseToAnyPublisher()
        )
        let snapshots = taskRunner.sequence
        Task {
            var querySnapshots = [DataStoreQuerySnapshot<Post>]()
            do {
                for try await querySnapshot in snapshots {
                    querySnapshots.append(querySnapshot)
                    if querySnapshots.count == 1 {
                        XCTAssertEqual(querySnapshot.items.count, 0)
                        XCTAssertEqual(querySnapshot.isSynced, false)
                        firstSnapshot.fulfill()
                    } else if querySnapshots.count == 2 {
                        XCTAssertEqual(querySnapshot.items.count, 0)
                        XCTAssertEqual(querySnapshot.isSynced, true)
                        secondSnapshot.fulfill()
                    } else if querySnapshots.count == 3 {
                        XCTFail("Should not receive third snapshot for a Model change")
                        thirdSnapshot.fulfill()
                    }
                }
            } catch {
                XCTFail("Failed with error \(error)")
            }
        }

        await fulfillment(of: [firstSnapshot], timeout: 5)

        dispatchedModelSyncedEvent.set(true)
        let modelSyncedEventPayload = HubPayload(
            eventName: HubPayload.EventName.DataStore.modelSynced,
            data: ModelSyncedEvent(
                modelName: Post.modelName,
                isFullSync: true,
                isDeltaSync: false,
                added: 0,
                updated: 0,
                deleted: 0
            )
        )
        Amplify.Hub.dispatch(to: .dataStore, payload: modelSyncedEventPayload)
        await fulfillment(of: [secondSnapshot], timeout: 10)

        let modelSyncedEventNotMatch = HubPayload(
            eventName: HubPayload.EventName.DataStore.modelSynced,
            data: ModelSyncedEvent.Builder().modelName
        )
        Amplify.Hub.dispatch(to: .dataStore, payload: modelSyncedEventNotMatch)
        await fulfillment(of: [thirdSnapshot], timeout: 10)
    }

    /// ObserveQuery will send the first snapshot with 2 items when storage engine
    /// is mocked to return 2 items.
    ///
    /// - Given: ObserveQuery starts
    /// - When:
    ///    -  ObserveQuery performs the initial query, two posts are queried through StorageEngine
    /// - Then:
    ///    - The items queried will return two posts in the first snapshot
    ///
    func testFirstSnapshotFromStorageQueryReturnsTwoPosts() async {
        let firstSnapshot = expectation(description: "firstSnapshot received")
        let taskRunner = ObserveQueryTaskRunner(
            modelType: Post.self,
            modelSchema: Post.schema,
            predicate: nil,
            sortInput: nil,
            storageEngine: storageEngine,
            dataStorePublisher: dataStorePublisher,
            dataStoreConfiguration: .testDefault(),
            dispatchedModelSyncedEvent: AtomicValue(initialValue: false),
            dataStoreStatePublisher: dataStoreStateSubject.eraseToAnyPublisher()
        )
        let post = Post(
            title: "model1",
            content: "content1",
            createdAt: .now()
        )
        storageEngine.responders[.query] = QueryResponder<Post>(callback: { _ in
            return .success([post, Post(title: "model1", content: "content1", createdAt: .now())])
        })
        let snapshots = taskRunner.sequence
        Task {
            var querySnapshots = [DataStoreQuerySnapshot<Post>]()
            do {
                for try await querySnapshot in snapshots {
                    querySnapshots.append(querySnapshot)
                    if querySnapshots.count == 1 {
                        XCTAssertEqual(querySnapshot.items.count, 2)
                        firstSnapshot.fulfill()
                    }
                }

            } catch {
                XCTFail("Failed with error \(error)")
            }
        }
        await fulfillment(of: [firstSnapshot], timeout: 10)
    }

    /// Multiple item changed observed will be returned in a single snapshot
    ///
    /// - Given:  The operation has started and the first query has completed.
    /// - When:
    ///    -  Observe multiple item changes.
    /// - Then:
    ///    - The items observed will be returned in the second snapshot
    ///
    func testMultipleItemChangesWillGenerateSecondSnapshot() async throws {
        let firstSnapshot = expectation(description: "first query snapshot")
        let secondSnapshot = expectation(description: "second query snapshot")

        let taskRunner = ObserveQueryTaskRunner(
            modelType: Post.self,
            modelSchema: Post.schema,
            predicate: nil,
            sortInput: nil,
            storageEngine: storageEngine,
            dataStorePublisher: dataStorePublisher,
            dataStoreConfiguration: .testDefault(),
            dispatchedModelSyncedEvent: AtomicValue(initialValue: false),
            dataStoreStatePublisher: dataStoreStateSubject.eraseToAnyPublisher()
        )
        let snapshots = taskRunner.sequence
        Task {
            var querySnapshots = [DataStoreQuerySnapshot<Post>]()
            do {
                for try await querySnapshot in snapshots {
                    querySnapshots.append(querySnapshot)
                    if querySnapshots.count == 1 {
                        XCTAssertEqual(querySnapshot.items.count, 0)
                        firstSnapshot.fulfill()
                    } else if querySnapshots.count == 2 {
                        XCTAssertEqual(querySnapshot.items.count, 3)
                        secondSnapshot.fulfill()
                    }
                }
            } catch {
                XCTFail("Failed with error \(error)")
            }
        }

        await fulfillment(of: [firstSnapshot], timeout: 1)

        let post1 = try createPost(id: "1")
        let post2 = try createPost(id: "2")
        let post3 = try createPost(id: "3")
        dataStorePublisher.send(input: post1)
        dataStorePublisher.send(input: post2)
        dataStorePublisher.send(input: post3)
        await fulfillment(of: [secondSnapshot], timeout: 10)
    }

    /// Multiple published objects (more than the `.collect` count of 1000) in a relatively short time window
    /// will cause the operation in test to exceed the limit of 1000 in its collection of items before sending a snapshot.
    /// The first snapshot will have 1000 items, and subsequent snapshots will follow as the remaining objects are published and processed.
    ///
    /// - Given:  The operation has started and the first query has completed.
    /// - When:
    ///    -  Observe 1100 item changes (beyond the `.collect` count of 1000)
    /// - Then:
    ///    - The items observed will perform a query and return 1000 items changed in the second query and the
    ///     remaining in the third query
    ///
    func testCollectOverMaxItemCountLimit() async throws {
        let firstSnapshot = expectation(description: "first query snapshot")
        let secondSnapshot = expectation(description: "second query snapshot")
        let thirdSnapshot = expectation(description: "third query snapshot")
        let validateSnapshotsComplete = expectation(description: "validate snapshots")
        let taskRunner = ObserveQueryTaskRunner(
            modelType: Post.self,
            modelSchema: Post.schema,
            predicate: nil,
            sortInput: nil,
            storageEngine: storageEngine,
            dataStorePublisher: dataStorePublisher,
            dataStoreConfiguration: .testDefault(),
            dispatchedModelSyncedEvent: AtomicValue(initialValue: false),
            dataStoreStatePublisher: dataStoreStateSubject.eraseToAnyPublisher()
        )
        let snapshots = taskRunner.sequence
        Task {
            var querySnapshots = [DataStoreQuerySnapshot<Post>]()
            do {
                for try await querySnapshot in snapshots {
                    querySnapshots.append(querySnapshot)
                    if querySnapshots.count == 1 {
                        firstSnapshot.fulfill()
                    } else if querySnapshots.count == 2 {
                        secondSnapshot.fulfill()
                    } else if querySnapshots.count == 3 {
                        thirdSnapshot.fulfill()
                    }
                }

                XCTAssertTrue(querySnapshots.count >= 3)
                XCTAssertTrue(querySnapshots[0].items.count <= querySnapshots[1].items.count)
                XCTAssertTrue(querySnapshots[1].items.count <= querySnapshots[2].items.count)
                XCTAssertTrue(querySnapshots[2].items.count <= 1_100)
                validateSnapshotsComplete.fulfill()
            } catch {
                XCTFail("Failed with error \(error)")
            }
        }

        await fulfillment(of: [firstSnapshot], timeout: 1)

        for postId in 1 ... 1_100 {
            let post = try createPost(id: "\(postId)")
            dataStorePublisher.send(input: post)
        }

        await fulfillment(of: [secondSnapshot, thirdSnapshot], timeout: 10)
        snapshots.cancel()
        await fulfillment(of: [validateSnapshotsComplete], timeout: 1.0)
    }

    /// Cancelling the subscription will no longer receive snapshots
    ///
    /// - Given:  subscriber to the operation
    /// - When:
    ///    - subscriber is cancelled
    /// - Then:
    ///    - no further snapshots are received
    ///
    func testSuccessfulSubscriptionCancel() async throws {
        let firstSnapshot = expectation(description: "first query snapshot")
        let secondSnapshot = expectation(description: "second query snapshot")
        secondSnapshot.isInverted = true

        let taskRunner = ObserveQueryTaskRunner(
            modelType: Post.self,
            modelSchema: Post.schema,
            predicate: nil,
            sortInput: nil,
            storageEngine: storageEngine,
            dataStorePublisher: dataStorePublisher,
            dataStoreConfiguration: .testDefault(),
            dispatchedModelSyncedEvent: AtomicValue(initialValue: false),
            dataStoreStatePublisher: dataStoreStateSubject.eraseToAnyPublisher()
        )
        let snapshots = taskRunner.sequence
        Task {
            var querySnapshots = [DataStoreQuerySnapshot<Post>]()
            do {
                for try await querySnapshot in snapshots {
                    querySnapshots.append(querySnapshot)
                    if querySnapshots.count == 1 {
                        firstSnapshot.fulfill()
                    } else if querySnapshots.count == 2 {
                        XCTFail("Should not receive second snapshot after cancelling")
                        secondSnapshot.fulfill()
                    }
                }
            } catch {
                XCTFail("Failed with error \(error)")
            }
        }

        await fulfillment(of: [firstSnapshot], timeout: 1)
        snapshots.cancel()
        let post1 = try createPost(id: "1")
        dataStorePublisher.send(input: post1)
        await fulfillment(of: [secondSnapshot], timeout: 1)
    }

    /// Cancelling the underlying operation will emit a completion to the subscribers
    ///
    /// - Given:  subscriber to the operation
    /// - When:
    ///    - operation is cancelled
    /// - Then:
    ///    - the subscriber receives a cancellation
    ///
    func testSuccessfulSequenceCancel() async throws {
        let firstSnapshot = expectation(description: "first query snapshot")
        let secondSnapshot = expectation(description: "second query snapshot")
        secondSnapshot.isInverted = true

        let completedEvent = expectation(description: "should have completed")
        let taskRunner = ObserveQueryTaskRunner(
            modelType: Post.self,
            modelSchema: Post.schema,
            predicate: nil,
            sortInput: nil,
            storageEngine: storageEngine,
            dataStorePublisher: dataStorePublisher,
            dataStoreConfiguration: .testDefault(),
            dispatchedModelSyncedEvent: AtomicValue(initialValue: false),
            dataStoreStatePublisher: dataStoreStateSubject.eraseToAnyPublisher()
        )
        let snapshots = taskRunner.sequence
        Task {
            var querySnapshots = [DataStoreQuerySnapshot<Post>]()
            do {
                for try await querySnapshot in snapshots {
                    querySnapshots.append(querySnapshot)
                    if querySnapshots.count == 1 {
                        firstSnapshot.fulfill()
                    } else if querySnapshots.count == 2 {
                        XCTFail("Should not receive second snapshot after cancelling")
                        secondSnapshot.fulfill()
                    }
                }
                completedEvent.fulfill()
            } catch {
                XCTFail("Failed with error \(error)")
            }
        }
        await fulfillment(of: [firstSnapshot], timeout: 1)
        snapshots.cancel()
        let post1 = try createPost(id: "1")
        dataStorePublisher.send(input: post1)

        await fulfillment(of: [secondSnapshot, completedEvent], timeout: 1)
    }

    ///  ObserveQuery's state should be able to be reset and initial query able to be started again.
    func testObserveQueryResetStateThenStartObserveQuery() async {
        let firstSnapshot = expectation(description: "first query snapshot")
        let secondSnapshot = expectation(description: "second query snapshot")
        let taskRunner = ObserveQueryTaskRunner(
            modelType: Post.self,
            modelSchema: Post.schema,
            predicate: nil,
            sortInput: nil,
            storageEngine: storageEngine,
            dataStorePublisher: dataStorePublisher,
            dataStoreConfiguration: .testDefault(),
            dispatchedModelSyncedEvent: AtomicValue(initialValue: false),
            dataStoreStatePublisher: dataStoreStateSubject.eraseToAnyPublisher()
        )
        let snapshots = taskRunner.sequence
        Task {
            var querySnapshots = [DataStoreQuerySnapshot<Post>]()
            do {
                for try await querySnapshot in snapshots {
                    querySnapshots.append(querySnapshot)
                    if querySnapshots.count == 1 {
                        firstSnapshot.fulfill()
                    } else if querySnapshots.count == 2 {
                        secondSnapshot.fulfill()
                    }
                }
            } catch {
                XCTFail("Failed with error \(error)")
            }
        }

        await fulfillment(of: [firstSnapshot], timeout: 1)
        dataStoreStateSubject.send(.stop)
        dataStoreStateSubject.send(.start(storageEngine: storageEngine))
        await fulfillment(of: [secondSnapshot], timeout: 1)
    }

    /// Multiple calls to start the observeQuery should not start again
    ///
    /// - Given: ObserverQuery operation is created, and then reset
    /// - When:
    ///    - operation.startObserveQuery twice
    /// - Then:
    ///    - Only one query should be performed / only one snapshot should be returned
    func testObserveQueryStaredShouldNotStartAgain() async {
        let firstSnapshot = expectation(description: "first query snapshot")
        let secondSnapshot = expectation(description: "second query snapshot")
        let thirdSnapshot = expectation(description: "third query snapshot")
        thirdSnapshot.isInverted = true

        let taskRunner = ObserveQueryTaskRunner(
            modelType: Post.self,
            modelSchema: Post.schema,
            predicate: nil,
            sortInput: nil,
            storageEngine: storageEngine,
            dataStorePublisher: dataStorePublisher,
            dataStoreConfiguration: .testDefault(),
            dispatchedModelSyncedEvent: AtomicValue(initialValue: false),
            dataStoreStatePublisher: dataStoreStateSubject.eraseToAnyPublisher()
        )
        let snapshots = taskRunner.sequence
        Task {
            var querySnapshots = [DataStoreQuerySnapshot<Post>]()
            do {
                for try await querySnapshot in snapshots {
                    querySnapshots.append(querySnapshot)
                    if querySnapshots.count == 1 {
                        firstSnapshot.fulfill()
                    } else if querySnapshots.count == 2 {
                        secondSnapshot.fulfill()
                    } else if querySnapshots.count == 3 {
                        thirdSnapshot.fulfill()
                    }
                }
            } catch {
                XCTFail("Failed with error \(error)")
            }
        }

        await fulfillment(of: [firstSnapshot], timeout: 1)
        dataStoreStateSubject.send(.stop)
        dataStoreStateSubject.send(.start(storageEngine: storageEngine))
        dataStoreStateSubject.send(.start(storageEngine: storageEngine))
        await fulfillment(of: [secondSnapshot, thirdSnapshot], timeout: 1)
        XCTAssertTrue(taskRunner.observeQueryStarted)
    }

    /// ObserveQuery operation entry points are `resetState`, `startObserveQuery`, and `onItemChanges(mutationEvents)`.
    /// Ensure concurrent random sequences of these API calls do not cause issues such as data race.
    func testConcurrent() async {
        let completeReceived = expectation(description: "complete received")
        completeReceived.isInverted = true
        let taskRunner = ObserveQueryTaskRunner(
            modelType: Post.self,
            modelSchema: Post.schema,
            predicate: nil,
            sortInput: nil,
            storageEngine: storageEngine,
            dataStorePublisher: dataStorePublisher,
            dataStoreConfiguration: .testDefault(),
            dispatchedModelSyncedEvent: AtomicValue(initialValue: false),
            dataStoreStatePublisher: dataStoreStateSubject.eraseToAnyPublisher()
        )
        let post = Post(
            title: "model1",
            content: "content1",
            createdAt: .now()
        )
        storageEngine.responders[.query] = QueryResponder<Post>(callback: { _ in
            return .success([post, Post(title: "model1", content: "content1", createdAt: .now())])
        })
        let snapshots = taskRunner.sequence
        Task {
            do {
                for try await _ in snapshots {
                }
                completeReceived.fulfill()
            } catch {
                XCTFail("Failed with error \(error)")
            }
        }
        await withTaskGroup(of: Void.self) { group in
            for _ in 0 ..< 3_000 {
                group.addTask {
                    let index = Int.random(in: 1 ... 5)
                    if index == 1 {
                        self.dataStoreStateSubject.send(.stop)
                    } else if index == 2 {
                        self.dataStoreStateSubject.send(.start(storageEngine: self.storageEngine))
                    } else {
                        do {
                            let itemChange = try self.createPost(id: post.id)
                            let itemChange2 = try self.createPost()
                            taskRunner.onItemsChangeDuringSync(mutationEvents: [itemChange, itemChange2])
                        } catch {
                            XCTFail("Failed to create post")
                        }
                    }
                }
            }
            for await _ in group {
            }
        }
        await fulfillment(of: [completeReceived], timeout: 10)
    }

    /// When a predicate like `title.beginsWith("title")` is given, the models that matched the predicate
    /// should be added to the snapshots, like `post` and `post2`. When `post2.title` is updated to no longer
    /// match the predicate, it should be removed from the snapshot.
    func testUpdatedModelNoLongerMatchesPredicateRemovedFromSnapshot() async throws {
        let firstSnapshot = expectation(description: "first query snapshots")
        let secondSnapshot = expectation(description: "second query snapshots")
        let thirdSnapshot = expectation(description: "third query snapshots")
        let fourthSnapshot = expectation(description: "fourth query snapshots")
        let taskRunner = ObserveQueryTaskRunner(
            modelType: Post.self,
            modelSchema: Post.schema,
            predicate: Post.keys.title.beginsWith("title"),
            sortInput: QuerySortInput.ascending(Post.keys.id).asSortDescriptors(),
            storageEngine: storageEngine,
            dataStorePublisher: dataStorePublisher,
            dataStoreConfiguration: .testDefault(),
            dispatchedModelSyncedEvent: AtomicValue(initialValue: true),
            dataStoreStatePublisher: dataStoreStateSubject.eraseToAnyPublisher()
        )

        let snapshots = taskRunner.sequence
        Task {
            var querySnapshots = [DataStoreQuerySnapshot<Post>]()
            do {
                for try await querySnapshot in snapshots {
                    querySnapshots.append(querySnapshot)
                    if querySnapshots.count == 1 {
                        // First snapshot is empty from the initial query
                        XCTAssertEqual(querySnapshot.items.count, 0)
                        firstSnapshot.fulfill()
                    } else if querySnapshots.count == 2 {
                        // Second snapshot contains `post` since it matches the predicate
                        XCTAssertEqual(querySnapshot.items.count, 1)
                        XCTAssertEqual(querySnapshot.items[0].id, "1")
                        secondSnapshot.fulfill()
                    } else if querySnapshots.count == 3 {
                        // Third snapshot contains both posts since they both match the predicate
                        XCTAssertEqual(querySnapshot.items.count, 2)
                        XCTAssertEqual(querySnapshot.items[0].id, "1")
                        XCTAssertEqual(querySnapshot.items[1].id, "2")
                        thirdSnapshot.fulfill()
                    } else if querySnapshots.count == 4 {
                        // Fourth snapshot no longer has the post2 since it was updated to not match the predicate
                        // and deleted at the same time.
                        XCTAssertEqual(querySnapshot.items.count, 1)
                        XCTAssertEqual(querySnapshot.items[0].id, "1")
                        fourthSnapshot.fulfill()
                    }
                }
            } catch {
                XCTFail("Failed with error \(error)")
            }
        }
        await fulfillment(of: [firstSnapshot], timeout: 5)

        let post = try createPost(id: "1", title: "title 1")
        dataStorePublisher.send(input: post)
        let post2 = try createPost(id: "2", title: "title 2")
        dataStorePublisher.send(input: post2)
        var updatedPost2 = try createPost(id: "2", title: "Does not match predicate")
        updatedPost2.mutationType = MutationEvent.MutationType.update.rawValue
        dataStorePublisher.send(input: updatedPost2)
        await fulfillment(of: [secondSnapshot, thirdSnapshot, fourthSnapshot], timeout: 10)
    }


    // MARK: - Helpers

    func createPost(id: String = UUID().uuidString, title: String? = nil) throws -> MutationEvent {
        try MutationEvent(
            model: Post(
                id: id,
                title: title ?? "model1",
                content: "content1",
                createdAt: .now()
            ),
            modelSchema: Post.schema,
            mutationType: MutationEvent.MutationType.create
        )
    }
}
