#include <ZividPython/ReleasableArray2D.h>
#include <ZividPython/ReleasablePointCloud.h>

#include <Zivid/Color.h>
#include <Zivid/PointCloud.h>

#include <pybind11/pybind11.h>

namespace py = pybind11;

namespace
{
#pragma pack(push)
    struct PointXYZColorRGBA
    {
        float x, y, z;
        uint8_t r, g, b, a;
    };

    struct PointXYZColorBGRA
    {
        float x, y, z;
        uint8_t b, g, r, a;
    };
#pragma pack(pop)

    template<typename NativeType, typename WrapperType, size_t depth>
    py::buffer_info pointCloudDataBuffer(ZividPython::ReleasableArray2D<NativeType> &arrayWrapper)
    {
        static_assert(depth > 0, "Depth of array must be one or higher");
        static_assert(sizeof(WrapperType) * depth == sizeof(NativeType), "Unexpected data format");

        constexpr py::ssize_t dim = (depth == 1) ? 2 : 3;
        const auto shape = { arrayWrapper.impl().height(), arrayWrapper.impl().width(), depth };
        const auto strides = { sizeof(WrapperType) * arrayWrapper.impl().width() * depth,
                               sizeof(WrapperType) * depth,
                               sizeof(WrapperType) };
        auto *dataPtr = static_cast<void *>(const_cast<NativeType *>(arrayWrapper.impl().data()));
        return py::buffer_info(dataPtr,
                               sizeof(WrapperType),
                               py::format_descriptor<WrapperType>::format(),
                               dim,
                               std::vector<py::ssize_t>(shape.begin(), shape.begin() + dim),
                               std::vector<py::ssize_t>(strides.begin(), strides.begin() + dim));
    }

    template<typename NativeType>
    std::unique_ptr<ZividPython::ReleasableArray2D<NativeType>> pointCloudDataCopier(
        ZividPython::ReleasablePointCloud &pointCloud)
    {
        return std::make_unique<ZividPython::ReleasableArray2D<NativeType>>(pointCloud.impl().copyData<NativeType>());
    }

    template<typename NativeType>
    void wrapImageClass(pybind11::class_<ZividPython::ReleasableArray2D<NativeType>> pyClass)
    {
        using WrapperType = uint8_t;
        static_assert(std::is_same_v<WrapperType, decltype(NativeType::r)>);
        static_assert(std::is_same_v<WrapperType, decltype(NativeType::g)>);
        static_assert(std::is_same_v<WrapperType, decltype(NativeType::b)>);
        static_assert(std::is_same_v<WrapperType, decltype(NativeType::a)>);

        pyClass.def(py::init<>(&pointCloudDataCopier<NativeType>))
            .def_buffer(pointCloudDataBuffer<NativeType, WrapperType, 4>);
    }
} // namespace

namespace ZividPython
{
    void wrapClass(pybind11::class_<ReleasableArray2D<Zivid::SNR>> pyClass)
    {
        using WrapperType = float;
        using NativeType = Zivid::SNR;
        static_assert(std::is_same_v<WrapperType, decltype(NativeType::value)>);

        pyClass.def(py::init<>(&pointCloudDataCopier<NativeType>))
            .def_buffer(pointCloudDataBuffer<NativeType, WrapperType, 1>);
    }

    void wrapClass(pybind11::class_<ReleasableArray2D<Zivid::ColorRGBA>> pyClass)
    {
        wrapImageClass<Zivid::ColorRGBA>(pyClass);
    }

    void wrapClass(pybind11::class_<ReleasableArray2D<Zivid::ColorBGRA>> pyClass)
    {
        wrapImageClass<Zivid::ColorBGRA>(pyClass);
    }

    void wrapClass(pybind11::class_<ReleasableArray2D<Zivid::ColorSRGB>> pyClass)
    {
        wrapImageClass<Zivid::ColorSRGB>(pyClass);
    }

    void wrapClass(pybind11::class_<ReleasableArray2D<Zivid::NormalXYZ>> pyClass)
    {
        using WrapperType = float;
        using NativeType = Zivid::NormalXYZ;
        static_assert(std::is_same_v<WrapperType, decltype(NativeType::x)>);
        static_assert(std::is_same_v<WrapperType, decltype(NativeType::y)>);
        static_assert(std::is_same_v<WrapperType, decltype(NativeType::z)>);

        pyClass.def(py::init<>(&pointCloudDataCopier<NativeType>))
            .def_buffer(pointCloudDataBuffer<NativeType, WrapperType, 3>);
    }

    void wrapClass(pybind11::class_<ReleasableArray2D<Zivid::PointXYZ>> pyClass)
    {
        using WrapperType = float;
        using NativeType = Zivid::PointXYZ;
        static_assert(std::is_same_v<WrapperType, decltype(NativeType::x)>);
        static_assert(std::is_same_v<WrapperType, decltype(NativeType::y)>);
        static_assert(std::is_same_v<WrapperType, decltype(NativeType::z)>);

        pyClass.def(py::init<>(&pointCloudDataCopier<NativeType>))
            .def_buffer(pointCloudDataBuffer<NativeType, WrapperType, 3>);
    }

