#include "image_encoder.h"
#include "image_codec.h"
#include "codec/encoder_nvenc.h"
#include "image_process/cuda_impl/pixel_convert.cuh"

#include <nlohmann/json.hpp>

using namespace nlohmann;
using namespace image_codec;

struct image_encoder::impl {

    create_config conf;

    struct encoder_store_type {
        impl *pimpl = nullptr;
        std::unique_ptr<encoder_nvenc> nvenc;
        timestamp_type last_clear_ts = 0;

        bool handle_idr() {
            if (auto req_ts = pimpl->last_idr_req_ts;
                    last_clear_ts < pimpl->last_idr_req_ts) {
                last_clear_ts = req_ts;
                return true;
            }
            return false;
        }
    };

    using encoder_map_type =
            std::unordered_map<size_t, encoder_store_type>;
    encoder_map_type enc_map;

    timestamp_type last_idr_req_ts = current_timestamp();

    data_type encode_nvenc(image_ptr img, size_t series) {
        // create NvEnc if needed
        assert(enc_map.contains(series));
        auto &enc_st = enc_map.at(series);
        auto &encoder = enc_st.nvenc;
        auto img_size = img->size();
        if (encoder == nullptr
            || encoder->frame_size() != img_size) [[unlikely]] {
            auto fps = img->get_meta_ext<size_t>(META_REFRESH_RATE);
            auto enc_conf = encoder_nvenc::create_config{
                    .frame_size = img_size,
                    .frame_rate = (int) fps,
                    .bitrate_mbps = conf.bitrate_mbps,
                    .save_file = conf.save_file,
                    .save_length = conf.save_length,
                    .ctx = conf.ctx,
                    .stream = conf.stream,
            };
            encoder = encoder_nvenc::create(enc_conf);
        }
        assert(encoder != nullptr);
        assert(encoder->frame_size() == img_size);

        if (img->cv_type() == CV_8UC3) {
            // rgb -> bgra
            auto img_bgra = create_image(img_size, CV_8UC4);
            call_cvt_rgb_bgra_u8(img->cuda<uchar3>(conf.stream),
                                 img_bgra->cuda<uchar4>(conf.stream),
                                 conf.stream->cuda);
            img_bgra->cuda_modified(conf.stream);
            img = img_bgra;
            assert(img->cv_type() == CV_8UC4);
        } else {
            assert(img->pixel_format() == PIX_NV12);
        }

        auto frame = encoder->encode(img, enc_st.handle_idr());
        return frame.data;
    }

    data_type encode(const image_ptr &img) {
        auto series = img->get_meta_ext<size_t>(META_SERIES_NAME);
        if (!enc_map.contains(series)) [[unlikely]] {
            enc_map.emplace(std::piecewise_construct,
                            std::forward_as_tuple(series),
                            std::forward_as_tuple(this));
        }

        auto enc_type = conf.type;
        size_t sp_id = 0;
        if (auto sp = img->get_meta_any(META_IMAGE_SPECIAL_CODEC); !sp.empty()) {
            enc_type = ENCODER_SPECIAL;
            sp_id = boost::any_cast<size_t>(sp);
        }

        auto head = json();
        head["series"] = series;
        head["type"] = enc_type;
        head["special"] = sp_id;
        head["width"] = img->width();
        head["height"] = img->height();
        head["pix_fmt"] = img->pixel_format();

        auto ret = data_type();
        switch (enc_type) {
            case ENCODER_NVENC: {
                ret = encode_nvenc(img, series);
                break;
            }
            case ENCODER_SPECIAL: {
                auto &enc_st = enc_map.at(series);
                ret = get_special_encoder(sp_id)(img, enc_st.handle_idr());
                break;
            }
            default: {
                RET_ERROR_E;
            }
        }

        auto writer = network_writer();
        writer.write_with_length(head);
        writer.write_with_length(ret);
        return writer.current_data();
    }

};

image_encoder::image_encoder(create_config conf)
        : pimpl(std::make_unique<impl>()) {
    pimpl->conf = conf;
}

image_encoder::~image_encoder() = default;

data_type image_encoder::encode(const image_ptr &img) {
    return pimpl->encode(img);
}

void image_encoder::clear_cache() {
    pimpl->last_idr_req_ts = current_timestamp();
}