// Copyright © 2023 Apple Inc.

#include <sstream>

#include <pybind11/pybind11.h>

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

namespace py = pybind11;
using namespace py::literals;
using namespace mlx::core;

void init_stream(py::module_& m) {
  py::class_<Stream>(m, "Stream")
      .def(py::init<int, Device>(), "index"_a, "device"_a)
      .def_readonly("device", &Stream::device)
      .def(
          "__repr__",
          [](const Stream& s) {
            std::ostringstream os;
            os << s;
            return os.str();
          })
      .def("__eq__", [](const Stream& s1, const Stream& s2) {
        return s1 == s2;
      });

  py::implicitly_convertible<Device::DeviceType, Device>();

  m.def("default_stream", &default_stream, "device"_a);
  m.def("set_default_stream", &set_default_stream, "stream"_a);
  m.def("new_stream", &new_stream, "device"_a);
}