#include "FlatDualMC.hpp"

#include "core/world/Elements.hpp"
#include <Tracy.hpp>   // NOLINT
#include <common/TracySystem.hpp> // NOLINT
#include <imgui.h>     // NOLINT
#include "core/utils/toml.hpp"
#include "dualmc.h"
#include "optimizer.hpp"

namespace contouring {
    constexpr auto PACK = [](const render::VertexData &v) {
        return render::PackedVertexData(meshopt_quantizeHalf(v.Position.x), meshopt_quantizeHalf(v.Position.y),
            meshopt_quantizeHalf(v.Position.z), v.Material,
            meshopt_quantizeHalf(v.Normal.x), meshopt_quantizeHalf(v.Normal.y),
            meshopt_quantizeHalf(v.Normal.z));
    };
    const std::vector<std::pair<float, float>> LEVELS = {{.001f, 9e-1f}, {.01f, 5e-1f}, {.1f, 1e-1f}, {.2, 1e-2f}, {.5, 5e-3f}};

    FlatDualMC::FlatDualMC(const std::string &str) : Abstract() {
        auto opt = toml::parse(str);
        loadDistance = opt["load_distance"].value_or(loadDistance);
        keepDistance = opt["keep_distance"].value_or(keepDistance);
        transparency = opt["transparency"].value_or(transparency);
        iso = opt["iso"].value_or(iso);
        manifold = opt["manifold"].value_or(manifold);
        reordering = opt["reordering"].value_or(reordering);
        lod_quality = opt["lod_quality"].value_or(lod_quality);
        lod_strength = opt["lod_strength"].value_or(lod_strength);
        if (const auto ptr = opt["lod_levels"].as_array(); ptr != nullptr && ptr->size() == LEVELS.size()) {
            const auto& arr = *ptr;
            for(const auto& v: arr) {
                lod_levels.push_back(v.value_or(true));
            }
        } else {
            lod_levels = {false, true, false, true, false};
        }
        for (size_t i = 0; i < LEVELS.size(); i++) {
            if(lod_levels[i])
                loadedLevels.push_back(LEVELS[i]);
        }

        for (size_t i = 1; i <= std::max<size_t>(1, std::thread::hardware_concurrency() / 2 - 1); i++) {
            workers.emplace_back([&] {
#if TRACY_ENABLE
                tracy::SetThreadName("Contouring");
#endif
                std::vector<render::VertexData> tmp;
                while (running) {
                    const auto area_priority = areaLoadQueue.currentWeight();
                    const auto part_priority = partLoadQueue.currentWeight();

                    if (part_priority > area_priority) {
                        if (std::pair<world::part_id, part_job> job; partLoadQueue.pop(job)) {
                            ZoneScopedN("Part");
                            render::Model::Data data;
                            render::Model::Data::indices_t idx;
                            for (uint8_t x = 0; x < job.second.size.x; x++) {
                            for (uint8_t y = 0; y < job.second.size.y; y++) {
                            for (uint8_t z = 0; z < job.second.size.z; z++) {
                                surrounding::corners sur;
                                surrounding::load(sur, job.second.size, glm::ucvec3(x, y, z), job.second.chunks);
                                idx.clear();
                                render(sur, idx, tmp, Layer::Both); //MAYBE: write inplace with sub-vector
                                simplify(idx, tmp);

                                data.indices.reserve(data.indices.size() + idx.size());
                                const auto idx_start = data.vertices.size();
                                std::transform(idx.begin(), idx.end(), std::back_inserter(data.indices), [&](glm::u16 i) { return i + idx_start; });
                                data.vertices.reserve(data.vertices.size() + tmp.size());
                                const auto vert_offset = glm::vec3(x * world::CHUNK_LENGTH, y * world::CHUNK_LENGTH, z * world::CHUNK_LENGTH);
                                std::transform(tmp.begin(), tmp.end(), std::back_inserter(data.vertices), [&](render::VertexData v) {
                                    v.Position += vert_offset; return PACK(v); });
                            }}}
                            optimize_fetch(data);
                            partLoadedQueue.emplace(job.first, std::move(data));
                        }
                    } else if (area_priority.has_value() && (!part_priority.has_value() || area_priority.value() >= part_priority.value())) {
                        if (std::pair<world::area_chunk_pos, surrounding::corners> ctx; areaLoadQueue.pop(ctx)) {
                            ZoneScopedN("Chunk");
                            area_models::data data;
                            if (transparency) {
                                render(ctx.second, data.main, tmp, Layer::Solid);
                                render(ctx.second, data.transparent, tmp, Layer::Transparent);
                            } else {
                                render(ctx.second, data.main, tmp, Layer::Both);
                            }
                            //MAYBE: direct upload with vulkan
                            areaLoadedQueue.emplace(ctx.first, std::move(data));
                        }
                    } else {
                        areaLoadQueue.wait();
                    }
                }
            });
        }
    }
    FlatDualMC::~FlatDualMC() {
        running = false;
        areaLoadQueue.notify_all();

        for(auto& worker: workers) {
            if (worker.joinable())
                worker.join();
        }
    }

