/*
 * .============.
 * // M A K E / \
 * // C++ DEV / \
 * // E A S Y / \/ \
 * ++ ----------. \/\ .
 * \\ \ \ /\ /
 * \\ \ \ /
 * \\ \ \ /
 * -============'
 *
 * Copyright (c) 2018 Hevake and contributors, all rights reserved.
 *
 * This file is part of cpp-tbox (https://github.com/cpp-main/cpp-tbox)
 * Use of this source code is governed by MIT license that can be found
 * in the LICENSE file in the root of the source tree. All contributing
 * project authors may be found in the CONTRIBUTORS.md file in the root
 * of the source tree.
 */
#include "thread_pool.h"

#include <cinttypes>
#include <array>
#include <map>
#include <set>
#include <deque>
#include <thread>
#include <mutex>
#include <algorithm>
#include <condition_variable>
#include <chrono>

#include <tbox/base/log.h>
#include <tbox/base/cabinet.hpp>
#include <tbox/base/assert.h>
#include <tbox/base/catch_throw.h>
#include <tbox/base/object_pool.hpp>
#include <tbox/base/wrapped_recorder.h>
#include <tbox/event/loop.h>

#undef MODULE_ID
#define MODULE_ID "tbox.thread_pool"

namespace tbox {
namespace eventx {

using Clock = std::chrono::steady_clock;

//! ThreadPool 的私有数据
struct ThreadPool::Data {
 event::Loop *wp_loop = nullptr; //!< 主线程

 bool is_ready = false; //! 是否已经初始化了

 size_t min_thread_num = 0; //!< 最少的线程个数
 size_t max_thread_num = 0; //!< 最多的线程个数

 std::mutex lock; //!< 互斥锁
 std::condition_variable cond_var; //!< 条件变量

 cabinet::Cabinet<Task> undo_tasks_cabinet;
 std::array<std::deque<TaskToken>, THREAD_POOL_PRIO_SIZE> undo_tasks_token; //!< 优先级任务列表,THREAD_POOL_PRIO_SIZE级
 std::set<TaskToken> doing_tasks_token; //!< 记录正在从事的任务

 size_t idle_thread_num = 0; //!< 空间线程个数
 cabinet::Cabinet<std::thread> threads_cabinet;
 bool all_threads_stop_flag = false; //!< 是否所有工作线程立即停止标记

 size_t undo_task_peak_num_ = 0;

 ObjectPool<Task> task_pool{64};
};

/**
 * 任务项
 */
struct ThreadPool::Task {
 TaskToken token;
 NonReturnFunc backend_task; //! 任务在工作线程中执行函数
 NonReturnFunc main_cb; //! 任务执行完成后由main_loop执行的回调函数
 Clock::time_point create_time_point;