    void wrapClass(pybind11::class_<ReleasableArray2D<Zivid::PointXYZW>> pyClass)
    {
        using WrapperType = float;
        using NativeType = Zivid::PointXYZW;
        static_assert(std::is_same_v<WrapperType, decltype(NativeType::x)>);
        static_assert(std::is_same_v<WrapperType, decltype(NativeType::y)>);
        static_assert(std::is_same_v<WrapperType, decltype(NativeType::z)>);
        static_assert(std::is_same_v<WrapperType, decltype(NativeType::w)>);

        pyClass.def(py::init<>(&pointCloudDataCopier<NativeType>))
            .def_buffer(pointCloudDataBuffer<NativeType, WrapperType, 4>);
    }

    void wrapClass(pybind11::class_<ReleasableArray2D<Zivid::PointZ>> pyClass)
    {
        using WrapperType = float;
        using NativeType = Zivid::PointZ;

        static_assert(std::is_same_v<WrapperType, decltype(NativeType::z)>);

        pyClass.def(py::init<>(&pointCloudDataCopier<NativeType>))
            .def_buffer(pointCloudDataBuffer<NativeType, WrapperType, 1>);
    }

    void wrapClass(pybind11::class_<ReleasableArray2D<Zivid::PointXYZColorRGBA>> pyClass)
    {
        using WrapperType = PointXYZColorRGBA;
        using NativeType = Zivid::PointXYZColorRGBA;
        using NativePoint = decltype(NativeType::point);
        using NativeColor = decltype(NativeType::color);
        static_assert(std::is_same_v<decltype(WrapperType::x), decltype(NativePoint::x)>);
        static_assert(std::is_same_v<decltype(WrapperType::y), decltype(NativePoint::y)>);
        static_assert(std::is_same_v<decltype(WrapperType::z), decltype(NativePoint::z)>);
        static_assert(std::is_same_v<decltype(WrapperType::r), decltype(NativeColor::r)>);
        static_assert(std::is_same_v<decltype(WrapperType::g), decltype(NativeColor::g)>);
        static_assert(std::is_same_v<decltype(WrapperType::b), decltype(NativeColor::b)>);
        static_assert(std::is_same_v<decltype(WrapperType::a), decltype(NativeColor::a)>);
        // Additional layout checks for composite types
        static_assert(offsetof(NativeType, point) == offsetof(WrapperType, x));
        static_assert(offsetof(NativeType, color) == offsetof(WrapperType, r));
        static_assert(offsetof(NativePoint, x) == offsetof(WrapperType, x));
        static_assert(offsetof(NativePoint, y) == offsetof(WrapperType, y));
        static_assert(offsetof(NativePoint, z) == offsetof(WrapperType, z));
        static_assert(offsetof(NativeColor, r) == 0);
        static_assert(offsetof(NativeColor, g) == offsetof(WrapperType, g) - offsetof(WrapperType, r));
        static_assert(offsetof(NativeColor, b) == offsetof(WrapperType, b) - offsetof(WrapperType, r));
        static_assert(offsetof(NativeColor, a) == offsetof(WrapperType, a) - offsetof(WrapperType, r));

        PYBIND11_NUMPY_DTYPE(WrapperType, x, y, z, r, g, b, a);
        pyClass.def(py::init<>(&pointCloudDataCopier<NativeType>))
            .def_buffer(pointCloudDataBuffer<NativeType, WrapperType, 1>);
    }

    void wrapClass(pybind11::class_<ReleasableArray2D<Zivid::PointXYZColorBGRA>> pyClass)
    {
        using WrapperType = PointXYZColorBGRA;
        using NativeType = Zivid::PointXYZColorBGRA;
        using NativePoint = decltype(NativeType::point);
        using NativeColor = decltype(NativeType::color);
        static_assert(std::is_same_v<decltype(WrapperType::x), decltype(NativePoint::x)>);
        static_assert(std::is_same_v<decltype(WrapperType::y), decltype(NativePoint::y)>);
        static_assert(std::is_same_v<decltype(WrapperType::z), decltype(NativePoint::z)>);
        static_assert(std::is_same_v<decltype(WrapperType::b), decltype(NativeColor::b)>);
        static_assert(std::is_same_v<decltype(WrapperType::g), decltype(NativeColor::g)>);
        static_assert(std::is_same_v<decltype(WrapperType::r), decltype(NativeColor::r)>);
        static_assert(std::is_same_v<decltype(WrapperType::a), decltype(NativeColor::a)>);
        // Additional layout checks for composite types
        static_assert(offsetof(NativeType, point) == offsetof(WrapperType, x));
        static_assert(offsetof(NativeType, color) == offsetof(WrapperType, b));
        static_assert(offsetof(NativePoint, x) == offsetof(WrapperType, x));
        static_assert(offsetof(NativePoint, y) == offsetof(WrapperType, y));
        static_assert(offsetof(NativePoint, z) == offsetof(WrapperType, z));
        static_assert(offsetof(NativeColor, b) == 0);
        static_assert(offsetof(NativeColor, g) == offsetof(WrapperType, g) - offsetof(WrapperType, b));
        static_assert(offsetof(NativeColor, r) == offsetof(WrapperType, r) - offsetof(WrapperType, b));
        static_assert(offsetof(NativeColor, a) == offsetof(WrapperType, a) - offsetof(WrapperType, b));

        PYBIND11_NUMPY_DTYPE(WrapperType, x, y, z, b, g, r, a);
        pyClass.def(py::init<>(&pointCloudDataCopier<NativeType>))
            .def_buffer(pointCloudDataBuffer<NativeType, WrapperType, 1>);
    }
} // namespace ZividPython
