#include <spdlog/spdlog.h>
#include "./request_scheduler.hpp"

namespace lsp::scheduler
{
    RequestScheduler::RequestScheduler(size_t concurrency) :
        executor_(concurrency)
    {
        spdlog::info("RequestScheduler initialized with {} threads", concurrency);
    }

    RequestScheduler::~RequestScheduler()
    {
        WaitAll();
    }

    void RequestScheduler::Submit(const std::string& request_id, TaskFunc task)
    {
        auto context = std::make_shared<TaskContext>();

        {
            std::lock_guard<std::mutex> lock(mutex_);

            // 取消旧任务
            auto it = running_tasks_.find(request_id);
            if (it != running_tasks_.end())
            {
                it->second->cancelled.store(true);
            }

            running_tasks_[request_id] = context;
        }

        executor_.async([this, request_id, task = std::move(task), context]() {
            try
            {
                if (context->cancelled.load())
                {
                    spdlog::debug("Task {} was cancelled", request_id);
                    return;
                }

                auto result = task();

                if (!context->cancelled.load() && result)
                {
                    SendResponse(*result);
                }
            }
            catch (const std::exception& e)
            {
                spdlog::error("Task {} failed: {}", request_id, e.what());
            }

            // 清理
            {
                std::lock_guard<std::mutex> lock(mutex_);
                running_tasks_.erase(request_id);
            }
        });
    }

    bool RequestScheduler::Cancel(const std::string& request_id)
    {
        std::lock_guard<std::mutex> lock(mutex_);

        auto it = running_tasks_.find(request_id);
        if (it != running_tasks_.end())
        {
            it->second->cancelled.store(true);
            return true;
        }

        return false;
    }

    void RequestScheduler::SetResponseCallback(ResponseCallback callback)
    {
        std::lock_guard<std::mutex> lock(mutex_);
        response_callback_ = std::move(callback);
    }

    void RequestScheduler::WaitAll()
    {
        executor_.wait_for_all();
    }

    void RequestScheduler::SendResponse(const std::string& response)
    {
        ResponseCallback callback;
        {
            std::lock_guard<std::mutex> lock(mutex_);
            callback = response_callback_;
        }

        if (callback)
            callback(response);
        else
            spdlog::error("No response callback set!");
    }
}
