#include "image_utility_v2_impl.h"
#include "core/image_utility_v2.h"


namespace image_utility_impl {

    std::type_index cv_type_id(int type) {
        switch (type) {
            // @formatter:off
            case CV_8UC1: { return typeid(uchar1); }
            case CV_8UC2: { return typeid(uchar2); }
            case CV_8UC3: { return typeid(uchar3); }
            case CV_8UC4: { return typeid(uchar4); }
            case CV_16UC1: { return typeid(ushort1); }
            case CV_32FC1: { return typeid(float1); }
            case CV_32FC2: { return typeid(float2); }
            // @formatter:on
            default: {
                RET_ERROR;
            }
        }
    }

}

void *image_memory::start_ptr(int component) const {
    switch (img->pixel_format()) {
        case PIX_NORMAL: {
            assert(component == 0);
            return ptr.get();
        }
        case PIX_NV12: {
            if (component == 0) { return ptr.get(); }
            if (component == 1) { return (uint8_t *) ptr.get() + pitch * img->height(); }
            RET_ERROR_E;
        }
        default: {
            RET_ERROR_E;
        }
    }
}

void *image_memory::at(int row, int col, int component) {
    auto sp = (uint8_t *) start_ptr(component);
    return sp + row * pitch + col * img->elem_size();
}

void image_memory::modified(smart_cuda_stream *stream) {
    auto &pimpl = img->pimpl;
    if (ptr == pimpl->store_host.ptr) {
        pimpl->q_this->host_modified(stream);
    } else {
        assert(ptr == pimpl->store_cuda.ptr);
        pimpl->q_this->cuda_modified(stream);
    }
}

void *generic_image::impl::storage_info::row_start(size_t row) {
    return (uint8_t *) ptr.get() + row * pitch;
}

void generic_image::impl::storage_info::reset() {
    ptr = nullptr;
    pitch = 0;
}

generic_image::impl::impl(generic_image::create_config conf) {
    type = conf.type;
    pix_fmt = conf.pixel;
    size = conf.size;

    // adjust display size to storage size
    switch (pix_fmt) {
        case PIX_NORMAL: {
            break;
        }
        case PIX_NV12: {
            size.height = nv12_storage_height(size.height);
            break;
        }
        default: {
            assert(false);
        }
    }
}

cv::Size generic_image::impl::display_size() const {
    cv::Size ret = size;
    switch (pix_fmt) {
        case PIX_NORMAL: {
            break;
        }
        case PIX_NV12: {
            ret.height = nv12_display_height(ret.height);
            break;
        }
        default: {
            RET_ERROR_E;
        }
    }
    return ret;
}

size_t generic_image::impl::elem_bytes() const {
    return CV_ELEM_SIZE(type);
}

size_t generic_image::impl::width_in_bytes() const {
    return size.width * elem_bytes();
}

size_t generic_image::impl::size_in_bytes() const {
    return size.height * width_in_bytes();
}

void generic_image::impl::create_host(smart_cuda_stream *stream) {
    if (store_host.ptr != nullptr) {
        SYNC_CREATE(store_host.ptr, stream);
        return;
    }
    store_host.ptr = ALLOC_PITCH_SHARED(
            uint8_t, width_in_bytes(), size.height, MEM_HOST, &store_host.pitch);
    if (store_cuda.ptr != nullptr) {
        SYNC_CREATE(store_cuda.ptr, stream);
        CUDA_API_CHECK(cudaMemcpy2DAsync(store_host.ptr.get(), store_host.pitch, // dst
                                         store_cuda.ptr.get(), store_cuda.pitch, // src
                                         width_in_bytes(), size.height,
                                         cudaMemcpyDeviceToHost, cuda_stream(stream)));
        REC_CREATE(store_host.ptr, stream);
    }
}

void generic_image::impl::create_cuda(smart_cuda_stream *stream) {
    if (store_cuda.ptr != nullptr) {
        SYNC_CREATE(store_cuda.ptr, stream);
        return;
    }
    store_cuda.ptr = ALLOC_PITCH_SHARED(
            uint8_t, width_in_bytes(), size.height, MEM_CUDA, &store_cuda.pitch);
    if (store_host.ptr != nullptr) {
        SYNC_CREATE(store_host.ptr, stream);
        CUDA_API_CHECK(cudaMemcpy2DAsync(store_cuda.ptr.get(), store_cuda.pitch, // dst
                                         store_host.ptr.get(), store_host.pitch, // src
                                         width_in_bytes(), size.height,
                                         cudaMemcpyHostToDevice, stream->cuda));
        REC_CREATE(store_cuda.ptr, stream);
    }
}