    std::string FlatDualMC::getOptions() const {
        std::ostringstream ss;
        toml::table tb({
            {"load_distance", loadDistance},
            {"keep_distance", keepDistance},
            {"transparency", transparency},
            {"iso", iso},
            {"manifold", manifold},
            {"reordering", reordering},
            {"lod_quality", lod_quality},
            {"lod_strength", lod_strength},
            {"lod_levels", toml::array()}
        });
        for(auto& v: lod_levels)
            tb["lod_levels"].as_array()->push_back(v);
        ss << tb;
        return ss.str();
    }
    std::pair<float, float> FlatDualMC::getFarRange() const {
        return std::make_pair((loadDistance - 1.5) * world::CHUNK_LENGTH, (keepDistance + .5) * world::CHUNK_LENGTH);
    }

    size_t FlatDualMC::getQueueSize() {
        return areaLoadQueue.size() + partLoadQueue.size();
    }

    void FlatDualMC::enqueue(const world::area_chunk_pos &pos, glm::ll offset, const world::ChunkContainer &data) {
        ZoneScopedN("EnqueueContouring");
        if (offset <= loadDistance * loadDistance) {
            surrounding::corners surrounding;
            if(surrounding::load(surrounding, pos.chunk, data)) {
                areaLoadQueue.push(pos, std::move(surrounding), -offset);
            }
        }
    }
    void FlatDualMC::onUpdate(const world::area_chunk_pos &pos, glm::ll offset, const world::ChunkContainer &data, geometry::Faces neighbors) {
        enqueue(pos, offset, data);
        if (neighbors && (geometry::Faces::Left | geometry::Faces::Down | geometry::Faces::Backward)) {
            for (size_t i = 1; i < 8; i++) {
                enqueue(world::area_chunk_pos{pos.area, pos.chunk - surrounding::g_corner_offsets[i]}, offset, data);
            }
        }
    }
    void FlatDualMC::onNotify(const world::area_chunk_pos &pos, glm::ll offset, const world::ChunkContainer &data) {
        const auto area = areas.find(pos.area);
        if(area == areas.end() || !area->second.second.contains(pos.chunk)) {
            enqueue(pos, offset, data);
        }
    }

    void FlatDualMC::onNotify(const world::part_id& id, glm::ll offset, const glm::ucvec3& size, const std::vector<std::shared_ptr<world::EdittableChunk>>& chunks, bool update) {
        if (parts.contains(id)) {
            if (!update)
                return;
        } else {
            parts.set_emplace(id, nullptr);
        }
        if (offset <= loadDistance * loadDistance) {
            partLoadQueue.push(id, part_job{size, chunks}, -offset);
        }
    }

