#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 <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 {
        auto left_host = left->host();
        auto right_host = right->host();
        assert(left_host.isContinuous());
        assert(right_host.isContinuous());

        auto ts_str = fmt::format("{}", system_timestamp());
        auto left_name = fmt::format("left_{}.raw", ts_str);
        auto right_name = fmt::format("right_{}.raw", ts_str);
        auto left_path = prefix_path / left_name;
        auto right_path = prefix_path / right_name;
        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());
    }

private:
    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);
    }
};

bool mono_mode = false;
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
bool save_next_raw = false;

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 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::wait_frame(simple_mq::index_type index) const {
    uint64_t cur_cnt;
    if (auto ptr = mq().query_variable_ptr<cv::Mat>(index, &cur_cnt);
            ptr == nullptr || cur_cnt <= raw_cnt) {
        mq().wait_variable(index, raw_cnt);
    }
}

void camera_related::process_frame(simple_mq::index_type index) {
    uint64_t cur_cnt;
    auto raw_ptr = mq().query_variable_ptr<cv::Mat>(index, &cur_cnt);
    assert(cur_cnt > raw_cnt);
    raw_cnt = cur_cnt;
    raw_img = *raw_ptr;

    // OpenCV debayer does not support alpha channel
    if (undistort_image) {
        use_crude_debayer = true;
    }

    // Undistort is currently not supported in mono mode
    if (mono_mode) {
        undistort_image = false;
    }

    // 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
    processor.process(raw_img.dev(stream.cv), &img_dev, process_conf, stream.cv);

    // 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"];
    stereo_info.transform =
            Eigen::Translation3f(
                    trans_info["tx"].as<float>(),
                    trans_info["ty"].as<float>(),
                    trans_info["tz"].as<float>()
            ) * Eigen::Quaternionf(
                    trans_info["qw"].as<float>(),
                    trans_info["qx"].as<float>(),
                    trans_info["qy"].as<float>(),
                    trans_info["qz"].as<float>());
    stereo_info.left = &left->intrinsic;
    stereo_info.right = &right->intrinsic;

    // 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->wait_frame(IMG_RAW_HOST_LEFT);
    right->wait_frame(IMG_RAW_HOST_RIGHT);
}

void process_camera_frames() {
    left->process_frame(IMG_RAW_HOST_LEFT);
    right->process_frame(IMG_RAW_HOST_RIGHT);

    // save image if required
    if (save_next_raw) {
        if (image_saver == nullptr) [[unlikely]] {
            image_saver = std::make_unique<simple_image_saver>();
        }
        image_saver->save_image_raw(&left->raw_img, &right->raw_img);
        save_next_raw = false;
    }
}

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);
            }
            ImGui::SameLine();
            if (ImGui::Button("Capture Raw")) {
                simple_eq.emplace([&] { save_next_raw = true; });
            }
        }
    }

    // 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 (!mono_mode && debug_options) {
        ImGui::SameLine();
        ImGui::Checkbox("Undistort", &undistort_image);
        {
            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");
        }

        // auto save raw config
//                ImGui::SeparatorText("Auto Shoot");
//                ImGui::PushID("Auto Shoot");
//
//                if (!auto_save_raw) {
//                    if (ImGui::Button("Start")) {
//                        auto_save_raw = true;
//                    }
//                } else {
//                    if (ImGui::Button("Stop")) {
//                        auto_save_raw = false;
//                    }
//                }
//
//                if (auto_save_raw_interval < 1) {
//                    auto_save_raw_interval = 1;
//                }
//
//                if (auto_save_raw) {
//                    ImGui::BeginDisabled();
//                    auto now_time = std::chrono::system_clock::now();
//                    if (now_time - last_save_raw_time >
//                        std::chrono::seconds{auto_save_raw_interval}) {
//                                camera.request_save_raw();
//                        last_save_raw_time = now_time;
//                    }
//                }
//                ImGui::InputInt("Shoot Interval (s)", &auto_save_raw_interval, 1, 10);
//                if (auto_save_raw) {
//                    ImGui::EndDisabled();
//                }
    }
}