#include "core/local_connection.h"
#include "frame_encoder/encoder_nvenc.h"
#include "imgui_utility.h"
#include "impl_types.h"
#include "mvs_camera.h"
#include "variable_defs.h"

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

#include <fmt/chrono.h>

#include <queue>

using namespace simple_mq_singleton;
using namespace sophiar;

struct simple_image_saver {
    std::filesystem::path prefix_path;

    simple_image_saver() {
        prefix_path = fmt::format("./capture_{:%Y_%m_%d_%H_%M_%S}",
                                  now_since_epoch_in_seconds());
        std::filesystem::create_directories(prefix_path);
    }

    void save_image_raw(smart_mat *left, smart_mat *right) const;

    void save_image_png(smart_mat *left, smart_mat *right) const;

    void save_image_png(const smart_texture &left,
                        const smart_texture &right) const;

private:
    auto get_save_path(std::string_view ext) const {
        auto ts_str = fmt::format("{}", system_timestamp());
        auto left_name = fmt::format("left_{}.{}", ts_str, ext);
        auto right_name = fmt::format("right_{}.{}", ts_str, ext);
        auto left_path = prefix_path / left_name;
        auto right_path = prefix_path / right_name;
        return std::make_tuple(left_path, right_path);
    }

    static void save_file(const char *file_path, void *data, size_t length) {
        static constexpr auto block_size = 512 * 1024; // 512KiB
        auto file = fopen(file_path, "wb");
        assert(file != nullptr);
        auto block_cnt = length / block_size;
        auto write_cnt = fwrite(data, block_size, block_cnt, file);
        assert(write_cnt == block_cnt);
        auto write_len = write_cnt * block_size;
        if (length > write_len) {
            fwrite((char *) data + write_len, length - write_len, 1, file);
        }
        fclose(file);
    }

    static void save_png(const std::string &file_path, const cv::Mat &img) {
        cv::Mat out_img = img;
        if (out_img.channels() == 3) { // RGB -> BGR
            cv::cvtColor(img, out_img, cv::COLOR_RGB2BGR);
        }
        cv::imwrite(file_path, out_img);
    }

    static void save_tex(const std::string &file_path, const smart_texture &tex) {
        // download texture
        GLenum format;
        GLsizei width, height;
        glBindTexture(GL_TEXTURE_2D, tex.id);
        glGetTexLevelParameteriv(GL_TEXTURE_2D, 0, GL_TEXTURE_INTERNAL_FORMAT, (int *) &format);
        glGetTexLevelParameteriv(GL_TEXTURE_2D, 0, GL_TEXTURE_WIDTH, &width);
        glGetTexLevelParameteriv(GL_TEXTURE_2D, 0, GL_TEXTURE_HEIGHT, &height);
        auto closer = sg::make_scope_guard([&] {
            glBindTexture(GL_TEXTURE_2D, 0);
        });

        switch (format) {
            case GL_RGBA8: {
                auto img = cv::Mat(height, width, CV_8UC4);
                assert(img.isContinuous());
                glGetTexImage(GL_TEXTURE_2D, 0, GL_BGRA, GL_UNSIGNED_BYTE, img.data);
                save_png(file_path, img);
                return;
            }
            default: {
                RET_ERROR;
            }
        }
    }
};

void simple_image_saver::save_image_raw(smart_mat *left, smart_mat *right) const {
    auto left_host = left->host();
    auto right_host = right->host();
    assert(left_host.isContinuous());
    assert(right_host.isContinuous());
    auto [left_path, right_path] = get_save_path(".raw");
    save_file(left_path.c_str(), left_host.data, left_host.total() * left_host.elemSize());
    save_file(right_path.c_str(), right_host.data, right_host.total() * right_host.elemSize());
}

void simple_image_saver::save_image_png(smart_mat *left, smart_mat *right) const {
    auto [left_path, right_path] = get_save_path(".p.png");
    save_png(left_path, left->host());
    save_png(right_path, right->host());
}

void simple_image_saver::save_image_png(const smart_texture &left,
                                        const smart_texture &right) const {
    auto [left_path, right_path] = get_save_path(".t.png");
    save_tex(left_path, left);
    save_tex(right_path, right);
}

static constexpr auto CAPTURE_RAW_BIT = 1 << 0;
static constexpr auto CAPTURE_PROCESSED_BIT = 1 << 1;
static constexpr auto CAPTURE_AUGMENT_BIT = 1 << 2;
uint8_t capture_flag = 0;
bool capture_raw = false;
bool capture_processed = false;
bool capture_augment = false;

bool mono_mode = true;
bool enhance_image = false;
bool use_crude_debayer = true; // debug options
bool undistort_image = true; // debug options
int preview_camera_index = 0;  // 0 for left, 1 for right

mvs::capture_config capture_conf = {};
stereo_camera_info stereo_info;
std::string left_camera_name, right_camera_name;
std::unique_ptr<mvs::camera> left_camera, right_camera;
std::unique_ptr<simple_image_saver> image_saver;

