// SPDX-License-Identifier: MPL-2.0
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <pybind11/stl.h>
#include <pybind11/complex.h>
#include "vibeqc/tensor/distributed_cpu.hpp"
namespace py=pybind11;
using namespace vibeqc::tensor;
using namespace vibeqc::tensor::distributed;
namespace {
void array_check(const py::array& a,const py::dtype& dtype) {
    require(a.dtype().is(dtype) && a.ndim()==1 && (a.flags()&py::array::c_style),"expected contiguous one-dimensional array without conversion");
}
py::array bytes_view(const std::shared_ptr<CpuTransferSession>& s,const std::uint8_t* p,std::size_t n) {
    py::array a(py::dtype::of<std::uint8_t>(),{n},{sizeof(std::uint8_t)},p,py::cast(s));
    a.attr("setflags")(false);return a;
}
}
void bind_tensor_distributed(py::module_& parent) {
    auto m=parent.def_submodule("_tensor_cpu","Bounded experimental CPU tensor transfer");
    py::enum_<Kind>(m,"Kind").value("domain",Kind::domain).value("tensor",Kind::tensor).value("source",Kind::source)
        .value("metric",Kind::metric).value("frame",Kind::frame).value("operation",Kind::operation);
#define FIELD(T,N) .def_readwrite(#N,&T::N)
    py::class_<Id>(m,"Id").def(py::init<>()) FIELD(Id,namespace_high) FIELD(Id,namespace_low) FIELD(Id,serial) FIELD(Id,kind);
    py::class_<Token>(m,"Token").def(py::init<>()) FIELD(Token,id) FIELD(Token,revision);
    py::class_<DomainRef>(m,"DomainRef").def(py::init<>()) FIELD(DomainRef,token) FIELD(DomainRef,parent)
        FIELD(DomainRef,parent_coordinate) FIELD(DomainRef,extent) FIELD(DomainRef,begin) FIELD(DomainRef,step)
        FIELD(DomainRef,dependent) FIELD(DomainRef,selected) FIELD(DomainRef,coordinates);
    py::class_<Descriptor>(m,"Descriptor").def(py::init<>()) FIELD(Descriptor,id) FIELD(Descriptor,schema)
        FIELD(Descriptor,structure_revision) FIELD(Descriptor,source) FIELD(Descriptor,metric) FIELD(Descriptor,frame)
        .def_property("axes",[](const Descriptor& d){return std::vector<DomainRef>(d.axes.begin(),d.axes.begin()+d.rank);},
            [](Descriptor& d,const std::vector<DomainRef>& axes){require(axes.size()<=d.axes.size(),"rank exceeds bound");d.rank=axes.size();std::copy(axes.begin(),axes.end(),d.axes.begin());});
    py::class_<Placement>(m,"Placement").def(py::init<>()) FIELD(Placement,ranks) FIELD(Placement,active)
        FIELD(Placement,executor_shift) FIELD(Placement,output_owner) FIELD(Placement,epoch);
    py::class_<Operation>(m,"Operation").def(py::init<>()) FIELD(Operation,id) FIELD(Operation,value_epoch)
        FIELD(Operation,left) FIELD(Operation,right) FIELD(Operation,output) FIELD(Operation,placement)
        FIELD(Operation,adjoint) FIELD(Operation,work_limit);
#undef FIELD
    py::class_<Published>(m,"Published").def_property_readonly("size",&Published::size).def("at",&Published::at)
        .def_property_readonly("descriptor",[](const Published& p){return Descriptor(p.descriptor());})
        .def_property_readonly("value_epoch",[](const Published& p){return p.value().epoch;});
    py::class_<CpuTransferSession,std::shared_ptr<CpuTransferSession>>(m,"Session")
        .def(py::init<Operation,std::size_t,std::size_t,std::size_t>())
        .def_static("required_bytes",&CpuTransferSession::required_bytes)
        .def_property_readonly("count",&CpuTransferSession::count)
        .def_property_readonly("ranks",&CpuTransferSession::ranks)
        .def_property_readonly("output_owner",&CpuTransferSession::output_owner)
        .def_property_readonly("phase_id",&CpuTransferSession::phase_id)
        .def_property_readonly("epoch",&CpuTransferSession::epoch)
        .def_property_readonly("packet_bytes",&CpuTransferSession::packet_bytes)
        .def_property_readonly("live",&CpuTransferSession::live)
        .def_property_readonly("peak",&CpuTransferSession::peak)
        .def_property_readonly("pending",&CpuTransferSession::pending)
        .def_property_readonly("failed",&CpuTransferSession::failed)
        .def_property_readonly("manifest",[](std::shared_ptr<CpuTransferSession> s){const auto& h=s->manifest();return bytes_view(s,h.bytes.data(),h.size);})
        .def("source",&CpuTransferSession::source)
        .def("verify_manifest",[](CpuTransferSession& s,py::array a){array_check(a,py::dtype::of<std::uint8_t>());s.verify_manifest(static_cast<const std::uint8_t*>(a.data()),a.size());})
        .def("produce",[](std::shared_ptr<CpuTransferSession> s,std::size_t p,std::uint64_t g,py::array l,py::array r){
            array_check(l,py::dtype::of<Complex>());array_check(r,py::dtype::of<Complex>());
            return bytes_view(s,s->produce(p,g,static_cast<const Complex*>(l.data()),l.size(),static_cast<const Complex*>(r.data()),r.size()),s->packet_bytes());})
        .def("consume",[](CpuTransferSession& s,std::size_t p,std::uint64_t g,py::array a){array_check(a,py::dtype::of<std::uint8_t>());s.consume(p,g,static_cast<const std::uint8_t*>(a.data()),a.size());})
        .def("retire",&CpuTransferSession::retire).def("ready",&CpuTransferSession::ready)
        .def("publish",&CpuTransferSession::publish).def("release_safe",&CpuTransferSession::release_safe,py::arg("success")=false);
}