#ifndef DEPTHGUIDE_BINARY_UTILITY_HPP
#define DEPTHGUIDE_BINARY_UTILITY_HPP

#include "core_v2/memory_manager.h"

#include <nlohmann/json.hpp>

#include <boost/core/noncopyable.hpp>
#include <boost/endian.hpp>

#include <bit>
#include <cassert>
#include <cstdint>
#include <memory>
#include <utility>

struct data_mem_type : private boost::noncopyable {

    host_memory_info mem;
    uint8_t *ptr = nullptr;
    size_t size = 0;

    explicit data_mem_type(size_t _size) {
        size = _size;
        mem = HOST_ALLOC(size);
        ptr = static_cast<uint8_t *>(mem.ptr);
    }

    uint8_t *start_ptr() const {
        return ptr;
    }

    uint8_t *end_ptr() const {
        return ptr + size;
    }
};

struct data_type {
    uint8_t *ptr = nullptr;
    size_t size = 0;

    data_type() = default;

    // pre_size: reserved size before ptr.
    explicit data_type(size_t _size, size_t pre_size = 0)
            : data_type(
            std::make_shared<data_mem_type>(_size + pre_size), pre_size, _size) {
    }

    data_type(size_t _size, void *data, size_t pre_size = 0)
            : data_type(_size, pre_size) {
        replace(0, _size, (uint8_t *) data);
    }

    bool empty() const {
        return ptr == nullptr;
    }

    auto clone(size_t pre_size = 0) const {
        auto ret = data_type(size, pre_size);
        ret.replace(0, *this);
        return ret;
    }

    /* if _size == -1:
     *   ret_size = -offset       (offset < 0)
     *   ret_size keeps end_ptr() (offset >= 0) */
    auto sub_data(ptrdiff_t offset, size_t _size = -1) const {
        // determine ret_size
        if (_size == -1) {
            if (offset < 0) {
                _size = -offset;
            } else {
                _size = size - offset;
            }
        }

        auto ret_ptr = ptr + offset;
        assert(ret_ptr >= mem->start_ptr());
        assert(ret_ptr + _size <= mem->end_ptr());
        return data_type(mem, ret_ptr - mem->ptr, _size);
    }

    void replace(size_t offset, size_t _size, uint8_t *data) {
        assert(offset + _size <= size);
        std::copy_n(data, _size, start_ptr() + offset);
    }

    void replace(size_t offset, const data_type &data) {
        replace(offset, data.size, data.start_ptr());
    }

    auto extend(size_t _size) const {
        return sub_data(0, size + _size);
    }

    void extend_self(size_t _size) {
        reserve(size + _size, true);
    }

    void reserve(size_t _size, bool keep_data = false) {
        if (_size <= size) [[likely]] {
            size = _size;
            return;
        }
        assert(_size > size);
        if (mem != nullptr &&
            start_ptr() + _size <= mem->end_ptr()) {
            size = _size;
            return;
        }

        auto next_size = std::bit_ceil(_size);
        auto next = data_type(next_size);
        next.shrink(_size);

        if (keep_data) {
            assert(next.size >= size);
            next.replace(0, size, start_ptr());
        }
        *this = next;
    }

    void shrink(size_t _size) {
        assert(_size <= size);
        size = _size;
    }

    template<typename T>
    T *at(size_t pos) {
        static_assert(std::is_trivial_v<T>);
        assert(size % sizeof(T) == 0);
        assert(pos < size / sizeof(T));
        return ((T *) start_ptr()) + pos;
    }

    uint8_t *start_ptr() const {
        return ptr;
    }

    uint8_t *end_ptr() const {
        return ptr + size;
    }

private:
    std::shared_ptr<data_mem_type> mem;

    data_type(const std::shared_ptr<data_mem_type> &_mem,
              size_t offset, size_t _size) {
        mem = _mem;
        ptr = mem->ptr + offset;
        size = _size;
        assert(offset + size <= mem->size);
    }
};

template<boost::endian::order net_order, typename T>
std::enable_if_t<std::is_arithmetic_v<T>>
inline swap_net_loc_endian(T &val) {
    if constexpr (boost::endian::order::native == net_order) {
        return;
    } else {
        boost::endian::endian_reverse_inplace(val);
    }
}

class versatile_io : public boost::noncopyable {
public:

    versatile_io() // enable dynamic memory if no data provided
            : versatile_io({}, true) {
    }

    explicit versatile_io(data_type _data,
                          bool _extendable = false)
            : data(std::move(_data)) {
        cur_ptr = start_ptr();
        extendable = _extendable;
    }

    // 从当前位置开始调整 offset
    void manual_offset(ptrdiff_t offset) {
        cur_ptr += offset;
        assert(cur_ptr >= start_ptr());
        assert(cur_ptr <= end_ptr());
    }

    auto current_offset() const {
        return cur_ptr - start_ptr();
    }

    auto remaining_bytes() const {
        return end_ptr() - cur_ptr;
    }

    bool empty() const {
        return cur_ptr == end_ptr();
    }

    data_type current_data() const {
        return data;
    }

protected:
    data_type data;
    uint8_t *cur_ptr = nullptr;
    bool extendable = false;

    uint8_t *start_ptr() const {
        return data.start_ptr();
    }

    uint8_t *end_ptr() const {
        return data.end_ptr();
    }

