#pragma once

#include "mlx/dtype.h"
#include "mlx/utils.h"

#include "rust/cxx.h"

namespace mlx_cxx {
    // TODO: add binding to print format?

    struct StreamOrDevice {
        enum class Tag: uint8_t {
            Default,
            Stream,
            Device,
        };

        union Payload {
            std::monostate default_payload;
            mlx::core::Stream stream;
            mlx::core::Device device;
        };

        Tag tag = Tag::Default;
        Payload payload = Payload{ std::monostate{} };

        std::variant<std::monostate, mlx::core::Stream, mlx::core::Device> to_variant();
    };

    mlx::core::Dtype result_type(rust::Slice<const std::unique_ptr<mlx::core::array>> arrays);

    std::unique_ptr<std::vector<int>> broadcast_shapes(
        const std::vector<int>& s1,
        const std::vector<int>& s2);

    bool is_same_shape(rust::Slice<const std::unique_ptr<mlx::core::array>> arrays);

    // template<typename T>
    // void push_opaque(
    //     std::vector<T>& vec,
    //     std::unique_ptr<T> item);

    // template<typename T>
    // std::unique_ptr<T> pop_opaque(
    //     std::vector<T>& vec);

    // template<typename T>
    // std::unique_ptr<std::vector<T>> std_vec_from_slice(
    //     rust::Slice<const std::unique_ptr<T>> slice);
    
    template<typename T>
    void push_opaque(
        std::vector<T>& vec,
        std::unique_ptr<T> item) {
        vec.push_back(*item);
    }

    template<typename T>
    std::unique_ptr<std::vector<T>> std_vec_from_slice(
        rust::Slice<const std::unique_ptr<T>> slice) {
        std::vector<T> vec;
        for (auto& item : slice) {
            vec.push_back(*item);
        }
        return std::make_unique<std::vector<T>>(vec);
    }
}