/*
    SPDX-FileCopyrightText: 2020 Volker Krause <vkrause@kde.org>

    SPDX-License-Identifier: LGPL-2.0-or-later
*/

#ifndef KANDROIDEXTRAS_JNIARRAY_H
#define KANDROIDEXTRAS_JNIARRAY_H

#include "jniargument.h"
#include "jnireturnvalue.h"
#include "jnitypetraits.h"

#include <QJniEnvironment>

namespace KAndroidExtras {

///@cond internal
namespace Internal {

/** Primitive type array type traits. */
template <typename T> struct array_trait {
    typedef jobjectArray type;
};

#define MAKE_ARRAY_TRAIT(base_type, type_name) \
template <> struct array_trait<base_type> { \
    typedef base_type ## Array type; \
    static inline type newArray(QJniEnvironment &env, jsize size) { return env->New ## type_name ## Array(size); } \
    static inline base_type* getArrayElements(QJniEnvironment &env, type array, jboolean *isCopy) { return env->Get ## type_name ## ArrayElements(array, isCopy); } \
    static inline void releaseArrayElements(QJniEnvironment &env, type array, base_type *data, jint mode) { return env->Release ## type_name ## ArrayElements(array, data, mode); } \
    static inline void setArrayRegion(QJniEnvironment &env, type array, jsize start, jsize length, const base_type *data) { env->Set ## type_name ## ArrayRegion(array, start, length, data); } \
};

MAKE_ARRAY_TRAIT(jboolean, Boolean)
MAKE_ARRAY_TRAIT(jbyte, Byte)
MAKE_ARRAY_TRAIT(jchar, Char)
MAKE_ARRAY_TRAIT(jshort, Short)
MAKE_ARRAY_TRAIT(jint, Int)
MAKE_ARRAY_TRAIT(jlong, Long)
MAKE_ARRAY_TRAIT(jfloat, Float)
MAKE_ARRAY_TRAIT(jdouble, Double)

#undef MAKE_ARRAY_TRAIT

/** Meta function for retrieving a JNI array .*/
template <typename Container, typename Value, bool is_primitive> struct FromArray {};

template <typename Container>
struct FromArray<Container, QJniObject, false>
{
    inline auto operator()(const QJniObject &array) const
    {
        if (!array.isValid()) {
            return Container{};
        }
        const auto a = static_cast<jobjectArray>(array.object());
        QJniEnvironment env;
        const auto size = env->GetArrayLength(a);
        Container r;
        r.reserve(size);
        for (auto i = 0; i < size; ++i) {
            r.push_back(QJniObject::fromLocalRef(env->GetObjectArrayElement(a, i)));
        }
        return r;
    }
};

template <typename Container, typename Value>
struct FromArray<Container, Value, false>
{
    inline auto operator()(const QJniObject &array) const
    {
        if (!array.isValid()) {
            return Container{};
        }
        const auto a = static_cast<jobjectArray>(array.object());
        QJniEnvironment env;
        const auto size = env->GetArrayLength(a);
        Container r;
        r.reserve(size);
        for (auto i = 0; i < size; ++i) {
            r.push_back(Jni::reverse_converter<Value>::type::convert(QJniObject::fromLocalRef(env->GetObjectArrayElement(a, i))));
        }
        return r;
    }
};

// specializations for primitive types
template <typename Container, typename Value>
struct FromArray<Container, Value, true>
{
    typedef array_trait<Value> _t;
    inline auto operator()(const QJniObject &array) const
    {
        if (!array.isValid()) {
            return Container{};
        }

        const auto a = static_cast<typename _t::type>(array.object());
        QJniEnvironment env;
        const auto size = env->GetArrayLength(a);
        Container r;
        r.reserve(size);

        auto data = _t::getArrayElements(env, a, nullptr);
        std::copy(data, data + size, std::back_inserter(r));
        _t::releaseArrayElements(env, a, data, JNI_ABORT);

        return r;
    }
};

// array wrapper, common base for primitive and non-primitive types
template <typename T>
class ArrayImplBase {
public:
    typedef T value_type;
    typedef jsize size_type;
    typedef jsize difference_type;

    ArrayImplBase() = default;
    inline ArrayImplBase(const QJniObject &array) : m_array(array) {}
    ArrayImplBase(const ArrayImplBase &) = default;
    ArrayImplBase(ArrayImplBase &&) = default;

    inline size_type size() const
    {
        if (!m_array.isValid()) {
            return 0;
        }
        const auto a = static_cast<typename _t::type>(m_array.object());
        QJniEnvironment env;
        return env->GetArrayLength(a);
    }

    inline operator QJniObject() const {
        return m_array;
    }
    inline QJniObject jniHandle() const { return m_array; }

protected:
    typedef array_trait<T> _t;

    typename _t::type handle() const { return static_cast<typename _t::type>(m_array.object()); }

    QJniObject m_array;
};

template <typename T, bool is_primitive>
class ArrayImpl {};

// array wrapper for primitive types
template <typename T>
class ArrayImpl<T, true> : public ArrayImplBase<T>
{
    static_assert(!Internal::is_invalid_primitive_type<T>::value, "Using an incompatible primitive type!");
public:
    inline ArrayImpl(const QJniObject &array)
        : ArrayImplBase<T>(array)
    {
        // ### do this on demand?
        getArrayElements();
    }

    /** Create a new array with @p size elements. */
    inline explicit ArrayImpl(jsize size)
    {
        QJniEnvironment env;
        ArrayImplBase<T>::m_array = QJniObject::fromLocalRef(ArrayImplBase<T>::_t::newArray(env, size));
        getArrayElements();
    }

    ArrayImpl() = default;
    ArrayImpl(const ArrayImpl&) = delete; // ### ref count m_data and allow copying?
    ArrayImpl(ArrayImpl&&) = default;
    ~ArrayImpl()
    {
        QJniEnvironment env;
        ArrayImplBase<T>::_t::releaseArrayElements(env, this->handle(), m_data, JNI_ABORT);
    }

    T operator[](jsize index) const
    {
        return m_data[index];
    }

    T* begin() const { return m_data; }
    T* end() const { return m_data + ArrayImplBase<T>::size(); }

private:
    inline void getArrayElements()
    {
        if (!ArrayImplBase<T>::m_array.isValid()) {
            return;
        }
        QJniEnvironment env;
        m_data = ArrayImplBase<T>::_t::getArrayElements(env, this->handle(), nullptr);
    }

    T *m_data = nullptr;
};

// array wrapper for non-primitive types
template <typename T>
class ArrayImpl<T, false> : public ArrayImplBase<T>
{
public:
    using ArrayImplBase<T>::ArrayImplBase;

    /** Create a new array with @p size elements initialized with @p value. */
    explicit inline ArrayImpl(jsize size, typename Internal::argument<T>::type value)
    {
        QJniEnvironment env;
        auto clazz = env.findClass(Jni::typeName<T>());
        ArrayImplBase<T>::m_array = QJniObject::fromLocalRef(env->NewObjectArray(size, clazz, Internal::argument<T>::toCallArgument(value)));
    }

    /** Create a new array with @p size null elements. */
    explicit inline ArrayImpl(jsize size, std::nullptr_t = nullptr)
    {
        QJniEnvironment env;
        auto clazz = env.findClass(Jni::typeName<T>());
        ArrayImplBase<T>::m_array = QJniObject::fromLocalRef(env->NewObjectArray(size, clazz, nullptr));
    }

    ArrayImpl() = default;
    ArrayImpl(const ArrayImpl&) = default;
    ArrayImpl(ArrayImpl&&) = default;

    auto operator[](jsize index) const
    {
        QJniEnvironment env;
        return Internal::return_wrapper<T>::toReturnValue(QJniObject::fromLocalRef(env->GetObjectArrayElement(this->handle(), index)));
    }

    class ref
    {
    public:
        inline operator auto()
        {
            QJniEnvironment env;
            return Internal::return_wrapper<T>::toReturnValue(QJniObject::fromLocalRef(env->GetObjectArrayElement(c.handle(), index)));
        }
        inline ref& operator=(typename Internal::argument<T>::type v)
        {
            QJniEnvironment env;
            env->SetObjectArrayElement(c.handle(), index, Internal::argument<T>::toCallArgument(v));
            return *this;
        }
    private:
        ArrayImpl<T, false> &c;
        jsize index;

        friend class ArrayImpl<T, false>;
        inline ref(jsize _i, ArrayImpl<T, false> &_c) : c(_c), index(_i) {}
    };
    ref operator[](jsize index)
    {
        return ref(index, *this);
    }

    class const_iterator
    {
        const ArrayImpl<T, false> &c;
        jsize i = 0;
    public:
        typedef jsize difference_type;
        typedef T value_type;
        typedef T& reference;
        typedef std::random_access_iterator_tag iterator_category;
        typedef T* pointer;

        const_iterator(const ArrayImpl<T, false> &_c, jsize _i) : c(_c), i(_i) {}

        difference_type operator-(const_iterator other) const { return i - other.i; }

        const_iterator& operator++() { ++i; return *this; }
        const_iterator operator++(int) { return const_iterator(c, i++); }

        bool operator==(const_iterator other) const { return i == other.i; }
        bool operator!=(const_iterator other) const { return i != other.i; }

        auto operator*() const {
            return c[i];
        }
    };

    const_iterator begin() const { return const_iterator(*this, 0); }
    const_iterator end() const { return const_iterator(*this, ArrayImplBase<T>::size()); }

};

}
///@endcond

namespace Jni {

/** Convert a JNI array to a C++ container.
    *  Container value types can be any of
    *  - QJniObject
    *  - a primitive JNI type
    *  - a type with a conversion defined with @c JNI_DECLARE_CONVERTER
    */
template <typename Container> constexpr __attribute__((__unused__)) Internal::FromArray<Container, typename Container::value_type, Jni::is_primitive_type<typename Container::value_type>::value> fromArray = {};

/** Container-like wrapper for JNI arrays. */
template <typename T>
class Array : public Internal::ArrayImpl<T, Jni::is_primitive_type<T>::value> {
public:
    using Internal::ArrayImpl<T, Jni::is_primitive_type<T>::value>::ArrayImpl;
    template <typename Container>
    inline operator Container() const {
        // ### should this be re-implemented in terms of Jni::Array API rather than direct JNI access?
        return Jni::fromArray<Container>(this->m_array);
    }
};
}

}

#endif
