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

import Foundation

/// A database record (or wrapper around one) representing an enqueued task
/// to be run. Contains whatever ``TaskRecordRunner`` needs to run the
/// task; all the generic ``TaskQueueLoader`` needs from it is some id
/// for uniqueing purposes.
public protocol TaskRecord {
    associatedtype IDType: Hashable

    var id: IDType { get }

    /// If non-nil, tasks will wait until this timestamp before running.
    /// Default to nil.
    var nextRetryTimestamp: UInt64? { get }
}

extension TaskRecord {
    var nextRetryTimestamp: UInt64? { return nil }
}

/// Intermediary between the ``TaskQueueLoader`` and the queue table in the database.
/// Used to peek (get next tasks to do) and remove tasks once complete.
public protocol TaskRecordStore {
    associatedtype Record: TaskRecord, Sendable

    /// Fetch the next `count` records to run, in priority order.
    ///
    /// Repeated results across peek calls is allowed; it is NOT expected that each peek call
    /// will retain an offset and return the N results _after_ the last peek.
    ///
    /// For example, the ``TaskQueueLoader`` might peek 4 and start running those 4 tasks.
    /// When the 1st task finishes, it will delete that record, then call peek 4 again.
    /// Tasks 2, 3, and 4 are still running and therefore not deleted from the database, so its
    /// expected that peek could return the same 3 plus the next record (e.g. returns [2,3,4,5].)
    /// Or it could return a totally new set of higher priority that were inserted since the last peek.
    func peek(count: UInt, tx: DBReadTransaction) throws -> [Record]

    /// Remove a record from the database table queue.
    ///
    /// Called by the ``TaskQueueLoader`` when the task either succeeds or fails with
    /// an unretryable error. Outside callers may remove queued records as they see fit.
    /// Removing a record that is currently running will not stop it from completing.
    func removeRecord(_ record: Record, tx: DBWriteTransaction) throws
}

/// The result of a ``TaskRecordRunner`` run.
public enum TaskRecordResult {
    /// Sucess! The task will be cleared from the queue by the loader (via the store).
    case success
    /// A retryable failure. The task will not be cleared from the queue by default; the
    /// runner gets an opportunity to respond to the error, including making any
    /// updates it desires to the record (including clearing it and reinserting if desired).
    case retryableError(Error)
    /// A non-retryable failure. The task will be cleared from the queue by the loader.
    /// Only return this result if it's expected that no retry will ever happen within
    /// the queue mechanism.
    case unretryableError(Error)
    /// Something happened between when the task was queued and when it was run
    /// that invalidated the need for any work (the task was "cancelled"). For example,
    /// perhaps the user deleted the owning thread, or whatever.
    /// Like a success, the row is removed and no error is assumed, but it's not
    /// _strictly_ a success so the result type and callback are differentiated.
    case obsolete
}

/// Used with a ``TaskQueueLoader``; this actually runs the task for each record
/// and performs any necessary cleanup on success or failure within the same write
/// transaction that clears the task record from the queue.
public protocol TaskRecordRunner {
    associatedtype Store: TaskRecordStore

    var store: Store { get }

    /// Run the task for a single record.
    /// Returned errors do NOT interrupt task processing; the loader will continue
    /// to process tasks until there are none left.
    ///
    /// - parameter loader: Provided so that runners can choose to call
    /// ``TaskQueueLoader/stop(reason:)`` if some error they encounter
    /// should stop all future tasks, not just the current one.
    func runTask(
        record: Store.Record,
        loader: TaskQueueLoader<Self>,
    ) async -> TaskRecordResult

    /// Called by ``TaskQueueLoader`` when the task completes successfully,
    /// with the same write transaction used to delete the task's database record.
    ///
    /// Thrown errors WILL interrupt task processing; it's assumed anything that goes
    /// wrong here is a severe error affecting the queue itself.
    func didSucceed(
        record: Store.Record,
        tx: DBWriteTransaction,
    ) throws

