#include "depth_guide.h"
#include "core/image_utility.hpp"
#include "core/imgui_utility.hpp"
#include "core/pc_utility.h"
#include "image_process/impl/versatile_convertor_impl.h"
#include "image_process/process_funcs.h"

app_depth_guide::app_depth_guide(const create_config &_conf) {
    conf = _conf;

    // initialize object manager
    OBJ_SAVE(img_color, image_u8c3());
    OBJ_SAVE(img_depth, image_f32c1());
    OBJ_SAVE(img_remap, image_ptr());
    OBJ_SAVE(img_depth_fake, image_u8c3());
    auto fake_info = fake_color_config{.mode = FAKE_800P, .lower = 200, .upper = 1000};
    OBJ_SAVE(img_depth_fake_info, versatile_convertor_impl::encode_config(fake_info));
    OBJ_SAVE(img_out, image_u8c4());
    OBJ_SAVE(img_obj_mask, image_ptr());
    OBJ_SAVE(pc_raw, pc_ptr());
    OBJ_SAVE(pc_obj_raw, pc_ptr());
    OBJ_SAVE(pc_obj, pc_ptr());

    OBJ_SIG(img_obj_mask)->connect([](auto _) {
        gen_pc_rgb_with_mask::call(
                {img_color, img_depth, img_obj_mask,
                 img_remap, pc_obj_raw, default_cuda_stream});
    });
    OBJ_SIG(pc_obj_raw)->connect([](auto _) {
        shrink_pc_zero_depth::call(
                {pc_obj_raw, pc_obj, default_cuda_stream});
    });

    // initialize modules
    auto orb_cam_conf = orb_camera_ui::create_config{
            .cf_name = img_color, .df_name = img_depth,
            .pc_name = pc_raw, .remap_name = img_remap,
            .ctx = conf.asio_ctx, .stream = default_cuda_stream,
    };
    orb_cam = std::make_unique<orb_camera_ui>(orb_cam_conf);

    auto fake_conf = versatile_convertor::create_config{
            .in_name = img_depth, .ext_in = img_depth_fake_info, .out_name = img_depth_fake,
            .cvt_opt = CVT_FAKE_ENCODE_800P, .stream = default_cuda_stream,
    };
    depth_encode = std::make_unique<versatile_convertor>(fake_conf);

    auto ai_conf = fast_sam::create_config{
            .img_in = img_color, .mask_out = img_obj_mask,
            .cuda_ctx = conf.cuda_ctx, .ctx = conf.asio_ctx,
    };
    ai_seg = std::make_unique<fast_sam>(ai_conf);

    auto out_conf = stereo_augment_helper::create_config{
            .left_name = img_color, .right_name = img_depth_fake, .out_name = img_out,
            .stream = default_cuda_stream
    };
    out_combiner = std::make_unique<stereo_augment_helper>(out_conf);
    out_combiner->fix_ui_config({.follow_image_size=true, .enable_halve_width=false});

    auto bg_viewer_conf = image_viewer::create_config{
            .mode = VIEW_STEREO, .flip_y = true,
            .stream = default_cuda_stream,
    };
    auto &bg_extra_conf = bg_viewer_conf.extra.stereo;
//    bg_extra_conf.c_name = img_color;
    bg_extra_conf.left_name = img_color;
    bg_extra_conf.c_fmt = COLOR_RGB;
    bg_extra_conf.right_name = img_obj_mask;
    bg_viewer = std::make_unique<image_viewer>(bg_viewer_conf);

    auto out_streamer_conf = image_streamer::create_config{
            .img_name = img_out,
            .ext_name = img_depth_fake_info, // comment this for tiny player to work
            .asio_ctx = conf.asio_ctx,
            .cuda_ctx = conf.cuda_ctx, .stream = default_cuda_stream
    };
    out_streamer = std::make_unique<image_streamer>(out_streamer_conf);

    auto saver_conf = image_saver::create_config{.ctx = conf.asio_ctx};
    saver_conf.img_list.emplace_back("Image", img_color);
    saver_conf.img_list.emplace_back("Image Out", img_out);
    saver_conf.img_list.emplace_back("PC Raw", pc_raw);
    saver_conf.img_list.emplace_back("PC Object", pc_obj);
    out_saver = std::make_unique<image_saver>(saver_conf);

    auto pc_conf = pc_viewer::create_config{
            .name = pc_raw, .stream = default_cuda_stream,
    };
    scene_viewer = std::make_unique<pc_viewer>(pc_conf);
}

void app_depth_guide::show_ui() {
    auto ctx = conf.asio_ctx;
    if (ImGui::Begin("Depth Guide Control")) {
        ImGui::PushItemWidth(200);

        if (ImGui::CollapsingHeader("Camera")) {
            auto id_guard = imgui_id_guard("camera");
            orb_cam->show();
        }

        if (ImGui::CollapsingHeader("AI Segmentation")) {
            auto id_guard = imgui_id_guard("ai_segment");
            ai_seg->show();
        }

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

        if (ImGui::CollapsingHeader("Debug")) {
            if (ImGui::TreeNode("Background")) {
                ImGui::RadioButton("Image", &enable_scene_viewer, false);
                ImGui::SameLine();
                ImGui::RadioButton("3D", &enable_scene_viewer, true);
                ImGui::TreePop();
            }
            if (enable_scene_viewer) {
                if (ImGui::TreeNode("3D Viewer")) {
                    scene_viewer->show();
                    ImGui::TreePop();
                }
            } else {
                if (ImGui::TreeNode("Image Viewer")) {
                    bg_viewer->show();
                    ImGui::TreePop();
                }
            }
            if (ImGui::TreeNode("Image Saver")) {
                out_saver->show();
                ImGui::TreePop();
            }
            if (ImGui::TreeNode("Memory Pool")) {
                if (ImGui::Button("Purge")) {
                    post(*ctx, [] { global_mp.purge(); });
                }
                ImGui::TreePop();
            }
            if (ImGui::TreeNode("Performance")) {
                ImGui::Text("UI Refresh Rate: %.2fms", perf_timer.query().interval);
                ImGui::TreePop();
            }
        }

        ImGui::PopItemWidth();
    }
    ImGui::End();
    perf_timer.record();
}

void app_depth_guide::render_background() {
    if (enable_scene_viewer) {
        scene_viewer->render();
    } else {
        bg_viewer->render();
    }
}