#pragma once
#include <atomic>

namespace kpn {

// ── SubmitGate ────────────────────────────────────────────────────────────────
//
// Decides, for one node, whether a wake must turn into a scheduler submission.
// Exactly one firing of a node may be in flight at a time, and a wake that
// arrives while one is already in flight must not be lost — it has to be
// honoured when that firing finishes, or the node sleeps holding work.
//
// 9c5ce5f wrote this as two independent atomics: queued_ said a firing was in
// flight, wake_pending_ recorded a wake that arrived during one. That cannot be
// made correct, because the release side has to read and write both, and a wake
// can land between the two operations:
//
//   producer (try_submit)              worker (release_and_recheck)
//   ------------------------           ----------------------------
//   CAS reads queued_ == true, fails
//                                      queued_.store(false)
//                                      wake_pending_.exchange(false) -> false
//   wake_pending_.store(true)
//
// End state: queued_ false, wake_pending_ true, nothing running and nothing
// scheduled. The node sleeps with a wake outstanding, which is precisely the
// invariant that commit set out to establish. It is not a memory-ordering
// subtlety — the interleaving above holds under seq_cst.
//
// It survived because every caller happened to follow release_and_recheck()
// with a level re-check (on_input_ready(), or outputs_have_space() on the
// parked path), which rediscovers the state a lost wake would have signalled.
// That is a property of the call sites, not of the mechanism, and any new early
// return that forgets the re-check turns it back into a hang.
//
// One atomic with three states makes the race unrepresentable: "idle" and "wake
// outstanding" are the same variable, so no interleaving can produce both.
//
//   Idle       nothing in flight
//   Queued     a firing is in flight or queued; no wake since it was claimed
//   QueuedWake a firing is in flight or queued, and a wake arrived meanwhile
//
class SubmitGate {
public:
    /// Register a wake. Returns true when the caller must submit the node;
    /// false when a firing is already in flight and the wake has been recorded
    /// against it instead.
    bool claim() noexcept {
        int cur = state_.load(std::memory_order_acquire);
        for (;;) {
            if (cur == kIdle) {
                if (state_.compare_exchange_weak(cur, kQueued,
                        std::memory_order_acq_rel, std::memory_order_acquire))
                    return true;
            } else if (cur == kQueued) {
                if (state_.compare_exchange_weak(cur, kQueuedWake,
                        std::memory_order_acq_rel, std::memory_order_acquire))
                    return false;
            } else {
                return false;            // a wake is already recorded
            }
        }
    }

    /// End the in-flight firing. Returns true when a wake arrived during it and
    /// the caller must submit again — in which case the gate stays claimed, so
    /// the node is handed straight from one firing to the next and is never
    /// momentarily idle with work outstanding. Returns false when the node is
    /// now idle.
    bool release() noexcept {
        int cur = state_.load(std::memory_order_acquire);
        for (;;) {
            if (cur == kQueuedWake) {
                if (state_.compare_exchange_weak(cur, kQueued,
                        std::memory_order_acq_rel, std::memory_order_acquire))
                    return true;
            } else {
                // kQueued, or kIdle if a stop already forced the gate down.
                if (state_.compare_exchange_weak(cur, kIdle,
                        std::memory_order_acq_rel, std::memory_order_acquire))
                    return false;
            }
        }
    }

    /// Drop the claim and any recorded wake. For stop paths only: honouring a
    /// wake there would resubmit a dead node.
    void force_idle() noexcept { state_.store(kIdle, std::memory_order_release); }

    bool queued()       const noexcept { return state_.load(std::memory_order_relaxed) != kIdle; }
    bool wake_pending() const noexcept { return state_.load(std::memory_order_relaxed) == kQueuedWake; }

private:
    static constexpr int kIdle       = 0;
    static constexpr int kQueued     = 1;
    static constexpr int kQueuedWake = 2;

    std::atomic<int> state_{kIdle};
};

} // namespace kpn