#include "vtk_viewer.h"
#include "core/utility.hpp"

#include <vtkCamera.h>
#include <vtkCellArray.h>
#include <vtkGenericOpenGLRenderWindow.h>
#include <vtkGenericRenderWindowInteractor.h>
#include <vtkInteractorStyleTrackballCamera.h>
#include <vtkMatrix4x4.h>
#include <vtkNew.h>
#include <vtkOBJReader.h>
#include <vtkOpenGLFramebufferObject.h>
#include <vtkPointPicker.h>
#include <vtkPoints.h>
#include <vtkPolyDataMapper.h>
#include <vtkProperty.h>
#include <vtkRenderer.h>
#include <vtkSTLReader.h>
#include <vtkTextureObject.h>

#include <spdlog/spdlog.h>

#include <filesystem>
#include <vector>

namespace vtk_viewer_impl {

    Eigen::Isometry3d to_eigen_transform(vtkMatrix4x4 *mat) {
        Eigen::Isometry3d ret;
        for (int i = 0; i < 4; ++i)
            for (int j = 0; j < 4; ++j)
                ret.matrix()(i, j) = mat->GetElement(i, j);
        return ret;
    }

}

using namespace vtk_viewer_impl;

struct vtk_viewer::impl {

    struct versatile_interaction_style;

    static constexpr auto default_focal_length = 8; // 8mm

    cv::Size last_size;
    vtkSmartPointer<vtkGenericOpenGLRenderWindow> window;
    vtkSmartPointer<vtkCamera> camera;
    vtkSmartPointer<vtkRenderWindowInteractor> controller;
    vtkSmartPointer<versatile_interaction_style> style;
    vtkSmartPointer<vtkRenderer> renderer;

    bool is_picking = false;
    std::optional<Eigen::Vector3d> picked_point;

    struct versatile_interaction_style
            : public vtkInteractorStyleTrackballCamera {
    vtkTypeMacro(versatile_interaction_style, vtkInteractorStyleTrackballCamera);

        static versatile_interaction_style *New() {
            auto ret = new versatile_interaction_style{};
            ret->InitializeObjectBase();
            return ret;
        }

        impl *pimpl = nullptr;
        Eigen::Vector2i mouse_down_pose;

        Eigen::Vector2i get_mouse_pos() {
            return {Interactor->GetEventPosition()[0],
                    Interactor->GetEventPosition()[1],
            };
        }

        void OnLeftButtonDown() override {
            mouse_down_pose = get_mouse_pos();
            vtkInteractorStyleTrackballCamera::OnLeftButtonDown();
        }

        bool try_picking() {
            assert(pimpl != nullptr);
            if (!pimpl->is_picking) return false;
            auto ret = Interactor->GetPicker()->Pick(mouse_down_pose.x(),
                                                     mouse_down_pose.y(),
                                                     0,
                                                     pimpl->renderer);
            if (ret == 0) return true; // does not pick anything
            Eigen::Vector3d point;
            Interactor->GetPicker()->GetPickPosition(point.data());
            pimpl->picked_point = point;
            SPDLOG_INFO("Picked point ({}, {}, {}).", point.x(), point.y(), point.z());
            return true;
        }

        void OnLeftButtonUp() override {
            if (get_mouse_pos() == mouse_down_pose) {
                try_picking();
            }
            vtkInteractorStyleTrackballCamera::OnLeftButtonUp();
        }
    };

    static cv::Size to_cv_size(ImVec2 size) {
        return {(int) size.x, (int) size.y};
    }

    impl() {
        window = vtkSmartPointer<vtkGenericOpenGLRenderWindow>::New();
        window->InitializeFromCurrentContext();
        window->SetOffScreenRendering(true);
        window->SetAlphaBitPlanes(true);
//        window->FramebufferFlipYOn();
        window->SetIsCurrent(true);
//        window->SwapBuffersOn();
//        window->SetFrameBlitModeToNoBlit();

        camera = vtkSmartPointer<vtkCamera>::New();
        camera->SetClippingRange(default_focal_length, 2000); // 8mm to 2m

        renderer = vtkSmartPointer<vtkRenderer>::New();
        renderer->SetUseDepthPeeling(true);
        renderer->SetActiveCamera(camera);
        window->AddRenderer(renderer);

        style = vtkSmartPointer<versatile_interaction_style>::New();
        style->pimpl = this;
        style->SetDefaultRenderer(renderer);
        controller = vtkSmartPointer<vtkGenericRenderWindowInteractor>::New();
        controller->SetInteractorStyle(style);
        controller->EnableRenderOff();
        window->SetInteractor(controller);
    }

    void set_camera_pose(const Eigen::Isometry3d &trans) {
        auto trans_part = trans.translation();
        auto rot_part = trans.rotation();
        auto focal_point = trans_part + rot_part.col(2) * default_focal_length;
        auto view_up = -rot_part.col(1);
        camera->SetPosition(trans_part.x(), trans_part.y(), trans_part.z());
        camera->SetFocalPoint(focal_point.x(), focal_point.y(), focal_point.z());
        camera->SetViewUp(view_up.x(), view_up.y(), view_up.z());
        camera->Modified();
    }

