#pragma once
#include "diagnostics.hpp"
#include <atomic>
#include <condition_variable>
#include <functional>
#include <memory>
#include <mutex>
#include <optional>
#include <queue>
#include <thread>
#include <vector>

namespace kpn {

// ── IScheduler ────────────────────────────────────────────────────────────────

struct IScheduler {
    virtual ~IScheduler() = default;

    // Submit a task with an optional priority in [0, 1]. Higher = run sooner.
    virtual void submit(std::function<void()> task, float priority = 0.5f) = 0;

    // Start worker threads. Must be called before submit().
    virtual void start() = 0;

    // Halt: signal workers to exit and join them. Pending tasks are discarded.
    virtual void stop() = 0;

    // Drain: block until all in-flight tasks complete. Workers keep running.
    virtual void drain() = 0;
};

// ── ThreadPool ────────────────────────────────────────────────────────────────
//
// Work-stealing thread pool with per-thread priority queues.
//
// Each worker owns a priority_queue (max-heap by priority, FIFO within equal
// priority via sequence number). submit() distributes via round-robin. When a
// worker's queue is empty it tries to steal from the most-loaded peer using
// try_lock to avoid blocking; if no work is found it sleeps on a shared CV.
//
// total_ counts tasks submitted-but-not-completed (queued + executing).
// drain() waits until total_ == 0.

class ThreadPool : public IScheduler, public IPoolProbe {
public:
    explicit ThreadPool(std::size_t thread_count) : thread_count_(thread_count) {}

    ~ThreadPool() {
        if (!stopped_.load(std::memory_order_relaxed))
            stop();
    }

    void start() override {
        stopped_.store(false, std::memory_order_relaxed);
        queues_.clear();
        for (std::size_t i = 0; i < thread_count_; ++i)
            queues_.push_back(std::make_unique<WorkerQueue>());
        workers_.reserve(thread_count_);
        for (std::size_t i = 0; i < thread_count_; ++i)
            workers_.emplace_back([this, i] { worker_loop(i); });
    }

    void stop() override {
        stopped_.store(true, std::memory_order_seq_cst);
        for (auto& q : queues_) {
            std::lock_guard lock(q->mx);
            std::size_t discarded = q->pq.size();
            while (!q->pq.empty()) q->pq.pop();
            total_.fetch_sub(discarded, std::memory_order_relaxed);
        }
        cv_.notify_all();
        for (auto& t : workers_) if (t.joinable()) t.join();
        workers_.clear();
        queues_.clear();
    }

    void drain() override {
        std::unique_lock lock(drain_mx_);
        drain_cv_.wait(lock, [this] {
            return total_.load(std::memory_order_acquire) == 0;
        });
    }

    void submit(std::function<void()> task, float priority = 0.5f) override {
        std::size_t target = next_.fetch_add(1, std::memory_order_relaxed) % thread_count_;
        {
            std::lock_guard lock(queues_[target]->mx);
            queues_[target]->pq.push(
                {std::move(task), priority, seq_.fetch_add(1, std::memory_order_relaxed)});
        }
        total_.fetch_add(1, std::memory_order_relaxed);
        submitted_.fetch_add(1, std::memory_order_relaxed);
        cv_.notify_one();
    }

    std::size_t thread_count() const { return thread_count_; }

    // ── IPoolProbe ────────────────────────────────────────────────────────────

    PoolSnapshot snapshot(const std::string& name) const override {
        std::size_t a = active_.load(std::memory_order_relaxed);
        std::size_t t = total_.load(std::memory_order_relaxed);
        return {
            name, thread_count_,
            t > a ? t - a : 0,   // queued (approximate)
            a,                    // executing
            submitted_.load(std::memory_order_relaxed),
            completed_.load(std::memory_order_relaxed),
        };
    }

private:
    struct Task {
        std::function<void()> fn;
        float    priority;
        uint64_t seq;
        // max-heap: higher priority runs first; older task wins tie
        bool operator<(const Task& o) const {
            if (priority != o.priority) return priority < o.priority;
            return seq > o.seq;
        }
    };

    // Separate cache lines to prevent false sharing between adjacent queues.
    struct alignas(64) WorkerQueue {
        std::priority_queue<Task> pq;
        std::mutex                mx;
    };

    std::optional<std::function<void()>> try_pop(WorkerQueue& q) {
        std::lock_guard lock(q.mx);
        if (q.pq.empty()) return std::nullopt;
        auto fn = std::move(const_cast<Task&>(q.pq.top()).fn);
        q.pq.pop();
        return fn;
    }

    std::optional<std::function<void()>> try_steal(std::size_t thief) {
        // Find the most-loaded peer without blocking — racy peek is fine.
        std::size_t victim = thief, best = 0;
        for (std::size_t i = 0; i < queues_.size(); ++i) {
            if (i == thief) continue;
            std::unique_lock lk(queues_[i]->mx, std::try_to_lock);
            if (!lk) continue;
            std::size_t n = queues_[i]->pq.size();
            if (n > best) { best = n; victim = i; }
        }
        if (victim == thief) return std::nullopt;
        return try_pop(*queues_[victim]);
    }

    void execute(std::function<void()>& fn) {
        active_.fetch_add(1, std::memory_order_relaxed);
        fn();
        completed_.fetch_add(1, std::memory_order_relaxed);
        active_.fetch_sub(1, std::memory_order_relaxed);
        // Notify drain() if this was the last in-flight task.
        // acq_rel ensures the decrement is visible before any drain() load.
        if (total_.fetch_sub(1, std::memory_order_acq_rel) == 1)
            drain_cv_.notify_all();
    }

    void worker_loop(std::size_t id) {
        while (true) {
            if (auto fn = try_pop(*queues_[id]))   { execute(*fn); continue; }
            if (auto fn = try_steal(id))            { execute(*fn); continue; }

            std::unique_lock lock(cv_mx_);
            cv_.wait(lock, [this] {
                return stopped_.load(std::memory_order_seq_cst)
                    || total_.load(std::memory_order_relaxed) > 0;
            });
            if (stopped_.load(std::memory_order_seq_cst)
                && total_.load(std::memory_order_relaxed) == 0)
                return;
        }
    }

    const std::size_t                          thread_count_;
    std::vector<std::unique_ptr<WorkerQueue>>  queues_;
    std::vector<std::thread>                   workers_;

    std::mutex              cv_mx_;
    std::condition_variable cv_;
    std::mutex              drain_mx_;
    std::condition_variable drain_cv_;

    std::atomic<bool>     stopped_{true};
    std::atomic<size_t>   total_{0};      // queued + executing
    std::atomic<size_t>   active_{0};     // executing only (for snapshot)
    std::atomic<size_t>   next_{0};       // round-robin submit cursor
    std::atomic<uint64_t> seq_{0};        // tie-break for equal-priority tasks
    std::atomic<uint64_t> submitted_{0};
    std::atomic<uint64_t> completed_{0};
};

} // namespace kpn