#ifndef REDUCE_NODE_H
#define REDUCE_NODE_H

#include <deque>
#include <hgraph/nodes/nested_evaluation_engine.h>
#include <hgraph/nodes/nested_node.h>
#include <hgraph/types/tsd.h>

namespace hgraph {
    void register_reduce_node_with_nanobind(nb::module_ & m);

    struct ReduceNode;
    using reduce_node_s_ptr = std::shared_ptr<ReduceNode>;

    /**
     * C++ implementation of PythonReduceNodeImpl.
     * This implements TSD reduction using an inverted binary tree with inputs at the leaves
     * and the result at the root. The inputs bound to the leaves can be moved as nodes come and go.
     *
     * Non-templated: uses Value/TypeMeta for type-erased key handling.
     */
    struct ReduceNode final : NestedNode {
        // Use Value for type-erased key storage (same as TSD)
        using key_map_type = std::unordered_map<value::Value, std::tuple<int64_t, int64_t>,
                                                 ValueHash, ValueEqual>;

        ReduceNode(int64_t node_ndx, std::vector<int64_t> owning_graph_id, NodeSignature::s_ptr signature,
                   nb::dict scalars,
                   graph_builder_s_ptr nested_graph_builder, const std::tuple<int64_t, int64_t> &input_node_ids,
                   int64_t output_node_id);

        // Non-copyable due to move-only Value members
        ReduceNode(const ReduceNode&) = delete;
        ReduceNode& operator=(const ReduceNode&) = delete;

        std::unordered_map<int, graph_s_ptr> &nested_graphs();

        TimeSeriesDictInputImpl::ptr ts();

        time_series_reference_input_ptr zero();

        // Expose attributes to allow us to more easily inspect the state in Python
        // Can make debugging easier.
        const graph_s_ptr &nested_graph() const;

        const std::tuple<int64_t, int64_t> &input_node_ids() const;

        int64_t output_node_id() const;

        // Returns Python dict for inspection (converts Value keys to Python)
        nb::dict py_bound_node_indexes() const;

        const std::vector<std::tuple<int64_t, int64_t> > &free_node_indexes() const;

        void enumerate_nested_graphs(const std::function<void(const graph_s_ptr&)>& callback) const override;

        VISITOR_SUPPORT()

    protected:
        void initialise() override;

        void do_start() override;

        void do_stop() override;

        void dispose() override;

        void eval() override;

        void do_eval() override {
        };

        TimeSeriesOutput::s_ptr last_output();

        // Uses View for key iteration from TSD
        void add_nodes_from_views(const std::vector<value::View> &keys);

        void remove_nodes_from_views(const std::vector<value::View> &keys);

        void re_balance_nodes();

        void grow_tree();

        void shrink_tree();

        void bind_key_to_node(const value::View &key, const std::tuple<int64_t, int64_t> &ndx);

        void zero_node(const std::tuple<int64_t, int64_t> &ndx);

        void swap_node(const std::tuple<int64_t, int64_t> &src_ndx, const std::tuple<int64_t, int64_t> &dst_ndx);

        int64_t node_size() const;

        int64_t node_count() const;

        std::vector<node_s_ptr> get_node(int64_t ndx);

    private:
        graph_s_ptr nested_graph_;
        graph_builder_s_ptr nested_graph_builder_;
        std::tuple<int64_t, int64_t> input_node_ids_; // LHS index, RHS index
        int64_t output_node_id_;
        key_map_type bound_node_indexes_;
        std::vector<std::tuple<int64_t, int64_t> > free_node_indexes_; // List of (ndx, 0(lhs)|1(rhs)) tuples

        // The python code uses the fact that you can randomly add properties to a python object and tracks
        // if an input is bound to a key or not using _bound_to_key.
        // C++ does not do that, so we can track if the ts is bound to a key using a set.
        std::unordered_set<TimeSeriesInput *> bound_to_key_flags_;
    };
} // namespace hgraph

#endif  // REDUCE_NODE_H
