#include "orb_camera_impl.h"
#include "image_utility.hpp"
#include "object_manager.h"
#include "utility.hpp"

#include <boost/asio/post.hpp>

#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>

using boost::asio::post;

namespace orb_camera_impl {

    ob::Context ob_ctx;

    const char *get_fmt_name(OBFormat fmt) {
        switch (fmt) {
            // @formatter:off
            case OB_FORMAT_YUYV: return "yuyv";
            case OB_FORMAT_MJPG: return "mjpg";
            case OB_FORMAT_RGB: return "rgb";
            case OB_FORMAT_Y16: return "y16";
            case OB_FORMAT_RLE: return "rle";
            case OB_FORMAT_Y14: return "y14";
            default: return "unknown";
            // @formatter:on
        }
        RET_ERROR_P;
    }

    // convert uncompressed video frame to image_xxx
    template<typename T>
    auto video_uc_frame_to_image(const video_frame_type &frame) {
        auto info = image_info_type<T>();
        info.ptr = std::shared_ptr<T>( // extend frame's lifetime
                (T *) frame->data(), [pf = frame](void *) {});
        info.loc = MEM_HOST;
        info.size = cv::Size(frame->width(), frame->height());
        info.pitch = frame->width() * sizeof(T);
        assert(frame->dataSize() == info.pitch * info.size.height);
        return create_image(info);
    }

    auto mjpeg_frame_to_image(const video_frame_type &frame) { // TODO: accelerate with CUDA
        assert(frame->format() == OB_FORMAT_MJPG);
        auto img_data = cv::_InputArray((uchar *) frame->data(), frame->dataSize());
        auto img_bgr = cv::imdecode(img_data, cv::IMREAD_UNCHANGED);
        auto img_rgb_info = create_image_info<uchar3>(img_bgr.size(), MEM_HOST);
        cv::cvtColor(img_bgr, img_rgb_info.as_mat(), cv::COLOR_BGR2RGB);
        return create_image(img_rgb_info);
    }

    image_f32c1 depth_y16_to_mm(const image_u16c1 &y16, float scale) { // TODO: accelerate with CUDA
        auto y16_info = y16->as_host_info();
        auto f32_info = create_image_info<float1>(y16_info.size, MEM_HOST);
        y16->as_host().convertTo(f32_info.as_mat(), CV_32FC1, scale);
        return create_image(f32_info);
    }

}

std::shared_ptr<ob::Device> orb_camera::impl::get_device(const char *sn) {
    auto dev_list = ob_ctx.queryDeviceList();
    return dev_list->getDeviceBySN(sn);
}

orb_camera::impl *orb_camera::impl::create(orb_camera::create_config conf) {
    auto dev = get_device(conf.sn_str);
    auto ret = new impl();
    ret->pipe = std::make_shared<ob::Pipeline>(dev);
    ret->ctx = conf.ctx;
    ret->stream = conf.stream;
    return ret;
}

orb_camera::impl::~impl() {
    if (is_capturing) {
        stop();
    }
}

void orb_camera::impl::create_video_info_list(
        const pf_list_type &pf_list, v_info_list_type *v_info) {
    v_info->reserve(pf_list->count());
    for (uint32_t k = 0; k < pf_list->count(); ++k) {
        auto pf = pf_list->getProfile(k)->as<ob::VideoStreamProfile>();
        assert(pf != nullptr);
        auto info = video_info{};
        info.index = k;
        info.fmt_name = get_fmt_name(pf->format());
        info.fps = pf->fps();
        info.width = pf->width();
        info.height = pf->height();
        v_info->push_back(info);
    }
}

orb_camera::impl::v_info_list_type *
orb_camera::impl::query_video_info(OBSensorType s_type) {
    switch (s_type) {
        case OB_SENSOR_COLOR: {
            if (c_pf_list == nullptr) [[unlikely]] {
                c_pf_list = pipe->getStreamProfileList(s_type);
                c_info = std::make_unique<v_info_list_type>();
                create_video_info_list(c_pf_list, c_info.get());
            }
            return c_info.get();
        }
        case OB_SENSOR_DEPTH: {
            if (d_pf_list == nullptr) [[unlikely]] {
                d_pf_list = pipe->getStreamProfileList(s_type);
                d_info = std::make_unique<v_info_list_type>();
                create_video_info_list(d_pf_list, d_info.get());
            }
            return d_info.get();
        }
    }
    RET_ERROR_P;
}