    void FlatDualMC::onLoad(const world::model_id& id, std::istream& in) {
        if (!models.contains(id)) {
            models.set_emplace(id, nullptr);
            modelLoadedQueue.emplace(id, render::Model::Data(in));
        }
    }
    void FlatDualMC::onUnload(const world::model_id& id) {
        models.vanish(id);
    }

    render::Model* FlatDualMC::getModel(const world::part_id& id) const {
        if(const auto it = parts.directly_at(id))
            return it->get();

        return nullptr;
    }
    render::Model* FlatDualMC::getModel(const world::model_id &id) const {
        if(const auto it = models.directly_at(id))
            return it->get();

        return nullptr;
    }

    void FlatDualMC::update(const world::voxel_pos& pos, const world::Elements& l) {
        ZoneScopedN("Ct");
        TracyPlot("CtLoad", static_cast<int64_t>(getQueueSize()));
        TracyPlot("CtLoaded", static_cast<int64_t>(areaLoadedQueue.size() + partLoadedQueue.size() + modelLoadedQueue.size()));

        size_t buffer_count = 0;
        { // Models
            modelLoadedQueue.extractor([&](const std::pair<world::model_id, render::Model::Data>& out) {
                if (const auto entity = models.directly_at(out.first))
                    *entity = render::Model::Create(out.second);
            });
            models.extract([&](const world::model_id &id, const std::unique_ptr<render::Model> &) {
                return !l.models.contains(id); });
            buffer_count += models.size();
        }
        { // Parts
            partLoadedQueue.extractor([&](const std::pair<world::part_id, render::Model::Data>& out) {
                if (const auto entity = parts.directly_at(out.first))
                    *entity = render::Model::Create(out.second);
            });
            parts.extract([&](const world::part_id &id, const std::unique_ptr<render::Model> &) {
                return !l.nodes.contains(id.val); });
            //MAYBE: or too far
            buffer_count += parts.size();
        }
        { // Areas
            areaLoadedQueue.extractor([&](const std::pair<world::area_chunk_pos, area_models::data> &out) {
                area_models mds;
                if (!out.second.main.first.empty())
                    mds.main = render::LodModel::Create(out.second.main);
                if (!out.second.transparent.first.empty())
                    mds.transparent = render::LodModel::Create(out.second.transparent);
                auto &chunks = areas[out.first.area].second; //NOTE: areas[n].first uninitialized: will be set by second part
                chunks.erase(out.first.chunk);
                chunks.emplace(out.first.chunk, std::move(mds));
            });
            ZoneScopedN("CtUpdate");
            const auto levelMax = loadedLevels.size();
            for (auto it_a = areas.begin(); it_a != areas.end();) {
                if (const auto node = world::NodeOf<world::Area>::Make(l.nodes.directly_at(it_a->first.val))) {
                    auto &infos = it_a->second.first;
                    infos.absolute = node->absolute;
                    const auto ptr = node->get();
                    infos.radius = ptr->getChunks().getRadius() * (world::cell_pos::value_type)world::CHUNK_LENGTH;
                    const auto diff = pos - (world::cell_pos)infos.absolute.position;
                    if (const auto surface = ptr->getCurvature()) {
                        // MAYBE: move to area or generator
                        infos.curvature = std::clamp<float>((glm::max_axis(glm::abs(diff)) - surface.value()) / (infos.radius - surface.value()), -0.1f, 1.f);
                    }
                    const auto chunkDiff = glm::dvec3(glm::divide(diff));
                    auto &chunks = it_a->second.second;
                    for (auto it_c = chunks.begin(); it_c != chunks.end();) {
                        const auto distRatio = glm::length(chunkDiff - (glm::qua<double>)infos.absolute.rotation * glm::dvec3(it_c->first)) / keepDistance;
                        if (distRatio <= 1) {
                            const auto level = std::clamp<size_t>((1 + lod_quality - distRatio) * levelMax * (1 + lod_strength), 0, levelMax);
                            if (it_c->second.main) {
                                it_c->second.main->setLevel(level);
                                buffer_count++;
                            }
                            if (it_c->second.transparent) {
                                it_c->second.transparent->setLevel(level);
                                buffer_count++;
                            }
                            ++it_c;
                        } else {
                            it_c = chunks.erase(it_c);
                        }
                    }
                }
                if (it_a->second.second.empty()) {
                    it_a = areas.erase(it_a);
                } else {
                    ++it_a;
                }
            }
        }
        TracyPlot("CtCount", static_cast<int64_t>(buffer_count));
    }