extern nvenc_config main_nvenc_conf;
extern bool debug_options;
extern float augment_render_angle;
extern float process_frame_rate;
extern bool augment_enable;

extern local_connection sophiar_conn;
extern std::queue<std::function<void()>> simple_eq;

extern std::unique_ptr<camera_related> left;
extern std::unique_ptr<camera_related> right;

void camera_related::load_intrinsic(YAML::Node conf) {
    intrinsic.fx = conf["fx"].as<float>();
    intrinsic.fy = conf["fy"].as<float>();
    intrinsic.cx = conf["cx"].as<float>();
    intrinsic.cy = conf["cy"].as<float>();
    intrinsic.k[0] = conf["k0"].as<float>();
    intrinsic.k[1] = conf["k1"].as<float>();
    intrinsic.width = conf["width"].as<int>();
    intrinsic.height = conf["height"].as<int>();

    process_conf.camera = &intrinsic;
}

void camera_related::get_next_frame(simple_mq::index_type index) {
    uint64_t cur_cnt;
    auto ptr = mq().query_variable_ptr<cv::Mat>(index, &cur_cnt);
    if (ptr == nullptr || cur_cnt <= raw_cnt) {
        mq().wait_variable(index, raw_cnt);
        ptr = mq().query_variable_ptr<cv::Mat>(index, &cur_cnt);
        assert(cur_cnt > raw_cnt);
    }
    raw_cnt = cur_cnt;
    raw_img = *ptr;
}

void camera_related::process_frame() {
    // OpenCV debayer does not support alpha channel
    if (undistort_image && !mono_mode) {
        use_crude_debayer = true;
    }

    // update process config
    process_conf.is_mono = mono_mode;
    process_conf.crude_debayer = use_crude_debayer;
    process_conf.undistort = undistort_image;
    process_conf.enhance = enhance_image;
    process_conf.resample_height = mq().query_variable<int>(OUTPUT_HEIGHT);

    // process image
    auto img_dev = img_mat.dev(stream.cv);
    processor.process(raw_img.dev(stream.cv), &img_dev, process_conf, stream.cv);
    img_mat.set_dev_mat(img_dev);

    // augment render
    process_augment();
}

void load_camera_config(YAML::Node conf) {
    // load camera names
    auto camera_names = conf["names"];
    left_camera_name = camera_names["left"].as<std::string>();
    right_camera_name = camera_names["right"].as<std::string>();
    auto capture_param = conf["capture"];
    capture_conf.frame_rate = capture_param["frame_rate"].as<int>();

    capture_conf.expo_time_ms = capture_param["expo_time_ms"].as<float>();
    capture_conf.gain_db = capture_param["gain_db"].as<float>();

    // load camera intrinsics
    auto intrinsic_conf = conf["intrinsic"];
    left->load_intrinsic(intrinsic_conf["left"]);
    right->load_intrinsic(intrinsic_conf["right"]);

    // load stereo trans
    auto trans_info = conf["stereo_trans"];
    left->transform =
            Eigen::Translation3d(
                    trans_info["tx"].as<double>(),
                    trans_info["ty"].as<double>(),
                    trans_info["tz"].as<double>()
            ) * Eigen::Quaterniond(
                    trans_info["qw"].as<double>(),
                    trans_info["qx"].as<double>(),
                    trans_info["qy"].as<double>(),
                    trans_info["qz"].as<double>());
    right->transform = Eigen::Isometry3d::Identity();
    stereo_info.left = &left->intrinsic;
    stereo_info.right = &right->intrinsic;
    stereo_info.transform = left->transform.cast<float>();

    // calculate valid resample range
    auto range = calc_valid_range(left->intrinsic, right->intrinsic, &augment_render_angle);
    augment_render_angle *= 180 / std::numbers::pi;
    left->process_conf.valid_range = range;
    right->process_conf.valid_range = range;
}

bool is_camera_opened() {
    assert((left_camera == nullptr) == (right_camera == nullptr));
    return left_camera != nullptr;
}

bool is_capturing() {
    if (!is_camera_opened()) return false;
    assert(left_camera->is_capture() == right_camera->is_capture());
    return left_camera->is_capture();
}

bool upload_capture_config_impl() {
    bool ok = true;
    ok &= left_camera->set_capture_config(capture_conf);
    ok &= right_camera->set_capture_config(capture_conf);
    return ok;
}

void upload_capture_config() {
    if (!is_camera_opened()) return;
    if (!upload_capture_config_impl()) {
        // TODO: show error msg
    }
}

void stop_capture() {
    left_camera->stop_capture();
    right_camera->stop_capture();
}

void start_capture() {
    assert(is_camera_opened());
    bool ok = true;
    ok &= upload_capture_config_impl();
    ok &= left_camera->start_capture();
    ok &= right_camera->start_capture();
    if (!ok) {
        // TODO: show error msg
        stop_capture();
    }
}

void close_cameras() {
    if (is_capturing()) {
        stop_capture();
    }
    left_camera.reset();
    right_camera.reset();
}