image_mem_info generic_image::impl::get_memory_v1(smart_cuda_stream *stream) const {
    auto ret = image_mem_info{
            .width = width_in_bytes(),
            .height = (size_t) size.height
    };
    if (store_cuda.ptr != nullptr) {
        ret.loc = MEM_CUDA;
        SYNC_CREATE(store_cuda.ptr, stream);
        ret.ptr = store_cuda.ptr;
        ret.pitch = store_cuda.pitch;
    } else {
        assert(store_host.ptr != nullptr);
        ret.loc = MEM_HOST;
        SYNC_CREATE(store_host.ptr, stream);
        ret.ptr = store_host.ptr;
        ret.pitch = store_host.pitch;
    }
    return ret;
}

image_memory generic_image::impl::get_memory(memory_location loc,
                                             smart_cuda_stream *stream) {
    auto ret = image_memory();
    ret.img = q_this->shared_from_this();
    ret.width = width_in_bytes();
    ret.height = size.height;

    switch (loc) {
        case MEM_HOST: {
            create_host(stream);
            ret.ptr = store_host.ptr;
            ret.pitch = store_host.pitch;
            break;
        }
        case MEM_CUDA: {
            create_cuda(stream);
            ret.ptr = store_cuda.ptr;
            ret.pitch = store_cuda.pitch;
            break;
        }
        default: {
            RET_ERROR_E;
        }
    }
    return ret;
}

cv::Mat generic_image::impl::get_cv_mat(smart_cuda_stream *stream) {
    create_host(stream);
    return cv::Mat(size, type, store_host.ptr.get(), store_host.pitch);
}

cv::cuda::GpuMat generic_image::impl::get_cv_gpumat(smart_cuda_stream *stream) {
    create_cuda(stream);
    return cv::cuda::GpuMat(size, type, store_cuda.ptr.get(), store_cuda.pitch);
}

template<typename T>
image_type_v2<T> generic_image::impl::get_image_type_v2(smart_cuda_stream *stream) {
    create_cuda(stream);
    assert(cv_type_id(type) == typeid(T));
    assert(size.width <= std::numeric_limits<ushort>::max());
    assert(size.height <= std::numeric_limits<ushort>::max());
    return image_type_v2<T>(
            (T *) store_cuda.ptr.get(), size.width, size.height, store_cuda.pitch);
}

template<typename T>
std::shared_ptr<smart_image<T>> generic_image::impl::get_image_v1() const {
    using ret_type = std::shared_ptr<smart_image<T>>;
    auto ret = ret_type();
    using info_type = image_info_type<T>;
    if (store_host.ptr != nullptr) {
        auto host_info = info_type{
                .ptr = std::reinterpret_pointer_cast<T>(store_host.ptr),
                .loc = MEM_HOST, .size = size, .pitch = store_host.pitch,
        };
        ret = std::make_shared<smart_image<T>>(host_info);
    }

    if (store_cuda.ptr != nullptr) {
        auto cuda_info = info_type{
                .ptr = std::reinterpret_pointer_cast<T>(store_cuda.ptr),
                .loc= MEM_CUDA, .size = size, .pitch = store_cuda.pitch,
        };
        if (ret == nullptr) {
            ret = std::make_shared<smart_image<T>>(cuda_info);
        } else {
            ret->cuda_info = cuda_info;
        }
    }

    assert(ret != nullptr);
    return ret;
}

void generic_image::impl::pixel_at(int row, int col, int component,
                                   void *dst, size_t _size, smart_cuda_stream *stream) {
    if (store_host.ptr != nullptr) {
        auto ptr = get_memory(MEM_HOST, stream).at(row, col, component);
        memcpy(dst, ptr, _size);
    } else {
        assert(store_cuda.ptr != nullptr);
        auto ptr = get_memory(MEM_CUDA, stream).at(row, col, component);
        CUDA_API_CHECK(cudaMemcpyAsync(dst, ptr, _size,
                                       cudaMemcpyDeviceToHost, stream->cuda));
        CUDA_API_CHECK(cudaStreamSynchronize(stream->cuda));
    }
}

