package ch.threema.domain.taskmanager

import ch.threema.base.utils.getThreemaLogger
import ch.threema.domain.protocol.connection.socket.ServerSocketCloseReason
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock

private val logger = getThreemaLogger("LocalTaskQueue")

internal class LocalTaskQueue(
    private val taskArchiver: TaskArchiver,
    private val getDebugString: Task<*, *>.() -> String,
) {
    /**
     * The mutex is required to access [taskQueue].
     */
    private val mutex = Mutex()

    /**
     * The task queue for tasks that have been initiated by this device. Note that this queue should
     * only be accessed when holding the [mutex], as both the scheduler and executor modify this
     * queue.
     */
    private val taskQueue by lazy {
        ArrayDeque<LocalTaskQueueElement<*>>().also { queue ->
            // First load all the persisted tasks and keep them persisted until they are completed
            taskArchiver.loadAllTasks().forEach {
                queue.add(LocalTaskQueueElement(it, CompletableDeferred()))
            }
        }
    }

    internal suspend fun <R> add(task: Task<R, TaskCodec>, done: CompletableDeferred<R>) {
        mutex.withLock {
            taskQueue.add(LocalTaskQueueElement(task, done))
            taskArchiver.addTask(task)
        }
    }

    internal fun getNextOrNull(): TaskQueue.TaskQueueElement? = runBlocking {
        mutex.withLock { taskQueue.firstOrNull() }
    }

    internal fun hasPendingTasks() = runBlocking {
        mutex.withLock { taskQueue.isNotEmpty() }
    }

    internal suspend fun removeDropOnDisconnectTasks(closeReason: ServerSocketCloseReason?) {
        val tasksToDrop = mutex.withLock {
            taskQueue.filter { taskQueueElement -> taskQueueElement.shouldDropOnDisconnect }
        }

        tasksToDrop.forEach { taskQueueElement ->
            taskQueueElement.completeExceptionally(
                ConnectionStoppedException(
                    "Dropping task before its execution because connection has been lost",
                    closeReason,
                ),
            )
        }
    }

    /**
     * An element of the local task queue. This is a task that has not been initiated by an incoming
     * message but has been scheduled locally.
     */
    private inner class LocalTaskQueueElement<R>(
        private val task: Task<R, TaskCodec>,
        private val done: CompletableDeferred<R>,
    ) : TaskQueue.TaskQueueElement {
        override val maximumNumberOfExecutions: Int = 5

        override val shouldDropOnDisconnect = task is DropOnDisconnectTask

        override suspend fun run(handle: TaskCodec) {
            logger.info("Running task {}", task.getDebugString())
            try {
                done.complete(task.invoke(handle))
            } catch (e: Exception) {
                when (e) {
                    is CancellationException -> logger.info("Cancelled task {}", task.getDebugString())
                    else -> logger.warn("Stopped task {} exceptionally", task.getDebugString())
                }
                throw e
            }
            logger.info("Completed task {} successfully", task.getDebugString())
            mutex.withLock {
                taskQueue.remove(this)
                taskArchiver.removeTask(task)
            }
        }

        override fun isCompleted() = done.isCompleted

        override suspend fun completeExceptionally(exception: Throwable) {
            logger.warn("Completed task {} exceptionally", task.getDebugString(), exception)
            done.completeExceptionally(exception)
            mutex.withLock {
                taskQueue.remove(this)
                taskArchiver.removeTask(task)
            }
        }
    }
}
