// Copyright Contributors to the OpenImageIO project.
// SPDX-License-Identifier: Apache-2.0
// https://github.com/AcademySoftwareFoundation/OpenImageIO

#ifndef PYOPENIMAGEIO_PY_OIIO_H
#define PYOPENIMAGEIO_PY_OIIO_H

// Python.h uses the 'register' keyword, don't warn about it being
// deprecated in C++17.
#if (__cplusplus >= 201703L && defined(__GNUC__)) || defined(__clang__)
#    pragma GCC diagnostic ignored "-Wregister"
#endif

// clang-format off
// Must include Python.h first to avoid certain warnings
#ifdef _POSIX_C_SOURCE
#  error "You must include Python.h (and therefore py_oiio.h) BEFORE anything that defines _POSIX_C_SOURCE"
#endif 
#include <Python.h>
// clang-format on

#include <memory>

// Avoid a compiler warning from a duplication in tiffconf.h/pyconfig.h
#undef SIZEOF_LONG

// Avoid a problem with copysign defined in pyconfig.h on Windows.
#ifdef copysign
#    undef copysign
#endif

#include <OpenImageIO/Imath.h>
#include <OpenImageIO/deepdata.h>
#include <OpenImageIO/imagebuf.h>
#include <OpenImageIO/imagecache.h>
#include <OpenImageIO/imageio.h>
#include <OpenImageIO/span.h>
#include <OpenImageIO/texture.h>
#include <OpenImageIO/typedesc.h>

#include <pybind11/numpy.h>
#include <pybind11/operators.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
namespace py = pybind11;


// Python3 is always unicode, so return a true str
#define PY_STR py::str



namespace pybind11 {
namespace detail {

    // This half casting support for numpy was all derived from discussions
    // here: https://github.com/pybind/pybind11/issues/1776

    // Similar to enums in `pybind11/numpy.h`. Determined by doing:
    // python3 -c 'import numpy as np; print(np.dtype(np.float16).num)'
    constexpr int NPY_FLOAT16 = 23;

    template<> struct npy_format_descriptor<half> {
        static pybind11::dtype dtype()
        {
            handle ptr = npy_api::get().PyArray_DescrFromType_(NPY_FLOAT16);
            return reinterpret_borrow<pybind11::dtype>(ptr);
        }
        static std::string format()
        {
            // following: https://docs.python.org/3/library/struct.html#format-characters
            return "e";
        }
        static constexpr auto name = _("float16");
    };

}  // namespace detail
}  // namespace pybind11



namespace PyOpenImageIO {

using namespace OIIO;

// clang-format off

void declare_imagespec (py::module& m);
void declare_imageinput (py::module& m);
void declare_imageoutput (py::module& m);
void declare_typedesc (py::module& m);
void declare_roi (py::module& m);
void declare_deepdata (py::module& m);
void declare_colorconfig (py::module& m);
void declare_imagecache (py::module& m);
void declare_imagebuf (py::module& m);
void declare_imagebufalgo (py::module& m);
void declare_paramvalue (py::module& m);
void declare_global (py::module& m);
void declare_wrap (py::module& m);
void declare_mipmpode (py::module& m);
void declare_interpmode (py::module& m);
void declare_textureopt (py::module& m);
void declare_texturesystem (py::module& m);

// bool PyProgressCallback(void*, float);
// object C_array_to_Python_array (const char *data, TypeDesc type, size_t size);
TypeDesc typedesc_from_python_array_code (string_view code);
// const char * python_array_code (TypeDesc format);  // unused


inline std::string
object_classname(const py::object& obj)
{
    return obj.attr("__class__").attr("__name__").cast<py::str>();
}



template<typename T> struct PyTypeForCType { };
template<> struct PyTypeForCType<int> { typedef py::int_ type; };
template<> struct PyTypeForCType<unsigned int> { typedef py::int_ type; };
template<> struct PyTypeForCType<short> { typedef py::int_ type; };
template<> struct PyTypeForCType<unsigned short> { typedef py::int_ type; };
template<> struct PyTypeForCType<char> { typedef py::int_ type; };
template<> struct PyTypeForCType<unsigned char> { typedef py::int_ type; };
template<> struct PyTypeForCType<int64_t> { typedef py::int_ type; };
template<> struct PyTypeForCType<uint64_t> { typedef py::int_ type; };
template<> struct PyTypeForCType<float> { typedef py::float_ type; };
template<> struct PyTypeForCType<half> { typedef py::float_ type; };
template<> struct PyTypeForCType<double> { typedef py::float_ type; };
template<> struct PyTypeForCType<const char*> { typedef PY_STR type; };
template<> struct PyTypeForCType<std::string> { typedef PY_STR type; };

// clang-format on


// Struct that holds OIIO style buffer info, constructed from
// py::buffer_info
struct oiio_bufinfo {
    TypeDesc format  = TypeUnknown;
    void* data       = nullptr;
    stride_t xstride = AutoStride, ystride = AutoStride, zstride = AutoStride;
    size_t size = 0;
    std::string error;

