#include "TextureDisplacementBakeJob.hpp"

#include <algorithm>
#include <functional>
#include <utility>
#include <string>
#include <vector>
#include "libslic3r/TextureDisplacement.hpp"
#include <exception>
#include <cstddef>
#include <cstdint>
#include "slic3r/GUI/Jobs/Worker.hpp"
#include <memory>

#include "libslic3r/Model.hpp"
#include "libslic3r/format.hpp"
#include "libslic3r/TriangleSelector.hpp"

#include "slic3r/GUI/GLCanvas3D.hpp"
#include "slic3r/GUI/GUI.hpp"
#include "slic3r/GUI/GUI_App.hpp"
#include "slic3r/GUI/GUI_ObjectList.hpp"
#include "slic3r/GUI/Gizmos/GLGizmoTextureDisplacement.hpp"
#include "slic3r/GUI/I18N.hpp"
#include "slic3r/GUI/NotificationManager.hpp"
#include "slic3r/GUI/Plater.hpp"
#include "slic3r/Utils/UndoRedo.hpp"
#include "libslic3r/ObjectID.hpp"
#include "libslic3r/TriangleMesh.hpp"
#include <array>

namespace Slic3r::GUI {

TextureDisplacementBakeJob::TextureDisplacementBakeJob(TextureDisplacementBakeInput &&input, std::function<void()> on_finished)
    : m_input(std::move(input)), m_on_finished(std::move(on_finished))
{
}

void TextureDisplacementBakeJob::process(Ctl &ctl)
{
    const std::string status = _u8L("Baking texture displacement");
    ctl.update_status(1, status);

    // Only ever touches m_input (captured by value before this job was queued) and local state -
    // never the live Model - so this is safe to run concurrently with the UI thread.
    //
    // The progress hook matters for more than cosmetics: the framework's progress notification only
    // grows a close button once it reaches 100%, so a job that reports 0 and nothing else leaves an
    // uncloseable notification pinned on screen. It also carries the Cancel button's effect into the
    // bake, which on a subdivided mesh can run for several seconds.
    // Colour, when any layer asks for it, is computed in the same pass as the displacement: both need
    // the same per-layer projection and UV work, and doing it twice would double the expensive part.
    TextureColorRequest  color_request;
    TextureColorRequest *color = nullptr;
    if (!m_input.color.empty()) {
        color_request.quantize = GLGizmoTextureDisplacement::make_palette_quantizer(m_input.color.palette);
        if (!m_input.color.palette_pure.empty())
            color_request.quantize_pure = GLGizmoTextureDisplacement::make_palette_quantizer(m_input.color.palette_pure);
        color_request.despeckle_passes = m_input.color.despeckle_passes;
        color_request.out_triangle     = &m_triangle_color;
        if (color_request.quantize)
            color = &color_request;
    }

    int last_reported = 1;
    m_result = TriangleMesh(build_texture_displacement(
        m_input.base_mesh, m_input.layers, m_input.facets_data, m_input.options,
        [&ctl, &status, &last_reported](int percent) {
            if (ctl.was_canceled())
                return false;
            // The notification repaints (and wakes the idle loop) on every call, so only push a
            // message when the displayed integer percentage actually moves.
            if (percent > last_reported) {
                last_reported = percent;
                ctl.update_status(percent, status);
            }
            return true;
        },
        color, m_input.volume_to_world, nullptr, &m_stats));

    // Always finish at 100: this is what closes the notification. Reported even on cancel, where
    // build_texture_displacement() returns an empty mesh and finalize() commits nothing.
    ctl.update_status(100, status);
}

void TextureDisplacementBakeJob::finalize(bool canceled, std::exception_ptr &eptr)
{
    struct OnExit
    {
        std::function<void()> fn;
        ~OnExit() { if (fn) fn(); }
    } on_exit{m_on_finished};

    if (canceled || eptr || m_result.empty())
        return;

    // A bake that moved nothing - no layer could be sampled, or every sample was zero - must not be
    // committed: committing is what clears the baked layers' paint, so the user would see the painted
    // region simply vanish with no relief in its place and no idea why. Keep the paint and say so.
    {
        const indexed_triangle_set &out = m_result.its;
        bool unchanged = out.indices.size() == m_input.base_mesh.indices.size() &&
                         out.vertices.size() == m_input.base_mesh.vertices.size();
        for (size_t i = 0; unchanged && i < out.vertices.size(); ++i)
            unchanged = (out.vertices[i] - m_input.base_mesh.vertices[i]).cwiseAbs().maxCoeff() < 1e-5f;
        if (unchanged) {
            show_error(nullptr, _u8L("The bake produced no displacement, so nothing was changed and the paint was kept. "
                                     "Check that the painted layer has a texture and a non-zero depth."));
            return;
        }
    }

    Plater *plater = wxGetApp().plater();

    const auto commit = [this, plater]() {
        ModelVolume *volume = get_model_volume(m_input.volume_id, plater->model().objects);
        if (volume == nullptr)
            return;

        // The filament each palette entry prints in. This is where a mix becomes a mixed filament slot,
        // and only a mix the bake actually painted with: the preview never creates one, so the project
        // gains only the slots this result needs. Done before anything below names them - adding a slot
        // runs ModelVolume::update_extruder_count(), which clamps paint above the old filament count.
        Sidebar &sidebar = plater->sidebar();
        const std::vector<int> filament_of = GLGizmoTextureDisplacement::palette_filaments(
            m_input.color.palette, m_triangle_color, [&sidebar](const PrintableColor &mix) {
                return sidebar.ensure_mixed_filament({ unsigned(mix.a + 1), unsigned(mix.b + 1) },
                                                     { mix.a_percent(), 100 - mix.a_percent() });
            });

        // Whether the bake replaced the triangles or only moved the vertices it was given. The classic
        // path is topology-preserving by contract, so every per-triangle annotation on the volume still
        // means what it did; the default pipeline refines, cuts and simplifies, so none of them do.
        // Both are checked: the flag is what the bake promised, the count is what it delivered. Only
        // the triangles matter - the annotations are per facet and reference no vertex index - so a
        // vertex count that changed on its own (the bake compacts unreferenced vertices away) is not a
        // reason to remap anything.
        const bool topology_changed = m_input.options.pipeline_v2 ||
                                      m_result.its.indices.size() != m_input.base_mesh.indices.size();

        // Carry the volume's existing paint - the model's own colours above all - across that rebuild.
        // Without it a bake left the mmu_segmentation_facets bitstream pointing at a triangle set that
        // no longer existed, which is what scrambled the colours of an imported coloured model even
        // with colour switched off on every layer and nothing in the bake that was meant to touch them.
        //
        // transfer_facet_paint() rather than ModelVolume::restore_painting(): the latter's remap wants
        // a target triangle overlapping the source paint to within 0.01 mm, which the baked surface is
        // nowhere near - it has moved by the relief depth. Nothing to do when the topology is intact:
        // triangle i is still triangle i, so the paint is already exactly right and carrying it at all
        // could only lose detail.
        std::array<TriangleSelector::TriangleSplittingData, 4> before;
        if (topology_changed) {
            before[0] = volume->supported_facets.get_data();
            before[1] = volume->seam_facets.get_data();
            before[2] = volume->mmu_segmentation_facets.get_data();
            before[3] = volume->fuzzy_skin_facets.get_data();
        }

        volume->set_mesh(std::move(m_result));
        volume->set_new_unique_id();
        volume->calculate_convex_hull();

        if (topology_changed) {
            const std::array<FacetsAnnotation *, 4> target{ &volume->supported_facets, &volume->seam_facets,
                                                            &volume->mmu_segmentation_facets,
                                                            &volume->fuzzy_skin_facets };
            for (size_t ch = 0; ch < target.size(); ++ch) {
                const std::vector<uint8_t> state =
                    transfer_facet_paint(m_input.base_mesh, before[ch], volume->mesh().its);
                // Nothing was painted in this channel, or there is nothing to paint it onto: either
                // way the annotation must not keep describing the mesh that is gone.
                if (state.empty()) {
                    target[ch]->reset();
                    continue;
                }
                TriangleSelector selector(volume->mesh());
                bool             any = false;
                for (size_t i = 0; i < state.size(); ++i)
                    if (state[i] != 0) {
                        selector.set_facet(int(i), EnforcerBlockerType(state[i]));
                        any = true;
                    }
                if (any)
                    target[ch]->set(selector);
                else
                    target[ch]->reset();
            }
        }

        // Colour lands in mmu_segmentation_facets, merged *over* whatever is already painted there
        // rather than replacing it: a triangle the texture does not colour keeps its existing filament,
        // and one the user never painted at all stays at NONE, which already means "the volume's own
        // filament". That is what confines the effect to the painted area without having to invent a
        // colour for everything outside it. What was already painted is whatever the transfer above
        // left in the annotation, so this bake's colour lands on top of it rather than instead of it.
        if (!m_triangle_color.empty() && m_triangle_color.size() == volume->mesh().its.indices.size()) {
            TriangleSelector selector(volume->mesh());
            const TriangleSelector::TriangleSplittingData &existing = volume->mmu_segmentation_facets.get_data();
            if (!existing.bitstream.empty())
                selector.deserialize(existing, false);
            for (size_t i = 0; i < m_triangle_color.size(); ++i) {
                const size_t entry = size_t(m_triangle_color[i]);
                if (entry > 0 && entry <= filament_of.size() && filament_of[entry - 1] >= 0)
                    selector.set_facet(int(i), EnforcerBlockerType(filament_of[entry - 1] + 1));
            }
            volume->mmu_segmentation_facets.set(selector);
        }

        // Clear the paint mask of every layer that was actually baked so a repeat bake (or the paint
        // overlay) doesn't act on triangles that no longer represent the same unbaked surface. The
        // texture layer definitions themselves are left untouched so the user can keep sculpting with
        // the same textures.
        //
        // A mask the bake did *not* consume only still means what it did if the topology is unchanged,
        // which is true of the classic path (it moves existing vertices) but not of the one-run
        // pipeline, which rebuilds and then simplifies the mesh. A mask left behind against the old
        // topology is exactly what makes the gizmo reload an empty selector over a non-empty mask and
        // then erase it on the next flush - see GLGizmoTextureDisplacement::update_from_model_object().
        for (int slot = 0; slot < int(TEXTURE_DISPLACEMENT_MAX_LAYERS); ++slot) {
            const auto it = std::find_if(m_input.layers.begin(), m_input.layers.end(),
                                         [slot](const TextureDisplacementLayer &l) { return l.slot == slot; });
            if (topology_changed || (it != m_input.layers.end() && !it->empty()))
                volume->texture_displacement_facet(slot).reset();
        }

        ModelObject *object = volume->get_object();
        if (object == nullptr)
            return;

        if (ObjectList *obj_list = wxGetApp().obj_list()) {
            const ModelObjectPtrs &objs = plater->model().objects;
            auto it = std::find(objs.begin(), objs.end(), object);
            if (it != objs.end())
                obj_list->update_info_items(size_t(it - objs.begin()));
        }

        plater->changed_object(*object);
    };

    // Standard mode's Bake is remesh -> subdivide -> displace under a single snapshot, and this job runs
    // long after that snapshot's scope has closed. Adding one here would put an undo step *between* the
    // subdivision and the displacement: the first Undo would land on a mesh carrying every added
    // triangle and no relief at all, and pressing Bake again from there would subdivide that mesh a
    // second time. So the caller says who owns the undo step.
    if (m_input.take_snapshot) {
        Plater::TakeSnapshot snapshot(plater, _u8L("Bake texture displacement"), UndoRedo::SnapshotType::GizmoAction);
        commit();
    } else {
        commit();
    }

    // The refinement went finer than the budget could keep: the simplification had to take detail back
    // out to fit, so what was baked carries less of the texture than the resolution asked for. Said
    // here, with the numbers, because it is the only place that knows them - and not as an error
    // dialog: the result is a usable mesh, just not the one the settings described.
    if (m_stats.budget_limited) {
        // Thousands under a million, so a small budget is not reported as "0.1 M".
        const auto count = [](size_t n) {
            return n >= 1000000 ? Slic3r::format("%1$.1f M", double(n) / 1000000.) :
                                  Slic3r::format("%1% k", (n + 500) / 1000);
        };
        wxGetApp().notification_manager()->push_notification(
            NotificationType::CustomNotification, NotificationManager::NotificationLevel::WarningNotificationLevel,
            Slic3r::format(_u8L("The triangle budget limited the detail: this resolution needs %1% triangles, "
                                "the budget kept %2%. Raise Budget or use a coarser Resolution for the full detail."),
                           count(m_stats.triangles_refined), count(m_stats.triangles_budget)));
    }
}

void queue_texture_displacement_bake(const ModelVolume &volume, const TextureColorSettings &color,
                                     std::function<void()> on_finished, bool take_snapshot)
{
    TextureDisplacementBakeInput input;
    input.color         = color;
    input.take_snapshot = take_snapshot;
    input.volume_id  = volume.id();
    input.base_mesh  = volume.mesh().its;
    input.layers     = volume.texture_displacement_layers;
    input.options    = volume.texture_displacement_options;
    input.volume_to_world = texture_displacement_volume_to_world(volume);
    for (int i = 0; i < int(TEXTURE_DISPLACEMENT_MAX_LAYERS); ++i)
        input.facets_data[size_t(i)] = volume.texture_displacement_facet(i).get_data();


    auto &worker = wxGetApp().plater()->get_ui_job_worker();
    queue_job(worker, std::make_unique<TextureDisplacementBakeJob>(std::move(input), std::move(on_finished)));
}

} // namespace Slic3r::GUI