/*
    Copyright Intel Corporation.
    
    This software and the related documents are Intel copyrighted materials, and
    your use of them is governed by the express license under which they were
    provided to you (License). Unless the License provides otherwise, you may
    not use, modify, copy, publish, distribute, disclose or transmit this
    software or the related documents without Intel's prior written permission.
    
    This software and the related documents are provided as is, with no express
    or implied warranties, other than those that are expressly stated in the
    License.
*/
#pragma once
#include <algorithm>
#include <cassert>
#include <numeric>
#include <iterator>
#include <sstream>
#include <stdexcept>

#include "oneapi/ccl/native_device_api/l0/base.hpp"
#include "oneapi/ccl/native_device_api/l0/subdevice.hpp"

namespace native {
inline std::ostream& operator<<(std::ostream& out, const ccl::device_index_type& index) {
    out << ccl::to_string(index);
    return out;
}
/**
 * Base RAII L0 handles wrappper
 * support serialize/deserialize concept
 */
#define TEMPLATE_DECL_ARG class handle_type, class resource_owner, class cl_context
#define TEMPLATE_DEF_ARG  handle_type, resource_owner, cl_context

template <TEMPLATE_DECL_ARG>
cl_base<TEMPLATE_DEF_ARG>::cl_base(handle_t h, owner_ptr_t parent, context_ptr_t ctx)
        : handle(h),
          owner(std::move(parent)),
          context(std::move(ctx)) {}

template <TEMPLATE_DECL_ARG>
cl_base<TEMPLATE_DEF_ARG>::cl_base(cl_base&& src) noexcept
        : handle(std::move(src.handle)),
          owner(std::move(src.owner)),
          context(std::move(src.context)) {}

template <TEMPLATE_DECL_ARG>
cl_base<TEMPLATE_DEF_ARG>& cl_base<TEMPLATE_DEF_ARG>::operator=(cl_base&& src) noexcept {
    handle = std::move(src.handle);
    owner = std::move(src.owner);
    context = std::move(src.context);
    return *this;
}

template <TEMPLATE_DECL_ARG>
CCL_API cl_base<TEMPLATE_DEF_ARG>::~cl_base() noexcept {
    auto lock = owner.lock();
    auto ctx = context.lock();
    if (lock && ctx) {
        auto ref = ctx->get();
        lock->on_delete(handle, ref);
    }
}

template <TEMPLATE_DECL_ARG>
std::shared_ptr<cl_base<TEMPLATE_DEF_ARG>> cl_base<TEMPLATE_DEF_ARG>::get_instance_ptr() {
    return this->shared_from_this();
}

template <TEMPLATE_DECL_ARG>
typename cl_base<TEMPLATE_DEF_ARG>::handle_t cl_base<TEMPLATE_DEF_ARG>::release() {
    handle_t ret;

    owner.reset();
    context.reset();

    std::swap(ret, handle);
    return ret;
}
template <TEMPLATE_DECL_ARG>
typename cl_base<TEMPLATE_DEF_ARG>::handle_t& cl_base<TEMPLATE_DEF_ARG>::get() noexcept {
    return const_cast<handle_t&>(static_cast<const self_t*>(this)->get());
}

template <TEMPLATE_DECL_ARG>
const typename cl_base<TEMPLATE_DEF_ARG>::handle_t& cl_base<TEMPLATE_DEF_ARG>::get()
    const noexcept {
    return handle;
}

template <TEMPLATE_DECL_ARG>
typename cl_base<TEMPLATE_DEF_ARG>::handle_t* cl_base<TEMPLATE_DEF_ARG>::get_ptr() noexcept {
    return &handle;
}

template <TEMPLATE_DECL_ARG>
const typename cl_base<TEMPLATE_DEF_ARG>::handle_t* cl_base<TEMPLATE_DEF_ARG>::get_ptr()
    const noexcept {
    return const_cast<handle_t*>(static_cast<const self_t*>(this)->get_ptr());
}

template <TEMPLATE_DECL_ARG>
const typename cl_base<TEMPLATE_DEF_ARG>::owner_ptr_t cl_base<TEMPLATE_DEF_ARG>::get_owner() const {
    return owner;
}

template <TEMPLATE_DECL_ARG>
const typename cl_base<TEMPLATE_DEF_ARG>::context_ptr_t cl_base<TEMPLATE_DEF_ARG>::get_ctx() const {
    return context;
}

template <TEMPLATE_DECL_ARG>
constexpr size_t cl_base<TEMPLATE_DEF_ARG>::get_size_for_serialize() {
    return resource_owner::get_size_for_serialize() + sizeof(handle_t);
}

template <TEMPLATE_DECL_ARG>
template <class... helpers>
size_t cl_base<TEMPLATE_DEF_ARG>::serialize(std::vector<uint8_t>& out,
                                            size_t from_pos,
                                            const helpers&... args) const {
    // check parent existence
    auto lock = owner.lock();
    if (!lock) {
        throw std::runtime_error("cannot serialize without owner");
    }

    constexpr size_t expected_bytes = sizeof(handle_t);

    // serialize from position
    size_t serialized_bytes =
        lock->serialize(out, from_pos, expected_bytes, args...); //resize vector inside

    uint8_t* data_start = out.data() + from_pos + serialized_bytes;
    *(reinterpret_cast<handle_t*>(data_start)) = handle;
    serialized_bytes += expected_bytes;
    return serialized_bytes;
}

template <TEMPLATE_DECL_ARG>
template <class type, class... helpers>
std::shared_ptr<type> cl_base<TEMPLATE_DEF_ARG>::deserialize(const uint8_t** data,
                                                             size_t& size,
                                                             std::shared_ptr<cl_context> ctx,
                                                             helpers&... args) {
    constexpr size_t expected_bytes = sizeof(handle);
    size_t initial_size = size;

    // recover parent handle at first
    auto owner = resource_owner::deserialize(data, size, args...).lock();

    size_t offset = initial_size - size;
    if (!owner or !offset) {
        throw std::runtime_error("cannot deserialize, owner if nullptr");
    }

    if (size < expected_bytes) {
        throw std::runtime_error("cannot deserialize, not enough data");
    }

    handle_t h = *(reinterpret_cast<const handle_t*>(*data));
    *data += expected_bytes;
    size -= expected_bytes;
    return std::shared_ptr<type>{ new type(h, owner, ctx) };
}

#undef TEMPLATE_DEF_ARG
#undef TEMPLATE_DECL_ARG

namespace detail {

template <ccl::device_index_enum index_id,
          class IndexedContainer,
          class value_type,
          class value_type_index_extractor>
indexed_storage<value_type> merge_indexed_values(const IndexedContainer& indexes,
                                                 std::vector<value_type>& values,
                                                 value_type_index_extractor functor) {
    indexed_storage<value_type> ret;
    try {
        // set indices to values at first
        indexed_storage<value_type> indexed_values;
        std::transform(values.begin(),
                       values.end(),
                       std::inserter(indexed_values, indexed_values.end()),
                       [functor](value_type& handle) ->
                       typename indexed_storage<value_type>::value_type {
                           auto index = functor(handle);
                           return std::make_pair(index, handle);
                       });

        // find requested device index in indexed values
        std::transform(
            indexes.begin(),
            indexes.end(),
            std::inserter(ret, ret.end()),
            [&indexed_values](const typename IndexedContainer::value_type& index) ->
            typename indexed_storage<value_type>::value_type {
                auto values_it = indexed_values.find(std::get<index_id>(index));
                if (values_it == indexed_values.end()) {
                    throw std::runtime_error(std::string(__PRETTY_FUNCTION__) +
                                             "Cannot find index: " + ccl::to_string(index) +
                                             ", from collected devices");
                }
                return std::make_pair(std::get<index_id>(index), values_it->second);
            });
    }
    catch (const std::exception& ex) {
        std::stringstream ss;
        ss << "Cannot merge indexed values\nRequested indices: ";
        for (const auto& index : indexes) {
            ss << index << ", ";
        }

        ss << "\nvalues to indexed: ";
        std::copy(values.begin(), values.end(), std::ostream_iterator<value_type>(ss, ", "));
        ss << "\nerror: " << ex.what();

        assert(false && "merge_indexed_values error");
        throw std::runtime_error(ss.str());
    }
    return ret;
}

template <ccl::device_index_enum index_id, class value_type, class value_type_index_extractor>
indexed_storage<value_type> collect_indexed_data(const ccl::device_indices_type& indexes,
                                                 std::vector<value_type>& collected_values,
                                                 value_type_index_extractor functor) {
    indexed_storage<value_type> ret;
    if (!indexes.empty()) //only requested indexes
    {
        ret = merge_indexed_values<index_id>(indexes, collected_values, functor);
    }
    else //all
    {
        std::transform(collected_values.begin(),
                       collected_values.end(),
                       std::inserter(ret, ret.end()),
                       [&functor](value_type handle) ->
                       typename indexed_storage<value_type>::value_type {
                           auto index = functor(handle);
                           return std::make_pair(index, handle);
                       });
    }
    return ret;
}

} // namespace detail
} // namespace native