    // Just raw buffer, no idea what to expect, treat like a flat array.
    // Only works for "contiguous" buffers.
    oiio_bufinfo(const py::buffer_info& pybuf);

    // Expect a certain layout, figure out how to make sense of the buffer.
    oiio_bufinfo(const py::buffer_info& pybuf, int nchans, int width,
                 int height, int depth, int pixeldims);

    // Retrieve presumed contiguous data value index i.
    template<typename T> T dataval(size_t i) { return ((const T*)data)[i]; }
};



// Suck up a tuple of presumed T values into a vector<T>. Works for T any of
// int, float, string, and works if the Python container is any of tuple,
// list.
template<typename T, typename PYT>
inline bool
py_indexable_pod_to_stdvector(std::vector<T>& vals, const PYT& obj)
{
    OIIO_ASSERT(py::isinstance<py::tuple>(obj)
                || py::isinstance<py::list>(obj));
    bool ok             = true;
    const size_t length = py::len(obj);
    vals.reserve(length);
    for (size_t i = 0; i < length; ++i) {
        auto elem = obj[i];
        if (std::is_same<T, float>::value && py::isinstance<py::float_>(elem)) {
            vals.emplace_back(T(elem.template cast<float>()));
        } else if ((std::is_same<T, float>::value || std::is_same<T, int>::value)
                   && py::isinstance<py::int_>(elem)) {
            vals.emplace_back(T(elem.template cast<int>()));
        } else if (std::is_same<T, unsigned int>::value
                   && py::isinstance<py::int_>(elem)) {
            vals.emplace_back(T(elem.template cast<unsigned int>()));
        } else {
            // FIXME? Other cases?
            vals.emplace_back(T(42));
            ok = false;
        }
    }
    return ok;
}


// Partial specialization for reading strings
template<typename PYT>
inline bool
py_indexable_pod_to_stdvector(std::vector<std::string>& vals, const PYT& obj)
{
    OIIO_ASSERT(py::isinstance<py::tuple>(obj)
                || py::isinstance<py::list>(obj));
    bool ok             = true;
    const size_t length = py::len(obj);
    vals.reserve(length);
    for (size_t i = 0; i < length; ++i) {
        auto elem = obj[i];
        if (py::isinstance<py::str>(elem)) {
            vals.emplace_back(elem.template cast<py::str>());
        } else {
            // FIXME? Other cases?
            vals.emplace_back("");
            ok = false;
        }
    }
    return ok;
}


// Partial specialization for reading TypeDesc
template<typename PYT>
inline bool
py_indexable_pod_to_stdvector(std::vector<TypeDesc>& vals, const PYT& obj)
{
    OIIO_ASSERT(py::isinstance<py::tuple>(obj)
                || py::isinstance<py::list>(obj));
    bool ok             = true;
    const size_t length = py::len(obj);
    vals.reserve(length);
    for (size_t i = 0; i < length; ++i) {
        auto elem = obj[i];
        if (py::isinstance<TypeDesc>(elem)) {
            vals.emplace_back(*(elem.template cast<TypeDesc*>()));
        } else if (py::isinstance<TypeDesc::BASETYPE>(elem)) {
            vals.emplace_back(*(elem.template cast<TypeDesc::BASETYPE*>()));
        } else if (py::isinstance<py::str>(elem)) {
            vals.emplace_back(
                TypeDesc(std::string(elem.template cast<py::str>())));
        } else {
            // FIXME? Other cases?
            vals.emplace_back(TypeUnknown);
            ok = false;
        }
    }
    return ok;
}



template<typename T>
inline bool
py_scalar_pod_to_stdvector(std::vector<T>& vals, const py::object& obj)
{
    using pytype = typename PyTypeForCType<T>::type;
    vals.clear();
    if (py::isinstance<pytype>(obj)) {
        vals.emplace_back(obj.template cast<pytype>());
        return true;
    } else {
        return false;
    }
}

// float specialization -- accept ints, too
template<>
inline bool
py_scalar_pod_to_stdvector(std::vector<float>& vals, const py::object& obj)
{
    vals.clear();
    if (py::isinstance<py::float_>(obj)) {
        vals.emplace_back(obj.template cast<py::float_>());
        return true;
    } else if (py::isinstance<py::int_>(obj)) {
        int i = obj.template cast<py::int_>();
        vals.emplace_back((float)i);
        return true;
    } else {
        return false;
    }
}

// TypeDesc specialization
template<>
inline bool
py_scalar_pod_to_stdvector(std::vector<TypeDesc>& vals, const py::object& obj)
{
    vals.clear();
    if (py::isinstance<TypeDesc>(obj)) {
        vals.emplace_back(*(obj.template cast<TypeDesc*>()));
        return true;
    } else if (py::isinstance<TypeDesc::BASETYPE>(obj)) {
        vals.emplace_back(*(obj.template cast<TypeDesc::BASETYPE*>()));
        return true;
    } else if (py::isinstance<py::str>(obj)) {
        vals.emplace_back(TypeDesc(std::string(obj.template cast<py::str>())));
        return true;
    } else {
        return false;
    }
}



template<typename T>
inline bool
py_buffer_to_stdvector(std::vector<T>& vals, const py::buffer& obj)
{
    OIIO_DASSERT(py::isinstance<py::buffer>(obj));
    oiio_bufinfo binfo(obj.request());
    bool ok = true;
    vals.reserve(binfo.size);
    for (size_t i = 0; i < binfo.size; ++i) {
        if (std::is_same<T, float>::value
            && binfo.format.basetype == TypeDesc::FLOAT) {
            vals.push_back(binfo.dataval<float>(i));
        } else if ((std::is_same<T, float>::value || std::is_same<T, int>::value)
                   && binfo.format.basetype == TypeDesc::INT) {
            vals.push_back(T(binfo.dataval<int>(i)));
        } else if (std::is_same<T, unsigned int>::value
                   && binfo.format.basetype == TypeDesc::UINT) {
            vals.push_back(T(binfo.dataval<unsigned int>(i)));
        } else if (std::is_same<T, unsigned char>::value
                   && binfo.format.basetype == TypeDesc::UINT8) {
            vals.push_back(T(binfo.dataval<unsigned char>(i)));
        } else if (std::is_same<T, unsigned short>::value
                   && binfo.format.basetype == TypeDesc::UINT16) {
            vals.push_back(T(binfo.dataval<unsigned short>(i)));
        } else {
            // FIXME? Other cases?
            vals.push_back(T(42));
            ok = false;
        }
    }
    return ok;
}


// Specialization for reading strings
template<>
inline bool
py_buffer_to_stdvector(std::vector<std::string>& /*vals*/,
                       const py::buffer& /*obj*/)
{
    return false;  // not supported
}


// Specialization for reading TypeDesc
template<>
inline bool
py_buffer_to_stdvector(std::vector<TypeDesc>& /*vals*/,
                       const py::buffer& /*obj*/)
{
    return false;  // not supported
}



// Suck up a tuple (or list, or single instance) of presumed T values into a
// vector<T>. Works for T any of int, float, string, and works if the Python
// container is any of tuple, list.
template<typename T>
inline bool
py_to_stdvector(std::vector<T>& vals, const py::object& obj)
{
    if (py::isinstance<py::tuple>(obj)) {  // if it's a Python tuple
        return py_indexable_pod_to_stdvector(vals, obj.cast<py::tuple>());
    }
    if (py::isinstance<py::list>(obj)) {  // if it's a Python list
        return py_indexable_pod_to_stdvector(vals, obj.cast<py::list>());
    }
    // Apparently a str can masquerade as a buffer object, so make sure to
    // exclude that from the buffer case.
    if (py::isinstance<py::buffer>(obj) && !py::isinstance<py::str>(obj)) {
        return py_buffer_to_stdvector(vals, obj.cast<py::buffer>());
    }

    // handle scalar case or bust
    return py_scalar_pod_to_stdvector(vals, obj);
}



template<typename T>
inline py::tuple
C_to_tuple(cspan<T> vals)
{
    size_t size = vals.size();
    py::tuple result(size);
    for (size_t i = 0; i < size; ++i)
        result[i] = typename PyTypeForCType<T>::type(vals[i]);
    return result;
}


template<typename T>
inline py::tuple
C_to_tuple(span<T> vals)
{
    size_t size = vals.size();
    py::tuple result(size);
    for (size_t i = 0; i < size; ++i)
        result[i] = typename PyTypeForCType<T>::type(vals[i]);
    return result;
}


template<typename T>
inline py::tuple
C_to_tuple(const T* vals, size_t size)
{
    py::tuple result(size);
    for (size_t i = 0; i < size; ++i)
        result[i] = typename PyTypeForCType<T>::type(vals[i]);
    return result;
}


template<>
inline py::tuple
C_to_tuple(cspan<unsigned char> vals)
{
    size_t size = vals.size();
    py::tuple result(size);
    for (size_t i = 0; i < size; ++i)
        result[i] = static_cast<unsigned char>(vals[i]);
    return result;
}


// Special case for TypeDesc
template<>
inline py::tuple
C_to_tuple<TypeDesc>(cspan<TypeDesc> vals)
{
    size_t size = vals.size();
    py::tuple result(size);
    for (size_t i = 0; i < size; ++i)
        result[i] = py::cast(vals[i]);
    return result;
}



// Convert an array of T values (described by type) into either a simple
// Python object (if it's an int, float, or string and a SCALAR) or a
// Python tuple.
template<typename T>
inline py::object
C_to_val_or_tuple(const T* vals, TypeDesc type, int nvalues = 1)
{
    OIIO_DASSERT(vals && nvalues);
    size_t n = type.numelements() * type.aggregate * nvalues;
    if (n == 1 && !type.arraylen)
        return typename PyTypeForCType<T>::type(vals[0]);
    else
        return C_to_tuple(cspan<T>(vals, n));
}



template<typename T, typename POBJ>
bool
attribute_typed(T& myobj, string_view name, TypeDesc type, const POBJ& dataobj)
{
    if (type.basetype == TypeDesc::INT) {
        std::vector<int> vals;
        bool ok = py_to_stdvector(vals, dataobj);
        ok &= (vals.size() == type.numelements() * type.aggregate);
        if (ok)
            myobj.attribute(name, type, &vals[0]);
        return ok;
    }
    if (type.basetype == TypeDesc::UINT) {
        std::vector<unsigned int> vals;
        bool ok = py_to_stdvector(vals, dataobj);
        ok &= (vals.size() == type.numelements() * type.aggregate);
        if (ok)
            myobj.attribute(name, type, &vals[0]);
        return ok;
    }
    if (type.basetype == TypeDesc::UINT8) {
        std::vector<unsigned char> vals;
        bool ok = py_to_stdvector(vals, dataobj);
        ok &= (vals.size() == type.numelements() * type.aggregate);
        if (ok)
            myobj.attribute(name, type, &vals[0]);
        return ok;
    }
    if (type.basetype == TypeDesc::FLOAT) {
        std::vector<float> vals;
        bool ok = py_to_stdvector(vals, dataobj);
        ok &= (vals.size() == type.numelements() * type.aggregate);
        if (ok)
            myobj.attribute(name, type, &vals[0]);
        return ok;
    }
    if (type.basetype == TypeDesc::STRING) {
        std::vector<std::string> vals;
        bool ok = py_to_stdvector(vals, dataobj);
        ok &= (vals.size() == type.numelements() * type.aggregate);
        if (ok) {
            std::vector<ustring> u;
            for (auto& val : vals)
                u.emplace_back(val);
            myobj.attribute(name, type, &u[0]);
        }
        return ok;
    }
    return false;
}



// `data` points to values of `type`. Make a python object that represents
// them.
py::object
make_pyobject(const void* data, TypeDesc type, int nvalues = 1,
              py::object defaultvalue = py::none());



template<typename T>
py::object
getattribute_typed(const T& obj, const std::string& name,
                   TypeDesc type = TypeUnknown)
{
    if (type == TypeUnknown)
        return py::none();
    char* data = OIIO_ALLOCA(char, type.size());
    bool ok    = obj.getattribute(name, type, data);
    if (!ok)
        return py::none();  // None
    return make_pyobject(data, type);
}



// TRANSFERS ownership of the data pointer!
// N.B. There is some evidence that this doesn't work properly with
// non-float arrays. Maybe a limitation of pybind11?
template<class T>
inline py::array_t<T>
make_numpy_array(T* data, int dims, size_t chans, size_t width, size_t height,
                 size_t depth = 1)
{
    const size_t size = chans * width * height * depth;
    T* mem            = data ? data : new T[size];

    // Create a Python object that will free the allocated memory when
    // destroyed:
    py::capsule free_when_done(mem, [](void* f) {
        delete[] (reinterpret_cast<T*>(f));
    });

    std::vector<size_t> shape, strides;
    if (dims == 4) {  // volumetric
        shape.assign({ depth, height, width, chans });
        strides.assign({ height * width * chans * sizeof(T),
                         width * chans * sizeof(T), chans * sizeof(T),
                         sizeof(T) });
    } else if (dims == 3 && depth == 1) {  // 2D+channels
        shape.assign({ height, width, chans });
        strides.assign(
            { width * chans * sizeof(T), chans * sizeof(T), sizeof(T) });
    } else if (dims == 2 && depth == 1
               && height == 1) {  // 1D (scanline) + channels
        shape.assign({ width, chans });
        strides.assign({ chans * sizeof(T), sizeof(T) });
    } else {  // punt -- make it a 1D array
        shape.assign({ size });
        strides.assign({ sizeof(T) });
    }
    return py::array_t<T>(shape, strides, mem, free_when_done);
}



inline py::object
make_numpy_array(TypeDesc format, void* data, int dims, size_t chans,
                 size_t width, size_t height, size_t depth = 1)
{
    if (format == TypeDesc::FLOAT)
        return make_numpy_array((float*)data, dims, chans, width, height,
                                depth);
    if (format == TypeDesc::UINT8)
        return make_numpy_array((unsigned char*)data, dims, chans, width,
                                height, depth);
    if (format == TypeDesc::UINT16)
        return make_numpy_array((unsigned short*)data, dims, chans, width,
                                height, depth);
    if (format == TypeDesc::INT8)
        return make_numpy_array((char*)data, dims, chans, width, height, depth);
    if (format == TypeDesc::INT16)
        return make_numpy_array((short*)data, dims, chans, width, height,
                                depth);
    if (format == TypeDesc::DOUBLE)
        return make_numpy_array((double*)data, dims, chans, width, height,
                                depth);
    if (format == TypeDesc::HALF)
        return make_numpy_array((half*)data, dims, chans, width, height, depth);
    if (format == TypeDesc::UINT)
        return make_numpy_array((unsigned int*)data, dims, chans, width, height,
                                depth);
    if (format == TypeDesc::INT)
        return make_numpy_array((int*)data, dims, chans, width, height, depth);
    delete[] (char*)data;
    return py::none();
}



template<typename C>
inline void
delegate_setitem(C& self, const std::string& key, py::object obj)
{
    if (py::isinstance<py::float_>(obj))
        self[key] = float(obj.template cast<py::float_>());
    else if (py::isinstance<py::int_>(obj))
        self[key] = int(obj.template cast<py::int_>());
    else if (py::isinstance<py::str>(obj))
        self[key] = std::string(obj.template cast<py::str>());
    else if (py::isinstance<py::bytes>(obj))
        self[key] = std::string(obj.template cast<py::bytes>());
    else
        throw std::invalid_argument("Bad type for __setitem__");
}


}  // namespace PyOpenImageIO

#endif  // PYOPENIMAGEIO_PY_OIIO_H
