/*
 * Copyright 2024 KylinSoft Co., Ltd.
 *
 * This program is free software: you can redistribute it and/or modify it under
 * the terms of the GNU General Public License as published by the Free Software
 * Foundation, either version 3 of the License, or (at your option) any later
 * version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
 * details.
 *
 * You should have received a copy of the GNU General Public License along with
 * this program. If not, see <https://www.gnu.org/licenses/>.
 */

#include "taskrunner.h"

#include <sys/types.h>

#include <mutex>
#include <thread>

#include "utils/logger.h"

namespace task {

TaskId generateTaskId() {
    static std::atomic<TaskId> taskId = 0;
    if (taskId == std::numeric_limits<std::atomic<TaskId>>::max()) {
        taskId = 1;
    } else {
        taskId++;
    }
    return taskId;
}

Task::Task(TaskId taskId, TaskType taskType, std::function<void()> task)
    : taskId(taskId), taskType(taskType), task(task) {}

void Task::operator()() const { task(); }

TaskRunner &TaskRunner::getInstance() {
    static TaskRunner taskRunner;
    return taskRunner;
}

TaskRunner::TaskRunner() { dispatchThread = new std::thread(&TaskRunner::run, this); }

TaskRunner::~TaskRunner() {
    stop = true;
    if (dispatchThread != nullptr) {
        dispatchThread->join();
        delete dispatchThread;
    }

    for (auto &future : futures_) {
        future.wait();
    }
}

void TaskRunner::addTask(Task task) {
    bool success = taskQueue.enqueue(std::move(task));
    if (not success) {
        LOGE("Add task failed.");
    }
}

void TaskRunner::deleteTask(TaskId taskId) {
    std::lock_guard<std::mutex> locker(mutex_);
    canceledTasks.insert(taskId);
}

void TaskRunner::waitForTasks() {
    while (taskQueue.size_approx() > 0) {
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }
}

void TaskRunner::setMaxConcurrentTasks(TaskType taskType, int maxTasks) { maxConcurrentTasks[taskType] = maxTasks; }

void TaskRunner::run() {
    while (!stop) {
        dispatchTasks();
        cleanTaskResult();
        std::this_thread::sleep_for(std::chrono::milliseconds(1));
    }
}

void TaskRunner::dispatchTasks() {
    Task task;
    if (taskQueue.try_dequeue(task)) {
        if (!isCanceledTask(task.taskId)) {
            dispatchTask(task);
        } else {
            removeCanceledTask(task.taskId);
        }
    } else {
        dispatchTasksLoggingBlockingException();
    }
}

void TaskRunner::dispatchTasksLoggingBlockingException() {
    // 当某个类型的任务达到该类型任务容量限制时，且每个任务都属于阻塞状态，会触
    // 发这里的代码逻辑，当某个类型的waitingQueue被连续10秒钟检测到大小不为0，认
    // 为该类型任务全部阻塞
    for (auto &[taskType, waitingQueue] : waitingQueue) {
        if (waitingQueue.size_approx() == 0) {
            counters[taskType] = 0;  // 重新计数
            continue;
        }
        counters[taskType] += 1;
        if (counters[taskType] == 10000) {
            // 约10秒钟
            LOGE("This type of task is blocked in its entirety, task type: {}", int(taskType));
            counters[taskType] = 0;  // 重新计数
        }
    }
}

void TaskRunner::dispatchTask(Task &task) {
    if (currentTasks[task.taskType] < maxConcurrentTasks[task.taskType]) {
        auto future = std::async(std::launch::async, [this, task]() {
            currentTasks[task.taskType]++;
            task();
            handleWaitingTask(task.taskType);
            currentTasks[task.taskType]--;
        });
        futures_.emplace_back(std::move(future));
    } else {
        waitingQueue[task.taskType].enqueue(task);
    }
}

void TaskRunner::handleWaitingTask(TaskType taskType) {
    Task task;
    while (waitingQueue[taskType].size_approx() != 0) {
        if (waitingQueue[taskType].try_dequeue(task)) {
            task();
        }
    }
}

bool TaskRunner::isCanceledTask(TaskId taskId) {
    std::lock_guard<std::mutex> locker(mutex_);
    return canceledTasks.find(taskId) != canceledTasks.end();
}

void TaskRunner::removeCanceledTask(TaskId taskId) {
    std::lock_guard<std::mutex> locker(mutex_);
    canceledTasks.erase(taskId);
}

void TaskRunner::cleanTaskResult() {
    futures_.erase(std::remove_if(futures_.begin(), futures_.end(),
                                  [](const std::future<void> &future) {
                                      return future.wait_for(std::chrono::seconds(0)) == std::future_status::ready;
                                  }),
                   futures_.end());
}
};  // namespace task
