#include "camera_augment_helper_v2_impl.h"
#include "core/math_helper.hpp"
#include "core/imgui_utility.hpp"
#include "module/guidance/cursor_guide.h"

#include <glm/gtc/matrix_transform.hpp>
#include <glm/gtc/type_ptr.hpp>
#include <glm/gtx/rotate_vector.hpp>

#include <boost/asio/post.hpp>

using boost::asio::post;

void camera_augment_helper_v2::impl::freedom_info_type::reset() {
    if (auto &mat = pimpl->tracked_transform; mat.has_value()) {
        static constexpr auto default_focal_length = 8.0f; // 8mm
        eye = glm::vec3((*mat)[3]);
        auto rot_part = glm::mat3(*mat);
        center = eye + rot_part[2] * default_focal_length;
        view_up = -rot_part[1];
    } else {
        eye = glm::vec3();
        center = glm::vec3(0.f, 0.f, 100.f);
        view_up = glm::vec3(0.0f, -1.0f, 0.0f);
    }
}

glm::mat4 camera_augment_helper_v2::impl::freedom_info_type::to_transform() {
    // normalize view up
    view_up -= view_dir() * glm::dot(view_up, view_dir());
    view_up = glm::normalize(view_up);

    auto ret = glm::inverse(glm::lookAt(eye, center, view_up));
    ret = glm::rotate(ret, glm::radians(180.0f), glm::vec3(1.0f, 0.0f, 0.0f)); // viewport -> camera
    return ret;
}

glm::vec3 camera_augment_helper_v2::impl::freedom_info_type::view_dir() const {
    return glm::normalize(center - eye);
}

float camera_augment_helper_v2::impl::freedom_info_type::view_dis() const {
    return glm::distance(center, eye);
}

void camera_augment_helper_v2::impl::freedom_info_type::translate(glm::vec3 offset) {
    eye += offset;
    center += offset;
}

void camera_augment_helper_v2::impl::freedom_info_type::forward(float dis) {
    translate(view_dir() * dis);
}

void camera_augment_helper_v2::impl::freedom_info_type::move_right(float dis) {
    auto axis = glm::cross(view_up, view_dir());
    translate(-axis * dis);
}

void camera_augment_helper_v2::impl::freedom_info_type::move_up(float dis) {
    translate(view_up * dis);
}

void camera_augment_helper_v2::impl::freedom_info_type::center_forward(float dis) {
    center = eye + view_dir() * (view_dis() + dis);
}

void camera_augment_helper_v2::impl::freedom_info_type::roll(float deg) {
    view_up = glm::rotate(view_up, glm::radians(deg), view_dir());
}

void camera_augment_helper_v2::impl::freedom_info_type::yaw(float deg) {
    auto next_view_dir = glm::rotate(view_dir(), glm::radians(deg), view_up);
    center = eye + next_view_dir * view_dis();
}

void camera_augment_helper_v2::impl::freedom_info_type::pitch(float deg) {
    auto rot_axis = glm::cross(view_up, view_dir());
    view_up = glm::rotate(view_up, glm::radians(deg), rot_axis);
    auto next_view_dir = glm::rotate(view_dir(), glm::radians(deg), rot_axis);
    center = eye + next_view_dir * view_dis();
}

glm::mat4 camera_augment_helper_v2::impl::fixed_info_type::extra_transform() const {
    return to_transform_mat(extra_offset, glm::radians(extra_rotation));
}

camera_augment_helper_v2::impl::impl(const create_config &conf) {
    manager = conf.manager;
    act_guide = conf.act_guide;
    sophiar_conn = conf.sophiar_conn;
    ctx = conf.ctx;

    pre_render_conn = manager->pre_render_sig.connect(
            [this](auto info) { pre_render_slot(info); });

    switch (conf.camera.index()) {
        case 1: {
            mode = MODE_FREEDOM;
            auto info = std::get<1>(conf.camera);
            transform_var = info.transform_var;
            freedom_info.pimpl = this;
            break;
        }
        case 2: {
            mode = MODE_FIXED;
            auto info = std::get<2>(conf.camera);
            fov = info.fov;
            transform_var = info.transform_var;
            break;
        }
        case 3: {
            mode = MODE_RELATIVE;
            auto info = std::get<3>(conf.camera);
            relative_info.parent = info.parent->pimpl.get();
            relative_info.transform = info.transform;
            break;
        }
        default: {
            RET_ERROR;
        }
    }
}

camera_augment_helper_v2::impl::~impl() {
    pre_render_conn.disconnect();
}

void camera_augment_helper_v2::impl::update_tracked_transform() {
    if (sophiar_conn == nullptr) return;
    auto trans = sophiar_conn->
            query_transform_variable(transform_var);
    is_missing = !trans.has_value();
    if (trans.has_value()) {
        auto cur_ts = current_timestamp();
        track_pool.add(to_mat4(*trans), cur_ts);
        tracked_transform = track_pool.query(cur_ts + time_offset);
        track_pool.remove_old(-time_offset);
    } else {
        tracked_transform = {};
    }
}