    void render(cv::Size size, bool interactive) {
        assert(size.area() > 0);
        if (size != last_size) [[unlikely]] {
            controller->SetSize(size.width, size.height);
            window->SetSize(size.width, size.height);
            last_size = size;
        }
        if (interactive) {
            process_event();
        }
        window->Render();
    }

    void process_event() {
        if (!ImGui::IsWindowFocused() || !ImGui::IsWindowHovered()) return;

        // set event position
        auto &io = ImGui::GetIO();
        io.ConfigWindowsMoveFromTitleBarOnly = true;
        auto view_pos = ImGui::GetWindowPos();
        auto x_pos = io.MousePos.x - view_pos.x;
        auto y_pos = io.MousePos.y - view_pos.y;

        // flip X axis
        x_pos = controller->GetSize()[0] - x_pos;
        controller->SetEventInformation(x_pos, y_pos, io.KeyCtrl, io.KeyShift);

        // dispatch event
        if (ImGui::IsWindowHovered()) {
            if (io.MouseClicked[ImGuiMouseButton_Left]) {
                controller->InvokeEvent(vtkCommand::LeftButtonPressEvent);
            }
            if (io.MouseClicked[ImGuiMouseButton_Middle]) {
                controller->InvokeEvent(vtkCommand::MiddleButtonPressEvent);
            }
            if (io.MouseClicked[ImGuiMouseButton_Right]) {
                controller->InvokeEvent(vtkCommand::RightButtonPressEvent);
            }
            if (io.MouseWheel > 0) {
                controller->InvokeEvent(vtkCommand::MouseWheelForwardEvent);
            } else if (io.MouseWheel < 0) {
                controller->InvokeEvent(vtkCommand::MouseWheelBackwardEvent);
            }
        }
        if (io.MouseReleased[ImGuiMouseButton_Left]) {
            controller->InvokeEvent(vtkCommand::LeftButtonReleaseEvent);
        }
        if (io.MouseReleased[ImGuiMouseButton_Middle]) {
            controller->InvokeEvent(vtkCommand::MiddleButtonReleaseEvent);
        }
        if (io.MouseReleased[ImGuiMouseButton_Right]) {
            controller->InvokeEvent(vtkCommand::RightButtonReleaseEvent);
        }
        controller->InvokeEvent(vtkCommand::MouseMoveEvent);
    }

    GLuint get_tex() const {
        return window->GetDisplayFramebuffer()
                ->GetColorAttachmentAsTextureObject(0)->GetHandle();
    }

    void show_imgui_window(const char *name, ImVec2 req_size) {
        ImGui::BeginChild(name, req_size, 0, no_scroll_flag);
        auto render_size = ImGui::GetContentRegionAvail();
        auto render_size_cv = to_cv_size(render_size);
        if (render_size_cv.area() <= 0)return;
        render(render_size_cv, true);
        ImGui::Image(reinterpret_cast<void *>(get_tex()), render_size, {1, 0}, {0, 1});
        ImGui::EndChild();
    }

};

vtk_viewer::vtk_viewer()
        : pimpl(std::make_unique<impl>()) {}

vtk_viewer::~vtk_viewer() = default;

void vtk_viewer::render(cv::Size size) {
    return pimpl->render(size, false);
}

GLuint vtk_viewer::get_tex() const {
    return pimpl->get_tex();
}

void vtk_viewer::add_actor(vtkActor *actor) {
    pimpl->renderer->AddActor(actor);
}

void vtk_viewer::remove_actor(vtkActor *actor) {
    pimpl->renderer->RemoveActor(actor);
}

void vtk_viewer::clear_actor() {
    pimpl->renderer->RemoveAllViewProps();
}

void vtk_viewer::show(const std::string &name) {
    pimpl->show_imgui_window(name.c_str(), ImGui::GetContentRegionAvail());
}

Eigen::Isometry3d vtk_viewer::get_camera_pose() {
    auto ret = to_eigen_transform( // world in camera, I think
            pimpl->camera->GetModelViewTransformMatrix());
    ret = ret.inverse() * Eigen::AngleAxisd{std::numbers::pi, Eigen::Vector3d::UnitX()};
    return ret;
}

void vtk_viewer::set_camera_pose(const Eigen::Isometry3d &trans) {
    pimpl->set_camera_pose(trans);
}

void vtk_viewer::set_camera_view_angle(double angle) {
    pimpl->camera->SetViewAngle(angle);
    pimpl->camera->Modified();
}

void vtk_viewer::start_picking() {
    pimpl->is_picking = true;
}

void vtk_viewer::stop_picking() {
    pimpl->is_picking = false;
}

bool vtk_viewer::is_picking() {
    return pimpl->is_picking;
}

std::optional<Eigen::Vector3d> vtk_viewer::get_picked_point() {
    auto ret = pimpl->picked_point;
    pimpl->picked_point = {};
    return ret;
}

void vtk_viewer::reset_camera() {
    pimpl->renderer->ResetCamera();
}

namespace vtk_viewer_helper {

