#include "remote_ar.h"
#include "core/imgui_utility.hpp"
#include "core/math_helper.hpp"
#include "core/yaml_utility.hpp"
#include "network/binary_utility.hpp"

#include <boost/asio/post.hpp>

using boost::asio::post;

app_remote_ar::app_remote_ar(const create_config &_conf) {
    conf = _conf.ext_conf;
    asio_ctx = _conf.asio_ctx;
    cuda_ctx = _conf.cuda_ctx;

    // initialize object manager
    OBJ_SAVE(raw_left, image_u8c1());
    OBJ_SAVE(raw_right, image_u8c1());
    OBJ_SAVE(rgb_left, image_u8c3());
    OBJ_SAVE(rgb_right, image_u8c3());
    OBJ_SAVE(aug_left, image_u8c3());
    OBJ_SAVE(aug_right, image_u8c3());
    OBJ_SAVE(img_out, image_u8c4()); // ARGB

    OBJ_SAVE(guide_combine, image_u8c1());
    OBJ_SAVE(guide_info, data_type());
    OBJ_SAVE(guide_combine_rgb, image_u8c3());
    OBJ_SAVE(guide_img, image_u8c3());
    OBJ_SAVE(guide_depth_fake, image_u8c3());
    OBJ_SAVE(guide_depth, image_f32c1());

    // process callbacks caused by OBJ_SAVE
    asio_ctx->poll();

    // initialize sophiar
    sophiar_thread = std::make_unique<std::thread>([conf_path = LOAD_STR("sophiar_config")] {
        sophiar::run_sophiar(conf_path);
    });
    sophiar_conn = std::make_unique<sophiar_conn_type>();
    sophiar_start_var = LOAD_STR("sophiar_start_var");

    // initialize modules
    auto mvs_conf = mvs_camera_ui::create_config{.ctx = asio_ctx};
    mvs_conf.cameras.push_back({.dev_name = LOAD_STR("left_camera_name"), .img_name = raw_left});
    mvs_conf.cameras.push_back({.dev_name = LOAD_STR("right_camera_name"), .img_name = raw_right});
    mvs_cam = std::make_unique<mvs_camera_ui>(mvs_conf);
    mvs_cam->cap_info_sig.connect([this](auto info) {
        out_streamer->change_frame_rate(info.frame_rate);
    });

    auto stereo_info = stereo_camera_info::from_yaml(LOAD_SUB("stereo_info"));
    float view_angle = 0.0f;
    auto img_range = calc_valid_range(stereo_info.left, stereo_info.right, &view_angle);
    view_angle = glm::degrees(view_angle);

    auto cam_left_conf = image_process_ui::create_config{
            .in_name = raw_left, .out_name = rgb_left, .stream = &cam_left.stream
    };
    cam_left_conf.dev_info = {.valid_range = img_range, .cam_int = stereo_info.left};
    cam_left.img_proc = std::make_unique<image_process_ui>(cam_left_conf);

    auto cam_right_conf = image_process_ui::create_config{
            .in_name = raw_right, .out_name = rgb_right, .stream = &cam_right.stream
    };
    cam_right_conf.dev_info = {.valid_range = img_range, .cam_int = stereo_info.right};
    cam_right.img_proc = std::make_unique<image_process_ui>(cam_right_conf);
    cam_right.img_proc->sync_with(cam_left.img_proc.get());

    mvs_cam->cap_info_sig.connect([this](auto info) {
        cam_left.img_proc->change_config({.is_mono = info.is_mono});
        cam_right.img_proc->change_config({.is_mono = info.is_mono});
    });

    auto guide_in_conf = image_player::create_config{
            .img_name = guide_combine, .ext_name = guide_info,
            .ctx = asio_ctx,
    };
    guide_player = std::make_unique<image_player>(guide_in_conf);

    auto guide_cvt_conf = versatile_convertor::create_config{
            .in_name = guide_combine, .out_name = guide_combine_rgb,
            .cvt_opt = CVT_NV12_RGB, .stream = default_cuda_stream,
    };
    guide_cvt = std::make_unique<versatile_convertor>(guide_cvt_conf);

    auto guide_split_conf = versatile_convertor::create_config{
            .in_name = guide_combine_rgb, .out_name = guide_img, .ext_out = guide_depth_fake,
            .cvt_opt = CVT_HALF_SPLIT, .stream = default_cuda_stream,
    };
    guide_split = std::make_unique<versatile_convertor>(guide_split_conf);

    auto guide_decode_conf = versatile_convertor::create_config{
            .in_name = guide_depth_fake, .ext_in = guide_info, .out_name = guide_depth,
            .cvt_opt = CVT_FAKE_DECODE_800P, .stream = default_cuda_stream,
    };
    guide_decode = std::make_unique<versatile_convertor>(guide_decode_conf);

    auto guide_control_conf = depth_guide_controller::create_config{
            .img_in = guide_img, .depth_in = guide_depth, .img_out = guide_final
    };
    guide_controller = std::make_unique<depth_guide_controller>(guide_control_conf);

    auto aug_list_v1 = augment_manager::item_list_from_yaml(LOAD_LIST("augment_list"));
    auto aug_list = augment_manager_v2::item_list_from_v1(aug_list_v1);
    auto aug_conf = augment_manager_v2::create_config{
            .item_list = aug_list,
            .sophiar_conn = sophiar_conn.get(),
            .stream = default_cuda_stream,
    };
    aug_manager = std::make_unique<augment_manager_v2>(aug_conf);

    auto left_aug_conf = camera_augment_helper_v2::create_config{
            .camera = camera_augment_helper_v2::create_config::fixed_camera_config{
                    .fov = view_angle, .transform_var = LOAD_STR("left_camera_transform_var"),
            },
            .sophiar_conn = sophiar_conn.get(), .manager = aug_manager.get()
    };
    cam_left.aug_helper = std::make_unique<camera_augment_helper_v2>(left_aug_conf);

    auto right_aug_conf = camera_augment_helper_v2::create_config{
            .camera = camera_augment_helper_v2::create_config::relative_camera_config{
                    .parent = cam_left.aug_helper.get(),
                    .transform = glm::inverse(to_mat4(stereo_info.transform)),
            },
            .sophiar_conn = sophiar_conn.get(), .manager = aug_manager.get()
    };
    cam_right.aug_helper = std::make_unique<camera_augment_helper_v2>(right_aug_conf);

//    auto right_aug_conf = left_aug_conf;
//    std::get<2>(right_aug_conf.camera).transform_var = LOAD_STR("right_camera_transform_var");
//    cam_right.aug_helper = std::make_unique<camera_augment_helper_v2>(right_aug_conf);
////    cam_right.aug_helper->sync_with(cam_left.aug_helper.get());

    auto left_aug_ren_conf = image_augment_helper::create_config{
            .in_name = rgb_left, .out_name = aug_left,
            .flip_image = true, .stream = &cam_left.stream
    };
    left_aug_ren_conf.render_func_list.emplace_back([this] {
        cam_left.aug_helper->render();
    });
    left_aug_ren_conf.render_func_list.emplace_back([this] {
        guide_controller->render({.stream = &cam_left.stream}); // TODO: add camera specific options
    });
    cam_left.aug_render = std::make_unique<image_augment_helper>(left_aug_ren_conf);

    auto right_aug_ren_conf = image_augment_helper::create_config{
            .in_name = rgb_right, .out_name = aug_right,
            .flip_image = true, .stream = &cam_right.stream
    };
    right_aug_ren_conf.render_func_list.emplace_back([this] {
        cam_right.aug_helper->render();
    });
    right_aug_ren_conf.render_func_list.emplace_back([this] {
        guide_controller->render({.stream = &cam_right.stream});
    });
    cam_right.aug_render = std::make_unique<image_augment_helper>(right_aug_ren_conf);

    auto output_size = cv::Size(
            LOAD_NUMBER(int, "output_width"),
            LOAD_NUMBER(int, "output_height"));
    auto stereo_aug_conf = stereo_augment_helper::create_config{
            .left_name = aug_left, .right_name = aug_right,
            .out_name = img_out, .fbo_size = output_size,
            .flip_image = true, .stream = default_cuda_stream
    };
    stereo_aug = std::make_unique<stereo_augment_helper>(stereo_aug_conf);

    auto streamer_conf = image_streamer::create_config{
            .img_name = img_out, .asio_ctx = asio_ctx,
            .cuda_ctx = cuda_ctx, .stream = default_cuda_stream
    };
    out_streamer = std::make_unique<image_streamer>(streamer_conf);
    out_streamer->sig_req_size.connect([this](auto size) {
        post(*asio_ctx, [=, this] { stereo_aug->resize(size); });
    });

    auto bg_viewer_conf = image_viewer::create_config{
            .mode = VIEW_STEREO, .flip_y = false,
            .stream = default_cuda_stream,
    };
    auto &stereo_conf = bg_viewer_conf.extra.stereo;
    stereo_conf.c_fmt = COLOR_RGB;
    stereo_conf.left_name = aug_left;
    stereo_conf.right_name = aug_right;
    bg_viewer = std::make_unique<image_viewer>(bg_viewer_conf);

//    auto bg_viewer_conf = image_viewer::create_config{
//            .mode = VIEW_COLOR_DEPTH, .flip_y = true,
//            .stream = default_cuda_stream,
//    };
//    auto &bg_extra_conf = bg_viewer_conf.extra.color_depth;
//    bg_extra_conf.c_fmt = COLOR_RGB;
//    bg_extra_conf.c_name = guide_img;
//    bg_extra_conf.d_name = guide_depth;
//    bg_viewer = std::make_unique<image_viewer>(bg_viewer_conf);

    auto saver_conf = image_saver::create_config{.ctx = asio_ctx};
    saver_conf.img_list.emplace_back("Left", rgb_left);
    saver_conf.img_list.emplace_back("Right", rgb_right);
    debug_saver = std::make_unique<image_saver>(saver_conf);

    auto reg_conf = registration_config{
            .conn = sophiar_conn.get(),
            .probe_model_path = LOAD_STR("probe_model"),
    };
    reg.reset(registration::create(reg_conf));
    for (auto reg_item: LOAD_LIST("registration_list")) {
        auto item_conf = registration_target::from_yaml(reg_item);
        reg->add_target(item_conf);
    }
}

