#include "frame_encoder/encoder_jpeg.h"
#include "frame_encoder/encoder_nvenc.h"
#include "frame_sender/sender_base.h"
#include "imgui_utility.h"
#include "impl_types.h"
#include "mvs_camera.h"
#include "simple_mq.h"
#include "variable_defs.h"
#include "utility.hpp"

#include <queue>

using namespace simple_mq_singleton;

bool encoder_save_file;
bool encoder_save_length;

encoder_type chosen_encoder = ENCODER_JPEG;
nvenc_config main_nvenc_conf;
nvjpeg_config main_nvjpeg_conf;

std::unique_ptr<std::thread> encoder_thread;

extern bool output_full_frame;
extern bool output_halve_width;
extern std::unique_ptr<smart_frame_buffer> output_fbo;

extern sender_type chosen_sender;
extern std::unique_ptr<sender_base> sender;

extern mvs::capture_config capture_conf;
extern std::unique_ptr<camera_related> left;
extern std::unique_ptr<camera_related> right;
extern std::queue<std::function<void()>> simple_eq;

void load_encoder_config(YAML::Node conf) {
    mq().update_variable(OUTPUT_WIDTH, conf["width"].as<int>());
    mq().update_variable(OUTPUT_HEIGHT, conf["height"].as<int>());
    main_nvenc_conf.bitrate_mbps = conf["hevc_bitrate"].as<float>();
    main_nvjpeg_conf.quality = conf["jpeg_quality"].as<int>();
}

void encoder_thread_work() {
    uint64_t last_conf_cnt;
    std::unique_ptr<encoder_base> encoder;
    bool save_file = false;
    switch (chosen_encoder) {
        case ENCODER_NVENC: {
            auto conf = mq().query_variable<nvenc_config>(ENCODER_CONFIG, &last_conf_cnt);
            save_file = conf.save_file;
            encoder.reset(encoder_nvenc::create(conf));
            break;
        }
        case ENCODER_JPEG: {
            auto conf = mq().query_variable<nvjpeg_config>(ENCODER_CONFIG, &last_conf_cnt);
            save_file = conf.save_file;
            encoder.reset(encoder_jpeg::create(conf));
            break;
        }
        default: {
            RET_ERROR;
        }
    }

    uint64_t frame_cnt = 0;
    for (;;) {
        mq().wait_variable(OUTPUT_FRAME, frame_cnt);
        auto frame = mq().query_variable_ptr<cv::cuda::GpuMat>(OUTPUT_FRAME, &frame_cnt);

        // test stop flag
        if (mq().query_variable<bool>(ENCODER_SHOULD_STOP)) break;

        if (!(save_file ||
              mq().query_variable<bool>(SENDER_CONNECTED))) {
            continue; // no need to encode frame
        }

        // check for config update
        uint64_t cur_conf_cnt;
        switch (chosen_encoder) {
            case ENCODER_NVENC: {
                auto conf = mq().query_variable<nvenc_config>(ENCODER_CONFIG, &cur_conf_cnt);
                if (cur_conf_cnt > last_conf_cnt) {
                    auto real_encoder = dynamic_cast<encoder_nvenc *>(encoder.get());
                    assert(real_encoder != nullptr);
                    real_encoder->change_config(conf);
                    last_conf_cnt = cur_conf_cnt;
                }
                break;
            }
            case ENCODER_JPEG: {
                auto conf = mq().query_variable<nvjpeg_config>(ENCODER_CONFIG, &cur_conf_cnt);
                if (cur_conf_cnt > last_conf_cnt) {
                    auto real_encoder = dynamic_cast<encoder_jpeg *>(encoder.get());
                    assert(real_encoder != nullptr);
                    real_encoder->change_config(conf);
                    last_conf_cnt = cur_conf_cnt;
                }
                break;
            }
            default : {
                RET_ERROR;
            }
        }

        bool force_idr = false;
        if (mq().query_variable<bool>(REQUEST_IDR)) {
            force_idr = true;
            mq().update_variable(REQUEST_IDR, false);
        }

        if (frame == nullptr) continue;
        auto frame_data = std::make_unique<video_nal>();
        encoder->encode(*frame, frame_data.get(), force_idr);
        if (mq().query_variable<bool>(SENDER_CONNECTED)) {
            assert(sender != nullptr);
            sender->send_frame(std::move(frame_data));
        }
    }
}

bool is_encoding() {
    return encoder_thread != nullptr;
}