    void FlatDualMC::onGui() {
        {
            int load = loadDistance;
            ImGui::SliderInt("Load Distance", &load, 1, keepDistance);
            loadDistance = load;
        }
        {
            int keep = keepDistance;
            ImGui::SliderInt("Keep Distance", &keep, loadDistance + 1, 21);
            keepDistance = keep;
        }
        ImGui::Checkbox("Transparency", &transparency);
        ImGui::SliderFloat("Iso", &iso, 0, 1);
        ImGui::Checkbox("Manifold", &manifold);
        ImGui::Checkbox("Reordering", &reordering);
        ImGui::SliderFloat("Lod quality", &lod_quality, -.5, 1);
        ImGui::SliderFloat("Lod strength", &lod_strength, -.5, .5);
        if (ImGui::CollapsingHeader("Lod levels")) {
            ImGui::Columns(3, nullptr, false);
            ImGui::Selectable("1", true, ImGuiSelectableFlags_Disabled);
            ImGui::NextColumn();
            for(int i = LEVELS.size() - 1; i >= 0; i--) {
                const auto str = std::to_string(static_cast<int>(1 / LEVELS[i].first));
                if (ImGui::Selectable(str.c_str(), &lod_levels[i])) {
                    loadedLevels.clear();
                    for (size_t i = 0; i < LEVELS.size(); i++) {
                        if(lod_levels[i])
                            loadedLevels.push_back(LEVELS[i]);
                    }
                }
                ImGui::NextColumn();
            }
            ImGui::Columns(1);
        }
        if (ImGui::Button("Flush buffers")) {
            areas.clear();
            parts.extract([](const world::part_id &, const std::unique_ptr<render::Model> &) { return true; });
            models.extract([](const world::model_id &, const std::unique_ptr<render::Model> &) { return true; });
        }
    }