void app_remote_ar::start_tracking() {
//    sophiar_conn->start_object(sophiar_start_var);
    // work in another thread to prevent blocking
    auto t = std::thread([this] {
        auto conn = sophiar::local_connection();
        conn.start_object(sophiar_start_var);
    });
    t.detach();
}

app_remote_ar::~app_remote_ar() {
    // sophiar
    sophiar::stop_sophiar();
    sophiar_thread->join();
}

void app_remote_ar::show_ui() {
    if (ImGui::Begin("Remote AR Control")) {
        ImGui::PushItemWidth(200);

        if (ImGui::CollapsingHeader("Camera")) {
            auto id_guard = imgui_id_guard("camera");
            mvs_cam->show();
            ImGui::SeparatorText("Post-Processing");
            cam_left.img_proc->show();
        }

        if (ImGui::CollapsingHeader("Augment")) {
            ImGui::SeparatorText("Tracker");
            if (ImGui::Button("Start")) {
                start_tracking();
            }
            ImGui::SameLine();
            ImGui::Checkbox("Registration", &enable_reg);
            {
                ImGui::SeparatorText("Scene");
                auto id_guard = imgui_id_guard("augment_scene");
                aug_manager->show();
            }
            {
                ImGui::SeparatorText("Camera");
                auto id_guard = imgui_id_guard("augment_camera");
                cam_left.aug_helper->show();
            }
        }

        if (ImGui::CollapsingHeader("Depth Guide")) {
            auto id_guard = imgui_id_guard("depth_guide");
            guide_player->show();
            ImGui::SeparatorText("Display");
            guide_controller->show();
        }

        if (ImGui::CollapsingHeader("Stereo")) {
            auto id_guard = imgui_id_guard("stereo");
            stereo_aug->show();
        }

        if (ImGui::CollapsingHeader("Streamer")) {
            auto id_guard = imgui_id_guard("streamer");
            out_streamer->show();
        }

        if (ImGui::CollapsingHeader("Debug")) {
            if (ImGui::TreeNode("Background")) {
                bg_viewer->show();
                ImGui::TreePop();
            }
            if (ImGui::TreeNode("Memory Pool")) {
                if (ImGui::Button("Purge")) {
                    post(*asio_ctx, [] { g_memory_manager->purify(); });
                }
                ImGui::TreePop();
            }
            if (ImGui::TreeNode("Performance")) {
                ImGui::Text("UI Refresh Rate: %.2fms", perf_timer.query().interval);
                ImGui::TreePop();
            }
            if (ImGui::TreeNode("Image Saver")) {
                debug_saver->show();
                ImGui::TreePop();
            }
        }

        ImGui::PopItemWidth();
    }
    ImGui::End();

    if (enable_reg) {
        reg->process();
        reg->show();
    }

    perf_timer.record();
}

void app_remote_ar::render_background() {
    bg_viewer->render();
}