#pragma once
#include "diagnostics.hpp"
#include <atomic>
#include <condition_variable>
#include <functional>
#include <memory>
#include <mutex>
#include <shared_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 {
        // Under the lifecycle lock for the same reason stop() is: submit()
        // reads queues_ and this rebuilds it. A network starts its nodes one at
        // a time, and a node already started fires into the next one's channel,
        // whose push callback submits — so a submission can genuinely land
        // while another pool is still inside start(). ThreadSanitizer reports
        // it as a read at submit() against this write, and the consequence is
        // worse than a torn read: push_back can reallocate the vector under a
        // reader that has already indexed it.
        //
        // Queues are all constructed before any worker is spawned, which is
        // what keeps worker_loop's own queues_[id] out of this — it never takes
        // the lock, so holding it across the spawn cannot deadlock.
        std::unique_lock lk(lifecycle_mx_);
        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 {
        // Close the pool to new work before touching anything, and do it under
        // the lifecycle lock so no submit() is midway through indexing queues_.
        {
            std::unique_lock lk(lifecycle_mx_);
            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);
            queued_.fetch_sub(discarded, std::memory_order_relaxed);
        }
        // Lock cv_mx_ before notifying so the stop signal can't be lost in the
        // gap between a worker's predicate check and its wait() (see submit()).
        { std::lock_guard<std::mutex> lk(cv_mx_); }
        cv_.notify_all();
        // Join without the lock: a worker's task may call submit(), which takes
        // it shared, and holding it here would deadlock against that.
        for (auto& t : workers_) if (t.joinable()) t.join();
        // Destroying the queues is what submit() must never race. By now
        // stopped_ is published, so any submit() that acquires the lock after
        // this point returns without touching them.
        std::unique_lock lk(lifecycle_mx_);
        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 {
        // A submission can arrive after this pool has been stopped, and did so
        // by an ordinary route: a node's space callback fires from whichever
        // thread drained the channel, which belongs to the *consumer*. Stop the
        // producer first — as a sources-first shutdown does — and the consumer
        // keeps draining its backlog, firing the producer's space callback into
        // a pool whose stop() has already run queues_.clear(). submit() then
        // indexed an empty vector: a segfault, reproducible about 12 runs in 20.
        //
        // The shared lock is what makes the check meaningful. Reading stopped_
        // alone leaves the window between the read and the indexing, which is
        // precisely where stop() clears the vector.
        std::shared_lock lk(lifecycle_mx_);
        if (stopped_.load(std::memory_order_acquire) || queues_.empty()) {
            rejected_.fetch_add(1, std::memory_order_relaxed);
            return;
        }
        // B9 — submit-to-self affinity. Round-robin hands every task to a
        // *different* worker, and on a pool of two or more that worker is
        // asleep, so each dispatch pays a futex wake: measured 182 ns/dispatch
        // on a 1-thread pool against 3197 ns on 20 threads, with voluntary
        // context switches per task rising 0.00 -> 1.19 in step.
        //
        // A submission originating on one of *our own* workers goes to that
        // worker's queue instead. It is about to return to worker_loop and
        // try_pop its own queue, so the work is already there and nothing
        // sleeps — the property that makes ThreadPool(1) fast, extended to
        // any pool size. Imbalance is corrected by the existing try_steal.
        //
        // The pool identity check is load-bearing: a worker of pool A
        // submitting into pool B must not use A's index, which may exceed B's
        // thread_count_ or alias an unrelated queue. Nested networks do
        // exactly this.
        std::size_t target;
        if (tls_pool == this && tls_worker < thread_count_) {
            target = tls_worker;
        } else {
            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);
        queued_.fetch_add(1, std::memory_order_relaxed);
        submitted_.fetch_add(1, std::memory_order_relaxed);
        // Synchronize with worker_loop's predicate evaluation: taking cv_mx_
        // here guarantees a worker is either before its predicate check (and
        // will observe total_ > 0) or already blocked in wait() (and will be
        // woken). Without this, notify_one() can slip into the gap between the
        // worker's predicate check and its wait(), and be lost — a deadlock.
        //
        // Skipped entirely when no worker is parked. waiters_ is incremented
        // *before* wait() releases cv_mx_ and decremented after it returns,
        // both under that mutex, so a worker on its way to sleep is already
        // counted here. Reading zero therefore means no worker can be in
        // wait(), and there is nothing a notify could reach — as opposed to
        // reading zero because we raced one, which the mutex prevents.
        //
        // This is the hot path for an already-busy pool: with B9 the work is
        // in the local queue and the submitting worker will find it itself,
        // so the lock round-trip and notify were pure overhead. Measured 1.00
        // voluntary context switches per task before this, on a pool where
        // only one task is ever in flight.
        if (waiters_.load(std::memory_order_seq_cst) != 0) {
            { std::lock_guard<std::mutex> lk(cv_mx_); }
            cv_.notify_one();
        }
    }

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

    /// Submissions dropped because the pool was stopped. See rejected_.
    uint64_t rejected() const { return rejected_.load(std::memory_order_relaxed); }

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

    PoolSnapshot snapshot(const std::string& name) const override {
        std::size_t a = active_.load(std::memory_order_relaxed);
        return {
            name, thread_count_,
            queued_.load(std::memory_order_relaxed),   // queued (exact)
            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();
        queued_.fetch_sub(1, std::memory_order_relaxed);
        return fn;
    }

    std::optional<std::function<void()>> try_steal(std::size_t thief) {
        // Find the most-loaded peer without blocking — racy peek is fine.
        //
        // The threshold is >0: a peer holding a single task is a valid victim.
        //
        // Raising it to >1 — to stop a thief winning the race for a task its
        // owner just submitted to itself (B9) — deadlocks. `latency` mode
        // hangs at 12 and 20 threads: an external submit() round-robins one
        // task onto an idle worker's queue, and if that worker is parked, no
        // peer will take it because a queue of one is no longer stealable.
        // Nothing else is coming to wake it, so the pool sits forever.
        // Measured before reverting: it also made steady-state *worse*,
        // 2229 -> 4546 ns at 12 threads.
        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.
        // Lock drain_mx_ before notifying to avoid a lost wakeup against
        // drain()'s predicate check (same hazard as submit()/cv_mx_).
        if (total_.fetch_sub(1, std::memory_order_acq_rel) == 1) {
            { std::lock_guard<std::mutex> lk(drain_mx_); }
            drain_cv_.notify_all();
        }
    }

    void worker_loop(std::size_t id) {
        // Identify this thread as one of our workers, for B9's affinity check
        // in submit(). Restored on exit rather than merely cleared: a pool
        // whose worker runs a task that itself starts and stops a nested pool
        // would otherwise come back with its identity erased.
        ThreadPool* const prev_pool   = tls_pool;
        const std::size_t prev_worker = tls_worker;
        tls_pool   = this;
        tls_worker = id;
        struct Restore {
            ThreadPool* p; std::size_t w;
            ~Restore() { tls_pool = p; tls_worker = w; }
        } restore{prev_pool, prev_worker};

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

            // B5 (bounded spin before parking) was tried here and removed: it
            // does not pay. Swept at 50/200/1000 rounds on a 12-thread pool,
            // steady state went 2123 / 2230 / 2574 ns against 2229 ns without
            // it, and voluntary context switches per task stayed at ~0.97
            // throughout. The spin cannot catch what it is aimed at, because
            // a peer is woken the moment queued_ becomes non-zero — which
            // happens before this worker reaches the spin at all.
            std::unique_lock lock(cv_mx_);
            // Counted under cv_mx_ and before the predicate is evaluated, so
            // that a submit() which reads waiters_ == 0 can be certain this
            // worker is not about to block: to get here we already hold the
            // mutex that submit() must take to notify.
            waiters_.fetch_add(1, std::memory_order_seq_cst);
            cv_.wait(lock, [this] {
                return stopped_.load(std::memory_order_seq_cst)
                    || queued_.load(std::memory_order_relaxed) > 0;
            });
            waiters_.fetch_sub(1, std::memory_order_seq_cst);
            // Exit on queued_, not total_: waiting for total_ to reach zero
            // meant waiting for someone else's task to finish, which this
            // worker cannot help with and would spin through until it did.
            if (stopped_.load(std::memory_order_seq_cst)
                && queued_.load(std::memory_order_relaxed) == 0)
                return;
        }
    }

    /// Which pool, and which of its workers, the calling thread is — or
    /// nullptr on any thread that is not a pool worker. Read by submit() to
    /// decide whether a local push is safe (B9). inline so the header stays
    /// header-only.
    static inline thread_local ThreadPool* tls_pool   = nullptr;
    static inline thread_local std::size_t tls_worker = 0;

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

    /// Guards the lifetime of queues_/workers_ against a concurrent submit().
    /// Shared by submit, exclusive by stop, so submissions still run in
    /// parallel with each other.
    mutable std::shared_mutex lifecycle_mx_;

    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 (drain() waits on this)
    /// Queued only — never counts a task that is already executing.
    ///
    /// The wait predicate used total_, which includes running tasks, so while
    /// any one task ran every *other* worker's predicate was true: wait()
    /// returned instantly and the worker spun through try_pop / try_steal /
    /// wait at full speed, try_lock-ing every peer queue on each pass. One slow
    /// task therefore pinned every other core and contended the very mutexes
    /// the working thread needed. Sleeping requires "no work is *waiting*",
    /// which is this.
    std::atomic<size_t>   queued_{0};     // waiting to run
    std::atomic<size_t>   active_{0};     // executing only (for snapshot)
    std::atomic<size_t>   next_{0};       // round-robin submit cursor
    /// Workers currently inside cv_.wait(), maintained under cv_mx_. Lets
    /// submit() skip the lock round-trip and notify when nobody is parked.
    std::atomic<size_t>   waiters_{0};
    std::atomic<uint64_t> seq_{0};        // tie-break for equal-priority tasks
    std::atomic<uint64_t> submitted_{0};
    /// Submissions refused because the pool was already stopped. Not an error —
    /// teardown races are expected — but silence here would hide a node that
    /// keeps trying to run after its pool is gone.
    std::atomic<uint64_t> rejected_{0};
    std::atomic<uint64_t> completed_{0};
};

} // namespace kpn