void upload_encoder_config() {
    switch (chosen_encoder) {
        case ENCODER_NVENC: {
            main_nvenc_conf.save_file = encoder_save_file;
            main_nvenc_conf.save_length = encoder_save_length;
            mq().update_variable(ENCODER_CONFIG, main_nvenc_conf);
            break;
        }
        case ENCODER_JPEG: {
            main_nvjpeg_conf.save_file = encoder_save_file;
            main_nvjpeg_conf.save_length = encoder_save_length;
            mq().update_variable(ENCODER_CONFIG, main_nvjpeg_conf);
            break;
        }
        default: {
            RET_ERROR;
        }
    }
}

void start_encoder() {
    // determine output frame size
    cv::Size frame_size;
    if (output_full_frame) {
        assert(!left->img_dev.empty() && !right->img_dev.empty());
        assert(left->img_dev.size() == right->img_dev.size());
        frame_size.width = left->img_dev.size().width * 2;
        frame_size.height = left->img_dev.size().height;
    } else {
        frame_size = {mq().query_variable<int>(OUTPUT_WIDTH),
                      mq().query_variable<int>(OUTPUT_HEIGHT)};
    }
    output_fbo->create(frame_size);

    // update config
    switch (chosen_encoder) {
        case ENCODER_NVENC: {
            main_nvenc_conf.frame_size = frame_size;
            main_nvenc_conf.frame_rate = capture_conf.frame_rate;
            break;
        }
        case ENCODER_JPEG: {
            break;
        }
        default: {
            RET_ERROR;
        }
    }
    upload_encoder_config();

    mq().update_variable(ENCODER_SHOULD_STOP, false);
    mq().update_variable(ENCODER_BUSY, false); // make variable exist
    encoder_thread = std::make_unique<std::thread>(encoder_thread_work);
}

void stop_encoder() {
    if (encoder_thread == nullptr) return;
    mq().update_variable(ENCODER_SHOULD_STOP, true);
    mq().update_variable_ptr<cv::cuda::GpuMat>(OUTPUT_FRAME, nullptr);
    encoder_thread->join();
    encoder_thread.reset();
}

void show_encoder_ui() {
    ImGui::SeparatorText("Method");
    {
        auto guard = imgui_disable_guard{is_encoding()};
        if (ImGui::RadioButton("NvEnc", chosen_encoder == ENCODER_NVENC)) {
            simple_eq.emplace([] {
                chosen_encoder = ENCODER_NVENC;
                if (chosen_sender == SENDER_UDP) { // NvEnc cannot be used with UDP
                    chosen_sender = SENDER_TCP;
                }
            });
        }
        ImGui::SameLine();
        if (ImGui::RadioButton("nvJPEG", chosen_encoder == ENCODER_JPEG)) {
            simple_eq.emplace([] { chosen_encoder = ENCODER_JPEG; });
        }
    }

    ImGui::SeparatorText("Actions");
    if (!is_encoding()) {
        if (ImGui::Button("Start")) {
            simple_eq.emplace(start_encoder);
        }
    } else {
        if (ImGui::Button("Close")) {
            simple_eq.emplace(stop_encoder);
        }
        ImGui::SameLine();
        if (ImGui::Button("Request IDR")) {
            simple_eq.emplace([] { mq().update_variable(REQUEST_IDR, true); });
        }
    }

    ImGui::SeparatorText("Configs");
    switch (chosen_encoder) {
        case ENCODER_NVENC: {
            if (ImGui::DragFloat("Bitrate (Mbps)", &main_nvenc_conf.bitrate_mbps,
                                 0.1, 1, 20, "%.01f")) {
                simple_eq.emplace(upload_encoder_config);
            }
            break;
        }
        case ENCODER_JPEG: {
            if (ImGui::DragInt("Quality (%)", &main_nvjpeg_conf.quality, 1, 1, 100)) {
                simple_eq.emplace(upload_encoder_config);
            }
            break;
        }
        default: {
            RET_ERROR;
        }
    }
    {
        auto guard = imgui_disable_guard{is_encoding()};
        ImGui::Checkbox("Full Resolution", &output_full_frame);
        if (!output_full_frame) {
            ImGui::SameLine();
            ImGui::Checkbox("Halve Width", &output_halve_width);
        }
//                ImGui::SameLine();
        ImGui::Checkbox("Save Video", &encoder_save_file);
        if (encoder_save_file) {
            ImGui::SameLine();
            ImGui::Checkbox("Save Frame Length", &encoder_save_length);
        }
    }
}