    vtkSmartPointer<vtkPolyData> load_stl(const std::string &path) {
        vtkNew<vtkSTLReader> reader;
        reader->SetFileName(path.c_str());
        reader->Update();
        return reader->GetOutput();
    }

    vtkSmartPointer<vtkPolyData> load_obj(const std::string &path) {
        vtkNew<vtkOBJReader> reader;
        reader->SetFileName(path.c_str());
        reader->Update();
        return reader->GetOutput();
    }

    vtkSmartPointer<vtkPolyData> load_any(const std::string &path) {
        auto path_fs = std::filesystem::path{path};
        auto ext = path_fs.extension();
        if (ext == ".stl") {
            return load_stl(path);
        } else if (ext == ".obj") {
            return load_obj(path);
        }
        RET_ERROR_P;
    }

    vtkSmartPointer<vtkActor> create_actor(vtkPolyData *data) {
        vtkNew<vtkPolyDataMapper> mapper;
        mapper->SetInputData(data);
        vtkNew<vtkActor> actor;
        actor->SetMapper(mapper);
        vtkNew<vtkMatrix4x4> pose;
        actor->SetUserMatrix(pose);
        return actor;
    }

    vtkSmartPointer<vtkActor> create_actor(const std::string &path) {
        return create_actor(load_stl(path));
    }

    void update_actor_pose(vtkActor *actor, const std::optional<Eigen::Isometry3d> &trans) {
        if (!trans.has_value()) {
            actor->VisibilityOff();
            return;
        }
        actor->VisibilityOn();
        auto &real_trans = trans.value();
        auto matrix = actor->GetUserMatrix();
        if (matrix == nullptr) {
            actor->SetUserMatrix(vtkMatrix4x4::New());
            matrix = actor->GetUserMatrix();
        }
        for (int i = 0; i < 4; ++i) {
            for (int j = 0; j < 4; ++j) {
                matrix->SetElement(i, j, real_trans(i, j));
            }
        }
        matrix->Modified();
        actor->Modified();
    }

}

struct smart_point_sets::impl {

    vtkSmartPointer<vtkPoints> points;
    vtkSmartPointer<vtkCellArray> vertices;
    vtkSmartPointer<vtkPolyData> poly;
    vtkSmartPointer<vtkActor> actor;

    std::vector<Eigen::Vector3d> points_store;

    impl() {
        actor = vtkSmartPointer<vtkActor>::New();
        actor->GetProperty()->SetRenderPointsAsSpheres(true);
        actor->GetProperty()->SetPointSize(10);
        reconstruct();
    }

    void add_point_to_poly(const Eigen::Vector3d &point) {
        vtkIdType pid[1];
        pid[0] = points->InsertNextPoint(point.data());
        points->Modified();
        vertices->InsertNextCell(1, pid);
        vertices->Modified();
        poly->Modified();
    }

    auto add_point(const Eigen::Vector3d &point) {
        points_store.emplace_back(point);
        add_point_to_poly(point);
    }

    void reconstruct() {
        points = vtkSmartPointer<vtkPoints>::New();
        vertices = vtkSmartPointer<vtkCellArray>::New();
        poly = vtkSmartPointer<vtkPolyData>::New();
        poly->SetPoints(points);
        poly->SetVerts(vertices);
        vtkNew<vtkPolyDataMapper> mapper;
        mapper->SetInputData(poly);
        actor->SetMapper(mapper);
        actor->Modified();

        for (auto &p: points_store) {
            add_point_to_poly(p);
        }
    }

    void for_each(const for_each_func_type &func) {
        for (auto iter = points_store.begin();
             iter != points_store.end();
             ++iter) {
            func(*(void **) &iter, *iter);
        }
    }

    void remove_point(void *token) {
        using iter_type = decltype(points_store)::iterator;
        points_store.erase(*(iter_type *) (&token));
        reconstruct();
    }

    Eigen::Vector3d pop_front() {
        assert(!points_store.empty());
        auto ret = points_store.front();
        points_store.erase(points_store.begin());
        reconstruct();
        return ret;
    }
};

smart_point_sets::smart_point_sets()
        : pimpl(std::make_unique<impl>()) {}

smart_point_sets::~smart_point_sets() = default;

void smart_point_sets::add_point(const Eigen::Vector3d &point) {
    pimpl->add_point(point);
}

vtkActor *smart_point_sets::get_actor() {
    return pimpl->actor;
}

void smart_point_sets::for_each(const for_each_func_type &func) {
    pimpl->for_each(func);
}

void smart_point_sets::remove_point(void *token) {
    pimpl->remove_point(token);
}

Eigen::Vector3d smart_point_sets::pop_front() {
    return pimpl->pop_front();
}

bool smart_point_sets::empty() const {
    return pimpl->points_store.empty();
}

size_t smart_point_sets::size() const {
    return pimpl->points_store.size();
}

std::vector<Eigen::Vector3d> smart_point_sets::to_vector() const {
    return pimpl->points_store;
}

Eigen::Vector3d smart_point_sets::operator[](size_t index) const {
    return pimpl->points_store[index];
}

void smart_point_sets::clear() const {
    pimpl->points_store.clear();
    pimpl->reconstruct();
}