    /// Called by ``TaskQueueLoader`` when the task fails.
    /// If the error is non-retryable, the record is deleted from the databse in the
    /// same transaction.
    /// Otherwise the record is left alone and it is up to the runner to modify it as
    /// needed. However it is modified, it will be retried by the loader the next time
    /// it is returned as a result of the store's `peek` method.
    ///
    /// Thrown errors WILL interrupt task processing; it's assumed anything that goes
    /// wrong here is a severe error affecting the queue itself.
    func didFail(
        record: Store.Record,
        error: Error,
        isRetryable: Bool,
        tx: DBWriteTransaction,
    ) throws

    /// Called by ``TaskQueueLoader`` when the task is obsoleted,
    /// with the same write transaction used to delete the task's database record.
    ///
    /// Thrown errors WILL interrupt task processing; it's assumed anything that goes
    /// wrong here is a severe error affecting the queue itself.
    func didObsolete(
        record: Store.Record,
        tx: DBWriteTransaction,
    ) throws

    /// Called by ``TaskQueueLoader`` when all tasks have finished.
    /// In other words, when ``TaskRecordStore`` peek returns an empty array.
    func didDrainQueue() async
}

extension TaskRecordRunner {
    public func didDrainQueue() async {}
}

/// Utility class that helps working with serialized queues for which each needs to run
/// an async task and then remove itself from the queue.
///
/// Works with a ``TaskRecordStore`` and ``TaskRecordRunner`` to load
/// ``TaskRecord``s from the database N at a time and run their associated async
/// tasks in parallel. When each task completes or fails, it then removes the task from
/// the queue (unless it fails with a retryable error) and moves on to the next task(s).
///
/// When should you use this type as opposed to ``JobRunner`` and ``JobRecord``?
/// This type is bring-your-own database table (or in memory queue), thus granting flexibility particularly
/// if you want to enforce ORDER BYs when querying the next row to run, which allows for more options
/// than just the FIFO ordering offered by ``JobRunner``.
/// On the other hand, ``JobRunner`` gives you more in-built functionality, particularly with
/// respect to error handling and retries. Callers of this class are left to do that on their own as well
/// as left the task of actually starting up the loader and managing the queue table.
public actor TaskQueueLoader<Runner: TaskRecordRunner & Sendable> {

    typealias Store = Runner.Store
    typealias Record = Store.Record

    public let maxConcurrentTasks: UInt

    public var isRunning: Bool {
        switch state {
        case .notRunning, .cleaningUp, .cancelled:
            return false
        case .running:
            return true
        }
    }

    private nonisolated let dateProvider: DateProvider
    private let db: any DB
    private let runner: Runner
    private var store: Store { runner.store }
    private let sleep: (_ nanoseconds: UInt64) async throws -> Void

    /// WARNING: the runner (and therefore any of its strong references) is strongly
    /// captured by this class and will be retained for its lifetime.
    init(
        maxConcurrentTasks: UInt,
        dateProvider: @escaping DateProvider,
        db: any DB,
        runner: Runner,
        sleep: @escaping (_ nanoseconds: UInt64) async throws -> Void,
    ) {
        self.maxConcurrentTasks = maxConcurrentTasks
        self.dateProvider = dateProvider
        self.db = db
        self.runner = runner
        self.sleep = sleep
    }

    /// WARNING: the runner (and therefore any of its strong references) is strongly
    /// captured by this class and will be retained for its lifetime.
    public init(
        maxConcurrentTasks: UInt,
        dateProvider: @escaping DateProvider,
        db: any DB,
        runner: Runner,
    ) {
        self.init(
            maxConcurrentTasks: maxConcurrentTasks,
            dateProvider: dateProvider,
            db: db,
            runner: runner,
            sleep: {
                try await Task.sleep(nanoseconds: $0)
            },
        )
    }

    private enum State {
        case notRunning
        case running(UUID, Task<Void, Error>)
        case cleaningUp(UUID, Task<Void, Error>)
        case cancelled(UUID, Task<Void, Error>)

        func isUnchanged(from oldState: State) -> Bool {
            switch (oldState, self) {
            case (.notRunning, .notRunning):
                return true
            case (.notRunning, _):
                return false
            case (.running(let oldId, _), .running(let newId, _)):
                return oldId == newId
            case (.running, _):
                return false
            case (.cleaningUp(let oldId, _), .cleaningUp(let newId, _)):
                return oldId == newId
            case (.cleaningUp, _):
                return false
            case (.cancelled(let oldId, _), .cancelled(let newId, _)):
                return oldId == newId
            case (.cancelled, _):
                return false
            }
        }
    }

    private var state = State.notRunning
    /// Random IDs of un-cancelled observers of the currently running task.
    private var runningTaskObservers = Set<UUID>()
    /// Error provided when the task was stopped; if not nil, throw it on callers of loadAndRunTasks.
    private var stoppedReason: Error?

    private var currentTaskIds = Set<Record.IDType>()

    /// Load tasks, N at a time, and begin running any that are not already running.
    /// (N = max concurrent tasks)
    /// Runs until all tasks are finished (finished = table peek returns empty).
    ///
    /// Throws an error IFF some database operation relating to the queue or post-task cleanup fails;
    /// within-task failures are handled by the runner and do NOT interrupt processing of subsequent tasks.
    ///
    /// Cancellation causes code execution to be returned to the caller but does not _necessarily_ cancel
    /// the execution of subsequent tasks. All callers of `loadAndRunTasks()` await a single runner.
    /// As long as some uncancelled task context is awaiting the result, the runner will continue to execute
    /// subsequent tasks. A single caller cancelling simply releases that caller from waiting on the runner
    /// to finish; the runner isn't cancelled until _all_ callers cancel.
    /// If some caller wishes to _force_ the runner to stop, call `stop()` instead of cancelling the task.
    public func loadAndRunTasks() async throws {
        let task: Task<Void, Error>
        let state = self.state
        switch state {
        case .notRunning:
            // Start a new task.
            let taskId = UUID()
            task = Task {
                try await self._loadAndRunTasks(taskId: taskId)
            }
            self.state = .running(taskId, task)
            self.runningTaskObservers = Set()
        case .running(_, let _task):
            task = _task
        case .cancelled(_, let cancelledTask):
            // We want to wait for cancellation to finish
            // applying, and then try again.
            try? await cancelledTask.value
            // We expect that state will have been cleaned
            // up by the time runningTask finishes, but
            // just in case do cleanup here.
            if self.state.isUnchanged(from: state) {
                self.state = .notRunning
            }
            return try await loadAndRunTasks()
        case .cleaningUp(_, let runningTask):
            // We want to wait for cleanup to finish
            // applying, and then try again.
            try? await runningTask.value
            // We expect that state will have been cleaned
            // up by the time runningTask finishes, but
            // just in case do cleanup here.
            if self.state.isUnchanged(from: state) {
                self.state = .notRunning
            }
            return try await loadAndRunTasks()
        }

        let observerId = UUID()
        runningTaskObservers.insert(observerId)
        // The cancellable continuation means if the calling context cancels
        // it will immediately resume while the inner task keeps running.
        // In addition, we use `withTaskCancellationHandler` to get the
        // onCancel callback so that we can track when _all_ observers
        // have cancelled and then we pass along the cancellation to
        // the actual running task.
        let continuation = CancellableContinuation<Void>()
        Task {
            do {
                try await task.value
                continuation.resume(with: .success(()))
            } catch let cancellationError as CancellationError {
                let error = self.stoppedReason ?? cancellationError
                continuation.resume(with: .failure(error))
            } catch let error {
                continuation.resume(with: .failure(error))
            }
        }

        try await withTaskCancellationHandler(
            operation: {
                try await continuation.wait()
            },
            onCancel: {
                continuation.cancel()
                Task {
                    try await observerDidCancel(observerId)
                }
            },
        )
    }

    private func observerDidCancel(_ id: UUID) throws {
        guard runningTaskObservers.contains(id) else {
            return
        }
        self.runningTaskObservers.remove(id)
        if runningTaskObservers.isEmpty {
            // We can stop and cancel the running task.
            try self.stop()
        }
    }

    public func stop(reason: Error? = nil) throws {
        switch state {
        case .notRunning:
            break
        case .cancelled:
            // Already cancelled; prefer the initial
            // reason (if any) and let it finish cancelling.
            return
        case .running(let taskId, let runningTask), .cleaningUp(let taskId, let runningTask):
            self.stoppedReason = reason
            runningTask.cancel()
            self.state = .cancelled(taskId, runningTask)
        }
    }

    private func _loadAndRunTasks(taskId: UUID) async throws {
        // Check cancellation at the start of each attempt.
        // This method is called recursively, so now is a good time to check.
        try Task.checkCancellation()

        if currentTaskIds.count >= maxConcurrentTasks {
            return
        }

        // _Always_ read as many tasks as we can run; afterwards we can
        // dedupe against the tasks already running and keep as many as
        // needed to fill up the max concurrent count.
        let recordCandidates = try db.read { tx in
            try store.peek(count: self.maxConcurrentTasks, tx: tx)
        }

        let records = recordCandidates.filter { record in
            !currentTaskIds.contains(record.id)
        }
        guard !records.isEmpty else {
            if currentTaskIds.isEmpty {
                switch self.state {
                case .notRunning:
                    return
                case .cancelled:
                    owsFailDebug("Cancel should have applied to the task context")
                    return
                case .cleaningUp:
                    return
                case .running(let stateTaskId, let runningTask):
                    if stateTaskId == taskId {
                        state = .cleaningUp(taskId, runningTask)
                        await self.runner.didDrainQueue()
                        switch self.state {
                        case .notRunning:
                            owsFailDebug("State is not running but we are, in fact, running")
                            return
                        case .running(let stateTaskId, _):
                            owsFailDebug("Not cleaning up? How?")
                            fallthrough
                        case .cancelled(let stateTaskId, _):
                            // Cancellation might've applied while running didDrainQueue;
                            // We're done now so just finish.
                            fallthrough
                        case .cleaningUp(let stateTaskId, _):
                            if stateTaskId == taskId {
                                state = .notRunning
                            }
                        }
                    }
                }
            }
            return
        }
        records.forEach { currentTaskIds.insert($0.id) }

        let runner = self.runner
        try await withThrowingTaskGroup(of: Void.self) { taskGroup in
            records.forEach { record in
                taskGroup.addTask {
                    let nowMs = self.dateProvider().ows_millisecondsSince1970
                    if
                        let nextRetryTimestamp = record.nextRetryTimestamp,
                        nowMs < nextRetryTimestamp
                    {
                        try await self.sleep(NSEC_PER_MSEC * (nextRetryTimestamp - nowMs))
                    }
                    let taskResult = await runner.runTask(record: record, loader: self)
                    switch taskResult {
                    case .success:
                        try await self.didSucceed(record: record)
                    case .retryableError(let error):
                        try await self.didFail(record: record, error: error, isRetryable: true)
                    case .unretryableError(let error):
                        try await self.didFail(record: record, error: error, isRetryable: false)
                    case .obsolete:
                        try await self.didObsolete(record: record)
                    }
                    // As soon as we finish any task, start loading more tasks to run.
                    try await self._loadAndRunTasks(taskId: taskId)
                }
            }
            try await taskGroup.waitForAll()
        }
    }

    private func didSucceed(record: Store.Record) async throws {
        try await db.awaitableWrite { tx in
            try self.store.removeRecord(record, tx: tx)
            try self.runner.didSucceed(record: record, tx: tx)
        }
        self.currentTaskIds.remove(record.id)
    }

    private func didFail(record: Store.Record, error: Error, isRetryable: Bool) async throws {
        try await db.awaitableWrite { tx in
            if !isRetryable {
                // Remove the record for non-retryable errors.
                try self.store.removeRecord(record, tx: tx)
            }
            try self.runner.didFail(record: record, error: error, isRetryable: isRetryable, tx: tx)
        }
        self.currentTaskIds.remove(record.id)
    }

    private func didObsolete(record: Store.Record) async throws {
        try await db.awaitableWrite { tx in
            try self.store.removeRecord(record, tx: tx)
            try self.runner.didObsolete(record: record, tx: tx)
        }
        self.currentTaskIds.remove(record.id)
    }
}