//
// ThreadPool.hpp - Priority-based thread pool with deadline support
//

#ifndef EAGLE0_THREADPOOL_HPP
#define EAGLE0_THREADPOOL_HPP

#include <atomic>
#include <chrono>
#include <condition_variable>
#include <functional>
#include <future>
#include <iostream>
#include <memory>
#include <mutex>
#include <queue>
#include <thread>
#include <vector>

namespace eagle0::common {

enum class TaskStatus { SUCCESS = 0, DEADLINE_EXCEEDED = 1, CANCELLED = 2 };

template<typename T>
struct TaskResult {
    T value;
    TaskStatus status;

    TaskResult() : value{}, status(TaskStatus::SUCCESS) {}
    TaskResult(T val) : value(std::move(val)), status(TaskStatus::SUCCESS) {}
    TaskResult(T val, TaskStatus stat) : value(std::move(val)), status(stat) {}

    // NO implicit conversion - this was causing infinite recursion
    // Use .value or .get() instead
    T get() const { return value; }

    bool succeeded() const { return status == TaskStatus::SUCCESS; }
    bool deadlineExceeded() const { return status == TaskStatus::DEADLINE_EXCEEDED; }
};

class ThreadPool {
public:
    using Clock = std::chrono::steady_clock;
    using TimePoint = Clock::time_point;

private:
    struct Task {
        std::function<void()> function;
        int priority;
        TimePoint deadline;
        bool has_deadline;

        Task(std::function<void()> f, int p, TimePoint d, bool has_d)
            : function(std::move(f)),
              priority(p),
              deadline(d),
              has_deadline(has_d) {}

        // Higher priority values and earlier deadlines have higher priority
        bool operator<(const Task& other) const {
            if (priority != other.priority) {
                return priority < other.priority;  // Lower priority values have lower priority in
                                                   // priority_queue
            }
            if (has_deadline && other.has_deadline) {
                return deadline > other.deadline;  // Later deadlines have lower priority
            }
            if (has_deadline && !other.has_deadline) {
                return false;  // Tasks with deadlines have higher priority
            }
            if (!has_deadline && other.has_deadline) {
                return true;  // Tasks without deadlines have lower priority
            }
            return false;  // Equal priority, no preference
        }
    };

    std::vector<std::thread> workers;
    std::priority_queue<Task> tasks;
    std::mutex queue_mutex;
    std::condition_variable condition;
    std::atomic<bool> stop{false};

public:
    explicit ThreadPool(size_t num_threads = std::thread::hardware_concurrency()) {
        for (size_t i = 0; i < num_threads; ++i) {
            workers.emplace_back([this] {
                while (true) {
                    Task task{nullptr, 0, TimePoint{}, false};
                    {
                        std::unique_lock<std::mutex> lock(queue_mutex);
                        condition.wait(lock, [this] { return stop.load() || !tasks.empty(); });

                        if (stop.load() && tasks.empty()) { return; }

                        if (!tasks.empty()) {
                            task = std::move(const_cast<Task&>(tasks.top()));
                            tasks.pop();
                        } else {
                            continue;
                        }
                    }

                    // Execute the task (deadline checking is now handled inside the task)
                    if (task.function) { task.function(); }
                }
            });
        }
    }

    // Enqueue a task with priority only
    template<class F, class... Args>
    auto enqueue(F&& f, Args&&... args, int priority = 0)
            -> std::future<TaskResult<std::invoke_result_t<F, Args...>>> {
        using return_type = std::invoke_result_t<F, Args...>;
        using result_type = TaskResult<return_type>;

        auto actualTask = std::bind(std::forward<F>(f), std::forward<Args>(args)...);

        auto task = std::make_shared<std::packaged_task<result_type()>>(
                [actualTask = std::move(actualTask)]() mutable -> result_type {
                    return result_type(actualTask());
                });

        std::future<result_type> result = task->get_future();

        {
            std::unique_lock<std::mutex> lock(queue_mutex);
            if (stop.load()) { throw std::runtime_error("enqueue on stopped ThreadPool"); }
            tasks.emplace([task]() { (*task)(); }, priority, TimePoint{}, false);
        }

        condition.notify_one();
        return result;
    }

    // Enqueue a task with priority and deadline
    template<class F, class... Args>
    auto enqueue_with_deadline(F&& f, Args&&... args, int priority, TimePoint deadline)
            -> std::future<TaskResult<std::invoke_result_t<F, Args...>>> {
        using return_type = std::invoke_result_t<F, Args...>;
        using result_type = TaskResult<return_type>;

        auto actualTask = std::bind(std::forward<F>(f), std::forward<Args>(args)...);

        auto task = std::make_shared<std::packaged_task<result_type()>>(
                [actualTask = std::move(actualTask), deadline]() mutable -> result_type {
                    if (Clock::now() > deadline) {
                        return result_type(return_type{}, TaskStatus::DEADLINE_EXCEEDED);
                    }
                    return result_type(actualTask());
                });

        std::future<result_type> result = task->get_future();

        {
            std::unique_lock<std::mutex> lock(queue_mutex);
            if (stop.load()) { throw std::runtime_error("enqueue on stopped ThreadPool"); }
            tasks.emplace([task]() { (*task)(); }, priority, deadline, true);
        }

        condition.notify_one();
        return result;
    }

    // Get current queue size (approximate, for monitoring)
    size_t queue_size() const {
        std::unique_lock<std::mutex> lock(const_cast<std::mutex&>(queue_mutex));
        return tasks.size();
    }

    // Get detailed queue information for debugging
    void debug_queue_state() const {
        std::unique_lock<std::mutex> lock(const_cast<std::mutex&>(queue_mutex));
        std::cout << "ThreadPool: Queue size: " << tasks.size() << '\n';
        if (!tasks.empty()) {
            // Create a copy to inspect priorities without modifying queue
            auto queue_copy = tasks;
            std::vector<int> priorities;
            while (!queue_copy.empty()) {
                priorities.push_back(queue_copy.top().priority);
                queue_copy.pop();
            }
            std::cout << "ThreadPool: Priorities in queue: ";
            for (int p : priorities) { std::cout << p << ' '; }
            std::cout << '\n';
        }
    }

    ~ThreadPool() {
        stop.store(true);
        condition.notify_all();
        for (std::thread& worker : workers) {
            if (worker.joinable()) { worker.join(); }
        }
    }
};

}  // namespace eagle0::common

#endif  // EAGLE0_THREADPOOL_HPP