bool orb_camera::impl::start(start_config conf) {
    assert(!is_capturing);
    auto ob_conf = std::make_shared<ob::Config>();

    if (conf.color.enable) {
        assert(c_pf_list != nullptr);
        ob_conf->enableStream(c_pf_list->getProfile(conf.color.config_index));
        c_name = conf.color.name;
    }
    if (conf.depth.enable) {
        assert(d_pf_list != nullptr);
        ob_conf->enableStream(d_pf_list->getProfile(conf.depth.config_index));
        d_name = conf.depth.name;
    }

    if (conf.color.enable && conf.depth.enable) {
        ob_conf->setAlignMode(ALIGN_D2C_HW_MODE);
    }

    pipe->start(ob_conf, [this](auto frames) {
        frames_callback(std::move(frames));
    });
    is_capturing = true;
    return true;
}

void orb_camera::impl::frames_callback(const frames_type &frames) {
    auto c_img = image_u8c3();
    if (auto c_frame = frames->colorFrame(); c_frame != nullptr) {
        assert(c_name != invalid_obj_name);
        auto fmt = c_frame->format();
        if (fmt == OB_FORMAT_RGB) {
            c_img = video_uc_frame_to_image<uchar3>(c_frame);
        } else {
            assert(fmt == OB_FORMAT_MJPG);
            c_img = mjpeg_frame_to_image(c_frame);
        }
    }

    auto d_img = image_f32c1();
    if (auto d_frame = frames->depthFrame(); d_frame != nullptr) {
        assert(d_name != invalid_obj_name);
        assert(d_frame->format() == OB_FORMAT_Y16);
        auto y16_img = video_uc_frame_to_image<ushort1>(d_frame);
        d_img = depth_y16_to_mm(y16_img, d_frame->getValueScale());
    }

    post(*ctx, [=, _c_name = c_name, _d_name = d_name] {
        if (c_img != nullptr) { OBJ_SAVE(_c_name, c_img); }
        if (d_img != nullptr) { OBJ_SAVE(_d_name, d_img); }
    });
}

void orb_camera::impl::stop() {
    assert(is_capturing);
    pipe->stop();
    is_capturing = false;
}

orb_camera *orb_camera::create(orb_camera::create_config conf) {
    auto pimpl = std::unique_ptr<impl>(impl::create(conf));
    if (pimpl == nullptr) return nullptr;
    auto ret = new orb_camera();
    ret->pimpl = std::move(pimpl);
    return ret;
}

orb_camera::~orb_camera() = default;

std::vector<orb_camera::device_info>
orb_camera::query_device_info() {
    auto dev_list = ob_ctx.queryDeviceList();
    std::vector<device_info> ret;
    ret.reserve(dev_list->deviceCount());
    for (auto k = 0; k < dev_list->deviceCount(); ++k) {
        ret.push_back({.sn_str = dev_list->serialNumber(k)});
    }
    return ret;
}

std::vector<orb_camera::video_info> *
orb_camera::query_video_info(video_type type) {
    switch (type) {
        case V_COLOR:
            return pimpl->query_video_info(OB_SENSOR_COLOR);
        case V_DEPTH:
            return pimpl->query_video_info(OB_SENSOR_DEPTH);
    }
    RET_ERROR_P;
}

bool orb_camera::start(orb_camera::start_config conf) {
    return pimpl->start(conf);
}

void orb_camera::stop() {
    pimpl->stop();
}

bool orb_camera::is_capturing() const {
    return pimpl->is_capturing;
}

orb_camera::video_info::operator std::string() {
    return fmt::format("{}x{} @ {}fps ({})", width, height, fps, fmt_name);
}