    void FlatDualMC::render(const surrounding::corners &surrounding, render::Model::Data::indices_t &out, std::vector<render::VertexData> &tmp, Layer layer) const {
        constexpr uint8_t SIZE = world::CHUNK_LENGTH + 3;
        std::array<dualmc::DualMC<float>::Point, SIZE * SIZE * SIZE> grid;
        const auto &materials = world::module::Registry::Get()->getMaterials();
        {
            ZoneScopedN("Load");
            const auto setCell = [&](glm::idx i, const world::Voxel &voxel) {
                auto &cell = grid[i];
                cell.w = voxel.material();
                cell.x = voxel.density_ratio() * (!materials.invisibilities[cell.w] &&
                    ((materials.transparencies[cell.w] && (layer && Layer::Transparent)) ||
                    (!materials.transparencies[cell.w] && (layer && Layer::Solid))));
            };
            size_t i = 0;
            for (uint8_t z = 0; z < SIZE; z++) {
            for (uint8_t y = 0; y < SIZE; y++) {
            for (uint8_t x = 0; x < SIZE; x++) {
                const auto &chunk = surrounding[(z >= world::CHUNK_LENGTH) + (y >= world::CHUNK_LENGTH)*2 + (x >= world::CHUNK_LENGTH)*4];
                const auto &voxel = chunk->get(glm::toIdx(x & glm::IDX_MASK, y & glm::IDX_MASK, z & glm::IDX_MASK));
                setCell(i, voxel);
                i++;
            }}}

            for (size_t i = 0; i < surrounding.size(); i++) {
                auto &edits = surrounding[i]->getEdits().getEdits();
                auto offset = glm::ivec3(surrounding::g_corner_offsets[i]) * world::CHUNK_LENGTH;
                for (auto it = edits.begin(); it != edits.end(); ++it) {
                    auto p = offset + glm::ivec3(glm::fromIdx(it->first));
                    if(p.x < SIZE && p.y < SIZE && p.z < SIZE) {
                        setCell(glm::toIdx(p, glm::ucvec3(SIZE)), it->second.value);
                    }
                }
            }
        }
        {
            std::vector<dualmc::Vertex> dmc_vertices;
            std::vector<dualmc::Tri> dmc_tris;

            dualmc::DualMC<float> builder;
            builder.buildTris(&grid.front(), SIZE, SIZE, SIZE, iso, materials.texture_ids.data(), materials.roughnesses.data(),
                manifold, dmc_vertices, dmc_tris);

            tmp.clear();
            tmp.reserve(dmc_vertices.size());
            std::transform(dmc_vertices.begin(), dmc_vertices.end(), std::back_inserter(tmp), [](const dualmc::Vertex &v) {
                constexpr auto HALF_MANTISSA = 10;
                return render::VertexData(glm::vec3(meshopt_quantizeFloat(v.x, HALF_MANTISSA),meshopt_quantizeFloat(v.y, HALF_MANTISSA),
                    meshopt_quantizeFloat(v.z, HALF_MANTISSA)), v.w, glm::vec3(0));
            });

            out.reserve(dmc_tris.size() * 3);
            for (const auto& t: dmc_tris) {
                glm::vec3 edge1 = tmp[t.i1].Position - tmp[t.i0].Position;
                glm::vec3 edge2 = tmp[t.i2].Position - tmp[t.i0].Position;
                glm::vec3 normal = glm::normalize(glm::cross(edge1, edge2));

                if(!reordering || glm::length2(edge1) > glm::length2(edge2)) {
                    out.push_back(t.i0);
                    out.push_back(t.i1);
                    out.push_back(t.i2);
                } else {
                    out.push_back(t.i2);
                    out.push_back(t.i0);
                    out.push_back(t.i1);
                }

                tmp[t.i0].Normal += normal;
                tmp[t.i1].Normal += normal;
                tmp[t.i2].Normal += normal;
            }

            for(auto& v: tmp) {
                v.Normal = glm::normalize(v.Normal);
            }
        }
    }
    void FlatDualMC::render(const surrounding::corners &surrounding, render::LodModel::LodData &out, std::vector<render::VertexData> &tmp, Layer layer) const {
        render(surrounding, out.first.indices, tmp, layer);

        out.second = simplify_lod(out.first.indices, tmp, loadedLevels);
        std::transform(tmp.begin(), tmp.end(), std::back_inserter(out.first.vertices), PACK);
        optimize_fetch(out.first);
    }

    void FlatDualMC::getTerrainModels(terrain_draw_call out, const std::optional<geometry::Frustum> &frustum, const world::cell_pos& offset, bool solid) const {
        for (const auto& area: areas) {
            const auto absolute = area.second.first.absolute;
            const auto rotation = glm::toMat4(area.second.first.absolute.rotation);
            const auto rot_abs = glm::abs(rotation);
            for (const auto& chunk: area.second.second) {
                const auto buffer = solid ? chunk.second.main.get() : chunk.second.transparent.get();
                if (!buffer) // TODO: Fix bad branch prediction
                    continue;

                const auto cPos = glm::multiply(chunk.first);
                if (frustum) {
                    const glm::vec3 fCenterPos = absolute.computeChild(cPos + world::cell_pos(world::CHUNK_LENGTH / 2)) - (world::world_pos)offset;
                    if (!frustum.value().contains(geometry::faabb::FromCenterAbs(fCenterPos, glm::vec3(world::CHUNK_LENGTH / 2), rot_abs)))
                        continue;
                }
                const glm::vec3 fMinPos = absolute.computeChild(cPos) - (world::world_pos)offset;
                out(glm::translate(glm::mat4(1), fMinPos) * rotation, buffer, area.second.first, cPos);
            }
        }
    }