void camera_augment_helper_v2::impl::update_freedom() {
    update_tracked_transform();
    transform = freedom_info.to_transform();
    is_missing = false;
}

void camera_augment_helper_v2::impl::update_fixed() {
    assert(sophiar_conn != nullptr);
    update_tracked_transform();
    if (tracked_transform.has_value()) {
        auto extra_mat = fixed_info.extra_transform();
        transform = *tracked_transform * extra_mat;
    }
}

void camera_augment_helper_v2::impl::update_relative() {
    auto par = relative_info.parent;
    par->update_transform();
    is_missing = par->is_missing;
    transform = par->transform * relative_info.transform;
    fov = par->fov;
}

void camera_augment_helper_v2::impl::update_transform() {
    switch (mode) {
        // @formatter:off
        case MODE_FREEDOM: { update_freedom(); break; }
        case MODE_RELATIVE: { update_relative(); break; }
        case MODE_FIXED: { update_fixed(); break; }
        // @formatter:on
    }
}

void camera_augment_helper_v2::impl::update_camera() {
    update_transform();
    last_camera = camera_info{
            .transform = transform, .fov = fov,
            .near = ui->near, .far = ui->far
    };
}

void camera_augment_helper_v2::impl::render() {
    update_camera();
    if (!is_missing || ui->ignore_missing) {
        last_vp_size = query_viewport_size();
        manager->record_current_camera_helper(q_this);
        manager->render(last_camera);
    }
}

void camera_augment_helper_v2::impl::render_image(output_config conf) {
    fbo_conf.size = conf.size;
    fbo.create(fbo_conf);

    if (img_downloader == nullptr) [[unlikely]] {
        img_downloader = std::make_unique<viewport_downloader>(
                viewport_downloader::create_config{.stream = conf.stream});
    }

    fbo.bind();
    render();
    auto out_img = img_downloader->download_v2();
    OBJ_SAVE(conf.img_name, out_img);
    fbo.unbind();
}

void camera_augment_helper_v2::impl::show_freedom() {
    auto &info = freedom_info;

    if (ImGui::Button("Reset")) {
        freedom_info.reset();
    }

    if (ImGui::TreeNode("Center")) {
        if (ImGui::Button("X+")) { info.move_right(info.linear_speed); }
        ImGui::SameLine();
        if (ImGui::Button("Y+")) { info.move_up(info.linear_speed); }
        ImGui::SameLine();
        if (ImGui::Button("Z+")) { info.forward(info.linear_speed); }
//        ImGui::SameLine();
//        if (ImGui::Button("D+")) { info.center_forward(info.linear_speed); }

        if (ImGui::Button("X-")) { info.move_right(-info.linear_speed); }
        ImGui::SameLine();
        if (ImGui::Button("Y-")) { info.move_up(-info.linear_speed); }
        ImGui::SameLine();
        if (ImGui::Button("Z-")) { info.forward(-info.linear_speed); }
//        ImGui::SameLine();
//        if (ImGui::Button("D-")) { info.center_forward(-info.linear_speed); }

        ImGui::TreePop();
    }

    if (ImGui::TreeNode("Eye")) {
        if (ImGui::Button("Ro+")) { info.roll(info.angle_speed); } // roll
        ImGui::SameLine();
        if (ImGui::Button("Ya+")) { info.yaw(info.angle_speed); } // yaw
        ImGui::SameLine();
        if (ImGui::Button("Pi+")) { info.pitch(info.angle_speed); }  // pitch

        if (ImGui::Button("Ro-")) { info.roll(-info.angle_speed); } // roll
        ImGui::SameLine();
        if (ImGui::Button("Ya-")) { info.yaw(-info.angle_speed); } // yaw
        ImGui::SameLine();
        if (ImGui::Button("Pi-")) { info.pitch(-info.angle_speed); } // pitch

        ImGui::TreePop();
    }

    if (ImGui::TreeNode("Extra")) {
        ImGui::SliderFloat("Angular Step", &info.angle_speed, 0.01f, 10.0f, "%.2f");
        ImGui::SliderFloat("Linear Step", &info.linear_speed, 0.01f, 100.0f, "%.2f");
        ImGui::DragFloat("FOV (deg)", &fov, 0.5f, 10.0f, 178.0f, "%.01f");
        ImGui::TreePop();
    }
}

void camera_augment_helper_v2::impl::show_fixed() {
    ImGui::DragFloat3("Offset (mm)", glm::value_ptr(fixed_info.extra_offset),
                      0.05f, 0.0f, 0.0f, "%.02f");
    ImGui::DragFloat3("Offset (deg)", glm::value_ptr(fixed_info.extra_rotation),
                      0.1f, -180.0f, 180.0f, "%.01f");
}

