#pragma once

#include <atomic>
#include <memory>
#include <utility>

namespace agent_smith
{

/// A copyable, cooperative signal that an in-progress operation should stop.
class CancellationToken
{
public:
    /// Create a token that is never cancelled.
    CancellationToken() = default;

    /// Return whether the associated cancellation source has been cancelled.
    bool cancelled() const noexcept
    {
        return (state_ != nullptr && state_->load()) ||
               (parent_ != nullptr && parent_->cancelled());
    }

private:
    friend class CancellationSource;

    explicit CancellationToken(std::shared_ptr<std::atomic_bool> state)
        : state_(std::move(state))
    {
    }

    std::shared_ptr<std::atomic_bool> state_;
    std::shared_ptr<const CancellationToken> parent_;
};

/// Owns and triggers a cancellation signal shared with its tokens.
class CancellationSource
{
public:
    /// Create a source with an initially non-cancelled token.
    CancellationSource() : state_(std::make_shared<std::atomic_bool>(false)) {}

    /// Create a source whose tokens also observe a parent cancellation token.
    explicit CancellationSource(CancellationToken parent)
        : CancellationSource()
    {
        parent_ = std::make_shared<CancellationToken>(std::move(parent));
    }

    /// Return a token that observes this source's cancellation state.
    CancellationToken token() const
    {
        CancellationToken result(state_);
        result.parent_ = parent_;
        return result;
    }

    /// Request cooperative cancellation for every token from this source.
    void cancel() noexcept
    {
        state_->store(true);
    }

private:
    std::shared_ptr<std::atomic_bool> state_;
    std::shared_ptr<const CancellationToken> parent_;
};

} // namespace agent_smith