    void FlatDualMC::getTerrainModels(terrain_draw_call out, const world::world_pos& from, float far_dist, const std::vector<glm::vec3> &occlusion, const world::cell_pos &offset, bool solid) const {
        const auto start = world::world_pos(glm::divide(from));
        const auto dist = far_dist / world::CHUNK_LENGTH;
        for (const auto& area: areas) {
            const auto &absolute = area.second.first.absolute;
            const auto area_offset = glm::divide(absolute.position);
            const auto rotation = glm::toMat4(absolute.rotation);
            robin_hood::unordered_set<world::chunk_pos> done;
            for (const auto& occ: occlusion) {
                geometry::Ray::iterator it(geometry::Ray(start, absolute.rotation * occ, dist));
                glm::llvec3 point;
                while (it.next(point)) {
                    const auto &chunks = area.second.second;
                    if (const auto it = chunks.find(glm::lvec3(point) - area_offset); it != chunks.end()) {
                        const auto buffer = solid ? it->second.main.get() : it->second.transparent.get();
                        if (buffer != NULL && done.insert(it->first).second) {
                            const auto cPos = glm::multiply(it->first);
                            const glm::vec3 fPos = absolute.computeChild(cPos) - (world::world_pos)offset;
                            out(glm::translate(glm::mat4(1), fPos) * rotation, buffer, area.second.first, cPos);
                            break;
                        }
                    }
                }
            }
        }
    }

    void FlatDualMC::getUniqueModels(unique_draw_call out, const world::Elements& l, const std::optional<geometry::Frustum>& frustum, const world::cell_pos& offset) const {
        for (const auto& ref: l.parts) {
            const world::part_id id = l.withFlag(ref.val).val;
            const auto node = l.findPart(id);

            if (const auto buffer = getModel(id)) {
                const auto rotation = glm::toMat4(node->absolute.rotation);
                const glm::vec3 fMinPos = node->absolute.position - (world::world_pos)offset;
                if (frustum) {
                    if (!frustum.value().contains(geometry::faabb::FromMin(fMinPos, node->get()->size, rotation)))
                        continue;
                }
                out(glm::translate(glm::mat4(1), fMinPos) * rotation, buffer);
            }
        }
    }

    void FlatDualMC::getInstancedModels(instanced_draw_call out, const world::Elements& l, const std::optional<geometry::Frustum>& frustum, const world::cell_pos& offset, world::node_id ignore) const {
        std::vector<glm::mat4> matrices;
        l.models.iter([&] (const world::model_id& id, const world::Model& model) {
            if (const auto buffer = getModel(id)) {
                for (const auto& ref: model.instances) {
                    const auto iid = l.withFlag(ref.val);
                    if (iid == ignore)
                        continue;

                    const auto node = world::NodeOf<world::Instance>::Make(l.findNode(iid));
                    const auto rotation = glm::toMat4(node->absolute.rotation);
                    const glm::vec3 fMinPos = node->absolute.position - (world::world_pos)offset;
                    if (frustum) {
                        if (!frustum.value().contains(model.collision.rotate(node->absolute.rotation) + fMinPos))
                            continue;
                    }
                    matrices.push_back(glm::scale(glm::translate(glm::mat4(1), fMinPos) * rotation, model.scale));
                }
                if (!matrices.empty()) {
                    out(matrices, buffer);
                    matrices.clear();
                }
            }
        });
    }

}