 Task *next = nullptr;
};

/////////////////////////////////////////////////////////////////////////////////

ThreadPool::ThreadPool(event::Loop *main_loop) :
 d_(new Data)
{
 d_->wp_loop = main_loop;
}

ThreadPool::~ThreadPool()
{
 if (d_->is_ready)
 cleanup();

 delete d_;
}

bool ThreadPool::initialize(ssize_t min_thread_num, ssize_t max_thread_num)
{
 if (d_->is_ready) {
 LogWarn("it has ready, cleanup() first");
 return false;
 }

 if (max_thread_num < 0 || min_thread_num < 0 ||
 min_thread_num > max_thread_num || max_thread_num == 0) {

 LogWarn("min_thread_num or max_thread_num invalid, min:%d, max:%d", min_thread_num, max_thread_num);
 return false;
 }

 {
 std::lock_guard<std::mutex> lg(d_->lock);
 d_->min_thread_num = min_thread_num;
 d_->max_thread_num = max_thread_num;

 for (ssize_t i = 0; i < min_thread_num; ++i)
 if (!createWorker())
 return false;
 }

 d_->all_threads_stop_flag = false;
 d_->is_ready = true;

 return true;
}

ThreadExecutor::TaskToken ThreadPool::execute(NonReturnFunc &&backend_task, NonReturnFunc &&main_cb, int prio)
{
 RECORD_SCOPE();
 TaskToken token;

 if (!d_->is_ready) {
 LogWarn("need initialize() first");
 return token;
 }

 if (prio < THREAD_POOL_PRIO_MIN)
 prio = THREAD_POOL_PRIO_MIN;
 else if (prio > THREAD_POOL_PRIO_MAX)
 prio = THREAD_POOL_PRIO_MAX;

 int level = prio + THREAD_POOL_PRIO_MAX;

 {
 std::lock_guard<std::mutex> lg(d_->lock);

 Task *item = d_->task_pool.alloc();
 item->backend_task = std::move(backend_task);
 item->main_cb = std::move(main_cb);
 item->create_time_point = Clock::now();
 item->token = token = d_->undo_tasks_cabinet.alloc(item);

 d_->undo_tasks_token.at(level).push_back(token);
 //! 如果空闲线程不够分配未认领的任务,且还可以再创建新的线程
 if (d_->undo_tasks_cabinet.size() > d_->idle_thread_num) {
 if (d_->threads_cabinet.size() < d_->max_thread_num) {
 createWorker();
 } else {
 if (d_->undo_task_peak_num_ < d_->undo_tasks_cabinet.size())
 d_->undo_task_peak_num_ = d_->undo_tasks_cabinet.size();
 }
 }
 }

 LogDbg("create task %u", token.id());
 d_->cond_var.notify_one();

 return token;
}

ThreadExecutor::TaskToken ThreadPool::execute(const NonReturnFunc &backend_task, const NonReturnFunc &main_cb, int prio)
{
 NonReturnFunc backend_task_copy(backend_task);
 NonReturnFunc main_cb_copy(main_cb);
 return execute(std::move(backend_task_copy), std::move(main_cb_copy), prio);
}

ThreadExecutor::TaskToken ThreadPool::execute(NonReturnFunc &&backend_task, int prio)
{
 return execute(std::move(backend_task), nullptr, prio);
}

ThreadExecutor::TaskToken ThreadPool::execute(const NonReturnFunc &backend_task, int prio)
{
 return execute(backend_task, nullptr, prio);
}

ThreadExecutor::TaskToken ThreadPool::execute(NonReturnFunc &&backend_task)
{
 return execute(std::move(backend_task), 0);
}

ThreadExecutor::TaskToken ThreadPool::execute(const NonReturnFunc &backend_task)
{
 return execute(backend_task, 0);
}

ThreadExecutor::TaskToken ThreadPool::execute(NonReturnFunc &&backend_task, NonReturnFunc &&main_cb)
{
 return execute(std::move(backend_task), std::move(main_cb), 0);
}

ThreadExecutor::TaskToken ThreadPool::execute(const NonReturnFunc &backend_task, const NonReturnFunc &main_cb)
{
 return execute(backend_task, main_cb, 0);
}

ThreadExecutor::TaskStatus ThreadPool::getTaskStatus(TaskToken task_token) const
{
 std::lock_guard<std::mutex> lg(d_->lock);

 if (d_->undo_tasks_cabinet.at(task_token) != nullptr)
 return TaskStatus::kWaiting;

 if (d_->doing_tasks_token.find(task_token) != d_->doing_tasks_token.end())
 return TaskStatus::kExecuting;

 return TaskStatus::kNotFound;
}

ThreadExecutor::CancelResult ThreadPool::cancel(TaskToken token)
{
 RECORD_SCOPE();
 std::lock_guard<std::mutex> lg(d_->lock);

 //! 如果正在执行
 if (d_->doing_tasks_token.find(token) != d_->doing_tasks_token.end())
 return CancelResult::kExecuting; //! 返回正在执行

 //! 从高优先级向低优先级遍历,找出优先级最高的任务
 for (size_t i = 0; i < d_->undo_tasks_token.size(); ++i) {
 auto &tasks_token = d_->undo_tasks_token.at(i);
 if (!tasks_token.empty()) {
 auto iter = std::find(tasks_token.begin(), tasks_token.end(), token);
 if (iter != tasks_token.end()) {
 tasks_token.erase(iter);
 d_->task_pool.free(d_->undo_tasks_cabinet.free(token));
 return CancelResult::kSuccess;
 }
 }
 }

 return CancelResult::kNotFound; //! 返回没有找到
}

void ThreadPool::cleanup()
{
 RECORD_SCOPE();
 if (!d_->is_ready)
 return;

 std::vector<std::thread*> thread_vec;
 {
 std::lock_guard<std::mutex> lg(d_->lock);
 //! 清空task中的任务
 for (size_t i = 0; i < d_->undo_tasks_token.size(); ++i) {
 auto &tasks_token = d_->undo_tasks_token.at(i);
 while (!tasks_token.empty()) {
 auto token = tasks_token.front();
 d_->task_pool.free(d_->undo_tasks_cabinet.free(token));
 tasks_token.pop_front();
 }
 }

 //! 将threads_cabinet中的线程搬到thread_vec
 thread_vec.reserve(d_->threads_cabinet.size());
 d_->threads_cabinet.foreach(
 [&](std::thread *t) {
 thread_vec.push_back(t);
 }
 );
 d_->threads_cabinet.clear();
 }

 d_->all_threads_stop_flag = true;
 d_->cond_var.notify_all();

 //! 等待所有的线程退出
 for (auto t : thread_vec) {
 t->join();
 delete t;
 }

 d_->is_ready = false;
}

ThreadPool::Snapshot ThreadPool::snapshot() const
{
 Snapshot ss;
 std::lock_guard<std::mutex> lg(d_->lock);

 ss.idle_thread_num = d_->idle_thread_num;
 ss.thread_num = d_->threads_cabinet.size();
 ss.doing_task_num = d_->doing_tasks_token.size();
 for (size_t i = 0; i < THREAD_POOL_PRIO_SIZE; ++i)
 ss.undo_task_num[i] = d_->undo_tasks_token[i].size();
 ss.undo_task_peak_num = d_->undo_task_peak_num_;

 return ss;
}

void ThreadPool::threadProc(ThreadToken thread_token)
{
 bool let_main_loop_join_me = false;

 LogDbg("thread %u start", thread_token.id());

 while (true) {
 Task* item = nullptr;
 {
 std::unique_lock<std::mutex> lk(d_->lock);

 /**
 * 如果当前空闲的线程数量大于等于未被领取的任务数,且当前的线程个数已超过长驻线程数,
 * 说明线程数据已满足现有要求则退出当前线程
 */
 if ((d_->idle_thread_num >= d_->undo_tasks_cabinet.size()) && (d_->threads_cabinet.size() > d_->min_thread_num)) {
 LogDbg("thread %u will exit, no more work.", thread_token.id());
 let_main_loop_join_me = true;
 break;
 }

 //! 等待任务
 ++d_->idle_thread_num;
 d_->cond_var.wait(lk, std::bind(&ThreadPool::shouldThreadExitWaiting, this));
 --d_->idle_thread_num;

 /**
 * 有两种情况会从 cond_var.wait() 退出
 * 1. 任务队列中有任务需要执行时
 * 2. 线程池 cleanup() 时要求所有工作线程退出时
 *
 * 所以,下面检查 all_threads_stop_flag 看是不是请求退出
 */
 if (d_->all_threads_stop_flag) {
 LogDbg("thread %u will exit, stop flag.", thread_token.id());
 break;
 }

 item = popOneTask(); //! 从任务队列中取出优先级最高的任务
 }

 //! 后面就是去执行任务,不需要再加锁了
 if (item != nullptr) {
 RECORD_SCOPE();
 {
 std::lock_guard<std::mutex> lg(d_->lock);
 d_->doing_tasks_token.insert(item->token);
 }

 LogDbg("thread %u pick task %u", thread_token.id(), item->token.id());

 auto exec_time_point = Clock::now();
 auto wait_time_cost = exec_time_point - item->create_time_point;

 {
 RECORD_SCOPE();
 CatchThrow(item->backend_task, "tbox::eventx::ThreadPool", true);
 }

 auto exec_time_cost = Clock::now() - exec_time_point;

 LogDbg("thread %u finish task %u, cost %" PRIu64 " + %" PRIu64 " us",
 thread_token.id(), item->token.id(),
 wait_time_cost.count() / 1000,
 exec_time_cost.count() / 1000);

 /**
 * 有时在妥托给WorkThread执行动作时,会在lamda中捕获智能指针,它所指向的
 * 对象的析构函数是有动作的,如:http的sp_ctx要在析构中发送HTTP回复,如果
 * 析构函数在子线程中执行,则会出现不希望见到的多线程竞争。为此,我们在main_cb
 * 中也让它持有这个智能指针,希望智能指针所指的对象只在主线程中析构。
 *
 * 为了保证main_cb中的持有的对象能够在main_loop线程中被析构,
 * 所以这里要先task_pool.free(),然后再runInLoop(std::move(main_cpp))
 */

 auto main_cb = std::move(item->main_cb);

 {
 std::lock_guard<std::mutex> lg(d_->lock);
 d_->doing_tasks_token.erase(item->token);
 d_->task_pool.free(item);
 }

 if (main_cb) {
 RECORD_SCOPE();
 d_->wp_loop->runInLoop(std::move(main_cb), "ThreadPool::threadProc, invoke main_cb");
 }
 }
 }

 LogDbg("thread %u exit", thread_token.id());

 if (let_main_loop_join_me) {
 //! 则将线程取出来,交给main_loop去join(),然后delete
 std::unique_lock<std::mutex> lk(d_->lock);

 auto t = d_->threads_cabinet.free(thread_token);
 TBOX_ASSERT(t != nullptr);
 d_->wp_loop->runInLoop(
 [t]{ t->join(); delete t; },
 "ThreadPool::threadProc, join and delete it"
 );
 //! 这个操作放到最后来做是为了减少主线程join()的等待时长
 }
}

bool ThreadPool::createWorker()
{
 RECORD_SCOPE();
 ThreadToken thread_token = d_->threads_cabinet.alloc();
 auto *new_thread = new std::thread(std::bind(&ThreadPool::threadProc, this, thread_token));
 if (new_thread != nullptr) {
 d_->threads_cabinet.update(thread_token, new_thread);
 LogDbg("create thread %u", thread_token.id());
 return true;

 } else {
 LogErr("new thread fail");
 return false;
 }
}

bool ThreadPool::shouldThreadExitWaiting() const
{
 if (d_->all_threads_stop_flag)
 return true;

 for (size_t i = 0; i < d_->undo_tasks_token.size(); ++i) {
 const auto &tasks_token = d_->undo_tasks_token.at(i);
 if (!tasks_token.empty()) {
 return true;
 }
 }

 return false;
}

ThreadPool::Task* ThreadPool::popOneTask()
{
 //! 从高优先级向低优先级遍历,找出优先级最高的任务
 for (size_t i = 0; i < d_->undo_tasks_token.size(); ++i) {
 auto &tasks_token = d_->undo_tasks_token.at(i);
 if (!tasks_token.empty()) {
 TaskToken token = tasks_token.front();
 tasks_token.pop_front();
 return d_->undo_tasks_cabinet.free(token);
 }
 }
 return nullptr;
}

}
}