void open_cameras() {
    auto pixel_type = mono_mode ? mvs::MONO_8 : mvs::RG_8;
    auto left_camera_conf = mvs::create_config{
            left_camera_name, pixel_type, IMG_RAW_HOST_LEFT
    };
    auto right_camera_conf = mvs::create_config{
            right_camera_name, pixel_type, IMG_RAW_HOST_RIGHT
    };
    left_camera.reset(mvs::camera::create(left_camera_conf));
    right_camera.reset(mvs::camera::create(right_camera_conf));
    if (left_camera == nullptr || right_camera == nullptr) {
        // TODO: show error msg
        close_cameras();
    }
}

void upload_encoder_config();

void wait_camera_frames() {
    assert(is_capturing());
    left->get_next_frame(IMG_RAW_HOST_LEFT);
    right->get_next_frame(IMG_RAW_HOST_RIGHT);
}

void set_capture_flag() {
    capture_flag = 0;
    capture_flag |= capture_raw ? CAPTURE_RAW_BIT : 0;
    capture_flag |= capture_processed ? CAPTURE_PROCESSED_BIT : 0;
    capture_flag |= capture_augment ? CAPTURE_AUGMENT_BIT : 0;

    if (image_saver == nullptr) {
        image_saver = std::make_unique<simple_image_saver>();
    }
}

void process_camera_frames() {
    if (capture_flag & CAPTURE_RAW_BIT) {
        image_saver->save_image_raw(&left->raw_img, &right->raw_img);
    }

    left->process_frame();
    right->process_frame();

    if (capture_flag & CAPTURE_PROCESSED_BIT) {
        image_saver->save_image_png(&left->img_mat, &right->img_mat);
    }
    if (capture_flag & CAPTURE_AUGMENT_BIT) {
        image_saver->save_image_png(left->augment_tex, right->augment_tex);
    }

    // reset capture flag
    capture_flag = 0;
}

void show_camera_ui() {
    // camera actions
    ImGui::SeparatorText("Actions");
    if (!is_camera_opened()) {
        if (ImGui::Button("Open")) {
            simple_eq.emplace(open_cameras);
        }
    } else { // cameras have been opened
        if (ImGui::Button("Close")) {
            simple_eq.emplace(close_cameras);
        }
        ImGui::SameLine();
        if (!is_capturing()) {
            if (ImGui::Button("Start")) {
                simple_eq.emplace(start_capture);
            }
        } else {
            if (ImGui::Button("Stop")) {
                simple_eq.emplace(stop_capture);
            }
        }
    }

    // camera configs
    ImGui::SeparatorText("Configs");
    if (ImGui::DragInt("Frame Rate (fps)", &capture_conf.frame_rate, 1, 1, 60)) {
        simple_eq.emplace(upload_capture_config);
        main_nvenc_conf.frame_rate = capture_conf.frame_rate;
        simple_eq.emplace(upload_encoder_config);
    }
    if (ImGui::DragFloat("Exposure Time (ms)", &capture_conf.expo_time_ms,
                         0.1, 0.1, 1e3f / capture_conf.frame_rate, "%.01f")) {
        simple_eq.emplace(upload_capture_config);
    }
    if (ImGui::DragFloat("Analog Gain (dB)", &capture_conf.gain_db,
                         0.1, 0, 23.4, "%.01f")) {
        simple_eq.emplace(upload_capture_config);
    }
    {
        auto guard = imgui_disable_guard{is_camera_opened()};
        ImGui::Checkbox("Mono", &mono_mode);
    }
    ImGui::SameLine();
    ImGui::Checkbox("Enhance", &enhance_image);
    if (debug_options) {
        ImGui::SameLine();
        ImGui::Checkbox("Undistort", &undistort_image);
        if (!mono_mode) {
            auto guard = imgui_disable_guard{undistort_image};
            ImGui::SameLine();
            ImGui::Checkbox("Crude Debayer", &use_crude_debayer);
        }
    }

    if (is_capturing()) {
        // preview config
        ImGui::SeparatorText("Preview Camera");
        ImGui::RadioButton("Left", &preview_camera_index, 0);
        ImGui::SameLine();
        ImGui::RadioButton("Right", &preview_camera_index, 1);

        ImGui::SeparatorText("Infos");
        {
            auto guard = imgui_disable_guard{};
            ImGui::DragFloat("Process Frame Rate (fps)", &process_frame_rate, 0, 0, 60, "%.01f");
        }

        ImGui::SeparatorText("Capture Config");
        ImGui::Checkbox("Raw", &capture_raw);
        ImGui::SameLine();
        ImGui::Checkbox("Processed", &capture_processed);
        if (augment_enable) {
            ImGui::SameLine();
            ImGui::Checkbox("Augment", &capture_augment);
        }
        ImGui::SameLine();
        if (ImGui::Button("Confirm")) {
            simple_eq.emplace([&] { set_capture_flag(); });
        }
    }
}