    void ensure_remaining(size_t size, bool may_extend = true) {
        if (remaining_bytes() >= size) [[likely]] return;
        if (extendable && may_extend) {
            auto offset = current_offset(); // underlying data may be changed after reserve
            data.reserve(offset + size, true);
            cur_ptr = start_ptr() + offset;
        } else {
            RET_ERROR;
        }
    }
};

// 分多次读取数据
template<boost::endian::order net_order>
class versatile_reader : public versatile_io {
public:

    using versatile_io::versatile_io;

    template<typename T>
    std::enable_if_t<std::is_arithmetic_v<T>, T>
    read_value() {
        T tmp_val;
        ensure_remaining(sizeof(T), false);
        std::copy_n(cur_ptr, sizeof(T), (uint8_t *) &tmp_val);
        swap_net_loc_endian<net_order>(tmp_val);
        cur_ptr += sizeof(T);
        return tmp_val;
    }

    template<typename T>
    std::enable_if_t<std::is_arithmetic_v<T>>
    read_value(T &val) {
        val = read_value<T>();
    }

    template<typename T, size_t Length>
    void read_value(std::array<T, Length> &arr) {
        for (auto &val: arr) {
            read_value(val);
        }
    }

    data_type read_data(size_t size) {
        ensure_remaining(size, false);
        auto offset = cur_ptr - start_ptr();
        auto ret = data.sub_data(offset, size);
        cur_ptr += size;
        return ret;
    }

    template<typename SizeType = uint32_t>
    data_type read_data_with_length() {
        auto size = read_value<SizeType>();
        return read_data(size);
    }

    void read_data(void *data, size_t size) {
        ensure_remaining(size, false);
        std::copy_n(cur_ptr, size, (uint8_t *) data);
        cur_ptr += size;
    }

    void read_data(const data_type &_data) {
        read_data(_data.start_ptr(), _data.size);
    }

    std::string read_std_string(size_t size) {
        ensure_remaining(size, false);
        auto ret = std::string((char *) cur_ptr, size);
        cur_ptr += size;
        return ret;
    }

    template<typename SizeType = uint32_t>
    nlohmann::json read_json_with_length() {
        auto j_size = read_value<SizeType>();
        auto j_str = read_std_string(j_size);
        return nlohmann::json::parse(j_str);
    }

    data_type read_remain() {
        auto offset = cur_ptr - start_ptr();
        auto size = end_ptr() - cur_ptr;
        auto ret = data.sub_data(offset, size);
        cur_ptr = end_ptr();
        return ret;
    }

    template<typename T>
    auto &operator>>(T &val) {
        read_value(val);
        return *this;
    }
};

// 分多次写入数据
template<boost::endian::order net_order>
class versatile_writer : public versatile_io {
public:

    using versatile_io::versatile_io;

    template<typename T>
    std::enable_if_t<std::is_arithmetic_v<T>>
    write_value(T val) {
        swap_net_loc_endian<net_order>(val);
        ensure_remaining(sizeof(T));
        std::copy_n((uint8_t *) &val, sizeof(T), cur_ptr);
        cur_ptr += sizeof(T);
    }

    template<typename T, size_t Length>
    void write_value(const std::array<T, Length> &arr) {
        for (auto val: arr) {
            write_value(val);
        }
    }

    void write_data(const void *data, size_t size) {
        ensure_remaining(size);
        std::copy_n((uint8_t *) data, size, cur_ptr);
        cur_ptr += size;
    }

    void write_data(const data_type &_data) {
        write_data(_data.start_ptr(), _data.size);
    }

    template<typename SizeType = uint32_t>
    void write_with_length(const data_type &_data) {
        auto size = (SizeType) _data.size;
        write_value(size);
        write_value(_data);
    }

    void write_value(const std::string &str) {
        write_data(str.data(), str.length());
    }

    void write_value(const data_type &_data) {
        write_data(_data);
    }

    template<typename SizeType = uint32_t>
    void write_with_length(const nlohmann::json &j) {
        auto j_str = j.dump();
        auto j_size = (SizeType) j_str.length();
        write_value(j_size);
        write_value(j_str);
    }

    template<typename T>
    auto &operator<<(const T &val) {
        write_value(val);
        return *this;
    }
};

static constexpr auto network_order = boost::endian::order::big;
using network_writer = versatile_writer<network_order>;
using network_reader = versatile_reader<network_order>;

template<typename T, boost::endian::order net_order = network_order>
static uint8_t *write_binary_number(uint8_t *ptr, T val) {
    static constexpr auto need_swap =
            (boost::endian::order::native != net_order);
    auto real_ptr = (T *) ptr;
    if constexpr (need_swap) {
        *real_ptr = boost::endian::endian_reverse(val);
    } else {
        *real_ptr = val;
    }
    return ptr + sizeof(T);
}

template<typename T, boost::endian::order net_order = network_order>
static uint8_t *read_binary_number(uint8_t *ptr, T *val) {
    static constexpr auto need_swap =
            (boost::endian::order::native != net_order);
    *val = *(T *) ptr;
    if constexpr (need_swap) {
        boost::endian::endian_reverse_inplace(*val);
    }
    return ptr + sizeof(T);
}

#endif //DEPTHGUIDE_BINARY_UTILITY_HPP