template<typename T>
void generic_image::impl::create_from_v1(const std::shared_ptr<smart_image<T>> &img) {
    if (img->host_info.ptr != nullptr) {
        store_host.ptr = img->host_info.ptr;
        store_host.pitch = img->host_info.pitch;
    }
    if (img->cuda_info.ptr != nullptr) {
        store_cuda.ptr = img->cuda_info.ptr;
        store_cuda.pitch = img->cuda_info.pitch;
    }
}

void generic_image::impl::sub_image_inplace(int row, int col, int width, int height) {
    if (width == -1) { width = size.width - col; }
    if (height == -1) { height = size.height - row; }
    assert(width + col <= size.width);
    assert(height + row <= size.height);

    if (pix_fmt == PIX_NV12) {
        assert(row == 0 && height == size.height);
    } else {
        // sub-image of other formats are not implemented
        assert(pix_fmt == PIX_NORMAL);
    }

    size = cv::Size(width, height);

    if (store_host.ptr != nullptr) {
        store_host.ptr = std::shared_ptr<void>(
                (uint8_t *) store_host.row_start(row) + col * elem_bytes(),
                [p = store_host.ptr](void *) {});
    }
    if (store_cuda.ptr != nullptr) {
        store_cuda.ptr = std::shared_ptr<void>(
                (uint8_t *) store_cuda.row_start(row) + col * elem_bytes(),
                [p = store_cuda.ptr](void *) {});
    }
}

void generic_image::impl::type_cast_inplace(int _type) {
    // bit-cast of other formats are not implemented
    assert(pix_fmt == PIX_NORMAL);
    auto _width = width_in_bytes() / CV_ELEM_SIZE(_type);
    assert(_width * CV_ELEM_SIZE(_type) == width_in_bytes());
    size.width = _width;
    type = _type;
}

void generic_image::impl::host_modified(smart_cuda_stream *stream) {
    assert(store_host.ptr != nullptr);
    store_cuda.reset();
    REC_CREATE(store_host.ptr, stream);
}

void generic_image::impl::cuda_modified(smart_cuda_stream *stream) {
    assert(store_cuda.ptr != nullptr);
    store_host.reset();
    REC_CREATE(store_cuda.ptr, stream);
}

bool generic_image::basic_info_type::operator==(const basic_info_type &o) const {
    if (size != o.size) return false;
    if (cv_type != o.cv_type) return false;
    if (pixel_format != o.pixel_format) return false;
    return true;
}

generic_image::generic_image(std::unique_ptr<impl> _pimpl)
        : meta_base(_pimpl.get()) {
    pimpl = std::move(_pimpl);
    assert(pimpl != nullptr);
    pimpl->q_this = this;
}

generic_image::pointer generic_image::create(create_config conf) {
    auto pimpl = std::make_unique<impl>(conf);
    auto ret = std::make_shared<generic_image>(std::move(pimpl));
    return ret;
}

generic_image::pointer generic_image::create(cv::Size size, int type, pixel_format_enum pixel) {
    auto conf = create_config{
            .size = size, .type = type, .pixel = pixel,
    };
    return create(conf);
}

template<typename T>
generic_image::pointer generic_image::create(const std::shared_ptr<smart_image<T>> &img) {
    auto ret = create(img->size(), get_cv_type<T>());
    ret->pimpl->create_from_v1(img);
    return ret;
}

// @formatter:off
template generic_image::pointer generic_image::create(const image_u8c1 &);
template generic_image::pointer generic_image::create(const image_u8c2 &);
template generic_image::pointer generic_image::create(const image_u8c3 &);
template generic_image::pointer generic_image::create(const image_u8c4 &);
template generic_image::pointer generic_image::create(const image_u16c1 &);
template generic_image::pointer generic_image::create(const image_f32c1 &);
// @formatter:on

cv::Size generic_image::size() const {
    return pimpl->display_size();
}

size_t generic_image::size_in_bytes() const {
    return pimpl->size_in_bytes();
}