void camera_augment_helper_v2::impl::show() {
    ImGui::Checkbox("Ignore Missing", &ui->ignore_missing); // TODO: enable sync with another
    if (allow_active_guide) {
        ImGui::SameLine();
        ImGui::Checkbox("Guidance", &ui->enable_active_guide);
    }
    ImGui::SliderFloat("Clip Near", &ui->near, 1.0f, ui->far, "%.f", ImGuiSliderFlags_Logarithmic);
    ImGui::SliderFloat("Clip Far", &ui->far, ui->near, 10000.0f, "%.f", ImGuiSliderFlags_Logarithmic);

    switch (mode) {
        // @formatter:off
        case MODE_FIXED: { show_fixed(); break; }
        case MODE_RELATIVE: { break; }
        case MODE_FREEDOM: { show_freedom(); break; }
        // @formatter:on
        default: {
            RET_ERROR;
        }
    }

    if (ui->enable_active_guide) {
        ImGui::SeparatorText("Cursor Guidance");
        act_guide->show();
    }
}

std::optional<glm::vec3> camera_augment_helper_v2::impl::
get_cursor_coordinate(glm::vec2 s_pos) {
    auto w = last_vp_size.width, h = last_vp_size.height;
    if (s_pos.x < 0 || s_pos.x > w) return {};
    if (s_pos.y < 0 || s_pos.y > h) return {};
    s_pos.y = h - s_pos.y; // flip y to texture coordinate

    float depth;
    glReadPixels(s_pos.x, s_pos.y, 1, 1, GL_DEPTH_COMPONENT, GL_FLOAT, &depth);
    if (depth == 1.f) return {};

    auto proj = last_camera.projection(last_vp_size.aspectRatio());
    auto ndc = glm::vec3(s_pos.x / w, s_pos.y / h, depth);
    auto c_pos_h = glm::inverse(proj) * to_homo(2.f * ndc - 1.f);
    auto c_pos = from_homo(c_pos_h);

    auto cam_mat = glm::rotate(last_camera.transform,
                               glm::radians(180.0f), glm::vec3(1.0f, 0.0f, 0.0f)); // camera -> viewport
    auto w_pos_h = cam_mat * to_homo(c_pos);
    auto w_pos = from_homo(w_pos_h);
    return w_pos;
}

void camera_augment_helper_v2::impl::pre_render_slot(const scene_ptr &info) {
    if (ui->enable_active_guide && allow_active_guide) {
        auto rec_info = scene_render_info::custom_info{
                .func = [this] {
                    if (manager->get_current_camera_helper() != q_this) {
                        return; // not current camera helper
                    }

                    auto s_pos = ImGui::GetMousePos();
                    auto info = cursor_guide::guide_info{
                            .position = get_cursor_coordinate(to_vec2(s_pos))
                    };
                    auto &io = ImGui::GetIO();
                    if (!io.WantCaptureMouse) {
                        info.mouse_left = io.MouseClicked[ImGuiMouseButton_Left];
                        info.mouse_right = io.MouseClicked[ImGuiMouseButton_Right];
                        info.mouse_wheel = io.MouseWheel;
                    }
                    act_guide->update(info);
                },
        };
        info->items.push_back({.info = rec_info});
        act_guide->pre_render_slot(info);
    }

    // TODO: not implemented
//    if (enable_passive_guide) {
//        auto rec_info = scene_render_info::custom_info{
//                .func = [=, this] {
//                    auto w_pos = get_cursor_coordinate(pas_info.pos);
//                    pas_guide.update({.position = w_pos});
//                },
//        };
//        info->items.push_back({.info = rec_info});
//        pas_guide.pre_render_slot(info);
//    }
}

camera_augment_helper_v2::camera_augment_helper_v2(const create_config &conf)
        : pimpl(std::make_unique<impl>(conf)) {
    pimpl->q_this = this;
}

camera_augment_helper_v2::~camera_augment_helper_v2() = default;

void camera_augment_helper_v2::render() {
    pimpl->render();
}

void camera_augment_helper_v2::render_image(output_config conf) {
    pimpl->render_image(conf);
}

void camera_augment_helper_v2::show() {
    pimpl->show();
}

void camera_augment_helper_v2::sync_with(camera_augment_helper_v2 *other) {
    pimpl->ui = other->pimpl->ui;
}

camera_info camera_augment_helper_v2::get_camera() {
    pimpl->update_camera();
    return pimpl->last_camera;
}

void camera_augment_helper_v2::set_fixed_transform_latency(int offset) {
    pimpl->time_offset = offset;
}

void camera_augment_helper_v2::set_active_guidance(bool flag) {
    pimpl->allow_active_guide = flag;
}

void camera_augment_helper_v2::set_passive_guidance(bool flag, const guide_info &info) {
    pimpl->enable_passive_guide = flag;
    pimpl->pas_info = info;
}