/*
 * 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/>.
 */

#ifndef TASKRUNNER_H
#define TASKRUNNER_H

#include <atomic>
#include <cstdio>
#include <functional>
#include <future>
#include <mutex>
#include <set>
#include <thread>
#include <unordered_map>

#include "concurrentqueue/concurrentqueue.h"

namespace task {

using TaskId = uint64_t;

constexpr TaskId INVALID_TASK_ID = 0;  // 0是无效任务ID

TaskId generateTaskId();

enum class TaskType {
    TextChat,
    OnceTextToSpeech,
    ContinuousTextToSpeech,
    OnceSpeechRecognition,
    ContinuousSpeechRecognition,
    TextToImage,
    TextRecognition,
    Initialization,
    Stop,
};

class Task {
public:
    Task() = default;
    Task(TaskId taskId, TaskType taskType, std::function<void()> task);
    void operator()() const;

public:
    std::function<void()> task;
    TaskId taskId;
    TaskType taskType;
};

class TaskRunner {
public:
    static TaskRunner &getInstance();
    ~TaskRunner();

    void addTask(Task task);
    void deleteTask(TaskId taskId);
    void waitForTasks();
    void setMaxConcurrentTasks(TaskType taskType, int maxTasks);

private:
    TaskRunner();
    void run();
    void dispatchTasks();
    void dispatchTasksLoggingBlockingException();
    void dispatchTask(Task &task);
    void handleWaitingTask(TaskType taskType);

    bool isCanceledTask(TaskId taskId);
    void removeCanceledTask(TaskId taskId);
    void cleanTaskResult();

private:
    std::thread *dispatchThread{nullptr};
    moodycamel::ConcurrentQueue<Task> taskQueue;
    std::set<TaskId> canceledTasks;
    std::unordered_map<TaskType, moodycamel::ConcurrentQueue<Task>> waitingQueue;
    std::atomic<bool> stop{false};
    std::unordered_map<TaskType, int> maxConcurrentTasks = {{TaskType::TextChat, 8},
                                                            {TaskType::OnceTextToSpeech, 2},
                                                            {TaskType::ContinuousTextToSpeech, 4},
                                                            {TaskType::OnceSpeechRecognition, 4},
                                                            {TaskType::ContinuousSpeechRecognition, 4},
                                                            {TaskType::TextToImage, 4},
                                                            {TaskType::TextRecognition, 8},
                                                            {TaskType::Initialization, 4}};
    std::unordered_map<TaskType, std::atomic_int> currentTasks;

    // 每个类型的waiting队列size不为0的计数器，为了当主任务队列全部被阻塞住的时候打印日志
    std::unordered_map<TaskType, int> counters{{TaskType::TextChat, 0},
                                               {TaskType::OnceTextToSpeech, 0},
                                               {TaskType::ContinuousTextToSpeech, 0},
                                               {TaskType::OnceSpeechRecognition, 0},
                                               {TaskType::ContinuousSpeechRecognition, 0},
                                               {TaskType::TextToImage, 0},
                                               {TaskType::TextRecognition, 0},
                                               {TaskType::Initialization, 0}};

    std::mutex mutex_;
    std::vector<std::future<void>> futures_;
};

}  // namespace task

#endif  // TASKRUNNER_H