size_t generic_image::width_in_bytes() const {
    return pimpl->width_in_bytes();
}

int generic_image::cv_type() const {
    return pimpl->type;
}

pixel_format_enum generic_image::pixel_format() const {
    return pimpl->pix_fmt;
}

generic_image::basic_info_type generic_image::basic_info() const {
    auto ret = basic_info_type();
    ret.size = size();
    ret.cv_type = cv_type();
    ret.pixel_format = pixel_format();
    return ret;
}

image_mem_info generic_image::memory_v1(smart_cuda_stream *stream) const {
    return pimpl->get_memory_v1(stream);
}

image_memory generic_image::memory(memory_location loc,
                                   smart_cuda_stream *stream) {
    return pimpl->get_memory(loc, stream);
}

cv::Mat generic_image::cv_mat(smart_cuda_stream *stream) {
    return pimpl->get_cv_mat(stream);
}

cv::cuda::GpuMat generic_image::cv_gpumat(smart_cuda_stream *stream) {
    return pimpl->get_cv_gpumat(stream);
}

template<typename T>
image_type_v2<T> generic_image::cuda(smart_cuda_stream *stream) {
    return pimpl->get_image_type_v2<T>(stream);
}

// @formatter:off
template image_type_v2<uchar1> generic_image::cuda(smart_cuda_stream *stream);
template image_type_v2<uchar2> generic_image::cuda(smart_cuda_stream *stream);
template image_type_v2<uchar3> generic_image::cuda(smart_cuda_stream *stream);
template image_type_v2<uchar4> generic_image::cuda(smart_cuda_stream *stream);
template image_type_v2<ushort1> generic_image::cuda(smart_cuda_stream *stream);
template image_type_v2<float1> generic_image::cuda(smart_cuda_stream *stream);
template image_type_v2<float2> generic_image::cuda(smart_cuda_stream *stream);
// @formatter:on

template<typename T>
std::shared_ptr<smart_image<T>> generic_image::v1() const {
    return pimpl->get_image_v1<T>();
}

// @formatter:off
template std::shared_ptr<smart_image<uchar1>> generic_image::v1() const;
template std::shared_ptr<smart_image<uchar2>> generic_image::v1() const;
template std::shared_ptr<smart_image<uchar3>> generic_image::v1() const;
template std::shared_ptr<smart_image<uchar4>> generic_image::v1() const;
template std::shared_ptr<smart_image<ushort1>> generic_image::v1() const;
template std::shared_ptr<smart_image<float1>> generic_image::v1() const;
// @formatter:on

generic_image::pointer generic_image::shallow_clone() const {
    auto pimpl_c = std::make_unique<impl>(*pimpl);
    auto ret = std::make_shared<generic_image>(std::move(pimpl_c));
    return ret;
}

void generic_image::pixel_at_impl(int row, int col, int component,
                                  void *dst, size_t size, smart_cuda_stream *stream) {
    return pimpl->pixel_at(row, col, component, dst, size, stream);
}

generic_image::pointer generic_image::sub_image(int row, int col, int width, int height) const {
    auto ret = shallow_clone();
    ret->pimpl->sub_image_inplace(row, col, width, height);
    return ret;
}

generic_image::pointer generic_image::bit_cast(int type) {
    auto ret = shallow_clone();
    ret->pimpl->type_cast_inplace(type);
    return ret;
}

void generic_image::host_modified(smart_cuda_stream *stream) {
    pimpl->host_modified(stream);
}

void generic_image::cuda_modified(smart_cuda_stream *stream) {
    pimpl->cuda_modified(stream);
}

image_ptr to_image(obj_name_type name) {
    if (name == invalid_obj_name) return nullptr;
    auto img_type = OBJ_TYPE(name);
    if (OBJ_TYPE(name) == typeid(image_ptr)) {
        return OBJ_QUERY(image_ptr, name);
    }

    // convert from v1
    auto ret = image_ptr();
    auto impl_func = [&](auto V) {
        using T = std::remove_cvref_t<decltype(V)>;
        if (img_type == typeid(T)) {
            auto img = OBJ_QUERY(T, name);
            if (img == nullptr) return;
            ret = create_image(img);
        }
    };
    FORALL_IMG_TYPE;

    OBJ_MERGE_META(name, ret.get());
    return ret;
}