#include "FlatDualMC.hpp"

#include "../../core/world/EdittableChunk.hpp"
#include "../../core/world/Area.hpp"
#include "../../core/world/materials.hpp"
#include <Tracy.hpp>   // NOLINT
#include <common/TracySystem.hpp> // NOLINT
#include <imgui.h>     // NOLINT
#include <toml.h>
#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) {
                    if (entity_area_job_t job; entityLoadQueue.pop(job)) {
                        ZoneScopedN("Entity");
                        render::Model::Data data;
                        render::Model::Data::indices_t idx;
                        for (uint8_t x = 0; x < job.size.x; x++) {
                        for (uint8_t y = 0; y < job.size.y; y++) {
                        for (uint8_t z = 0; z < job.size.z; z++) {
                            surrounding::corners sur;
                            surrounding::load(sur, job.size, glm::ucvec3(x, y, z), job.area);
                            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 * CHUNK_LENGTH, y * CHUNK_LENGTH, z * 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);
                        entityLoadedQueue.emplace(job.id, data);
                    } else if (std::pair<area_<chunk_pos>, surrounding::corners> ctx; loadQueue.pop(ctx)) {
                        ZoneScopedN("Chunk");
                        std::pair<render::LodModel::LodData, render::LodModel::LodData> data;
                        if (transparency) {
                            render(ctx.second, data.first, tmp, Layer::Solid);
                            render(ctx.second, data.second, tmp, Layer::Transparent);
                        } else {
                            render(ctx.second, data.first, tmp, Layer::Both);
                        }
                        //TODO: direct upload with vulkan
                        loadedQueue.emplace(ctx.first, data);
                    } else {
                        loadQueue.wait();
                    }
                }
            });
        }
    }
    FlatDualMC::~FlatDualMC() {
        running = false;
        loadQueue.notify_all();

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

        //TODO: prefer unique_ptr
        for(auto& buffer: buffers) {
            for(auto& val: buffer.second.second) {
                delete val.second.first;
                delete val.second.second;
            }
        }
    }

    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) * CHUNK_LENGTH, (keepDistance + .5) * CHUNK_LENGTH);
    }

    size_t FlatDualMC::getQueueSize() {
        return loadQueue.size();
    }

    void FlatDualMC::enqueue(const area_<chunk_pos> &pos, const chunk_pos &offset, const world::ChunkContainer &data)
    {
        ZoneScopedN("EnqueueContouring");
        const auto dist2 = glm::length2(offset - pos.second);
        if (dist2 <= loadDistance * loadDistance) {
            surrounding::corners surrounding;
            if(surrounding::load(surrounding, pos.second, data)) {
                loadQueue.push(pos, surrounding, -dist2);
            }
        }
     }

    void FlatDualMC::onUpdate(const area_<chunk_pos> &pos, const chunk_pos &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(std::make_pair(pos.first, pos.second - surrounding::g_corner_offsets[i]), offset, data);
            }
        }
    }

    void FlatDualMC::onNotify(const area_<chunk_pos> &pos, const chunk_pos &offset, const world::ChunkContainer &data) {
        const auto it = buffers.find(pos.first);
        if(it == buffers.end() || it->second.second.find(pos.second) == it->second.second.end()) {
            enqueue(pos, offset, data);
        }
    }

    void FlatDualMC::onEntityLoad(size_t id, std::istream &in) {
        if(!entities.get(id))
            entities.set_emplace(id, nullptr);
        entityLoadedQueue.emplace(id, render::Model::Data(in));
    }
    void FlatDualMC::onEntityLoad(size_t id, const glm::ucvec3& size, const std::vector<std::shared_ptr<world::Chunk>>& area) {
        if(!entities.get(id))
            entities.set_emplace(id, nullptr);
        entityLoadQueue.emplace(entity_area_job_t{id, size, area});
        loadQueue.notify_one();
    }
    void FlatDualMC::onEntityUnload(size_t id) {
        entities.erase(id);
    }

    void FlatDualMC::update(const voxel_pos& pos, const world::client::area_map& areas) {
        ZoneScopedN("Ct");
        TracyPlot("CtLoad", static_cast<int64_t>(loadQueue.size() + entityLoadQueue.size()));
        //MAYBE: clear out of range loadQueue.trim(keepDistance * keepDistance)
        TracyPlot("CtLoaded", static_cast<int64_t>(loadedQueue.size() + entityLoadedQueue.size()));

        {
            std::pair<size_t, render::Model::Data> out;
            for(auto handle = entityLoadedQueue.extractor(); handle.first(out);) {
                if (auto entity = entities.get(out.first))
                    *entity = render::Model::Create(out.second);
            }
        }

        std::pair<area_<chunk_pos>, std::pair<render::LodModel::LodData, render::LodModel::LodData>> out;
        for(auto handle = loadedQueue.extractor(); handle.first(out);) {
            const auto bufferSolid = out.second.first.first.empty() ? NULL : render::LodModel::Create(out.second.first).release();
            const auto bufferTrans = out.second.second.first.empty() ? NULL : render::LodModel::Create(out.second.second).release();
            auto &bfs = buffers[out.first.first].second; //NOTE: buffer.first uninitialized (will be set in clear())
            if (const auto it = bfs.find(out.first.second); it != bfs.end()) {
                if (it->second.first != NULL)
                    delete it->second.first;
                if (it->second.second != NULL)
                    delete it->second.second;

                it->second.first = bufferSolid;
                it->second.second = bufferTrans;
            } else {
                bfs.emplace(out.first.second, std::make_pair(bufferSolid, bufferTrans));
            }
        }
        size_t buffer_count = 0;
        {
            ZoneScopedN("CtUpdate");
            const auto levelMax = loadedLevels.size();
            auto it_a = buffers.begin();
            while (it_a != buffers.end()) { // Remove out of range buffers
                if (const auto area = areas.find(it_a->first); area != areas.end()) {
                    //Update
                    std::get<0>(it_a->second.first) = area->second->getOffset();
                    std::get<1>(it_a->second.first) = area->second->getChunks().getRadius() * static_cast<long long>(CHUNK_LENGTH);
                    const auto centerV = pos - std::get<0>(it_a->second.first).as_voxel();
                    if (const auto surface = area->second->getCurvature()) {
                        const auto dist = glm::max_axis(glm::abs(centerV));
                        std::get<2>(it_a->second.first) = std::clamp<float>((dist - surface.value()) / (std::get<1>(it_a->second.first) - surface.value()), -0.1f, 1.f);
                    }
                    const auto center = glm::divide(centerV);
                    auto &bfs = it_a->second.second;
                    auto it = bfs.begin();
                    while(it != bfs.end()) {
                        if (const auto distRatio = glm::length(glm::dvec3(center - it->first)) / keepDistance; distRatio > 1) {
                            if(it->second.first != NULL)
                                delete it->second.first;
                            if (it->second.second != NULL)
                                delete it->second.second;

                            it = bfs.erase(it);
                        } else {
                            const auto level = std::clamp<size_t>((1 + lod_quality - distRatio) * levelMax * (1 + lod_strength), 0, levelMax);
                            if (it->second.first != NULL) {
                                it->second.first->setLevel(level);
                                buffer_count++;
                            }
                            if (it->second.second != NULL) {
                                it->second.second->setLevel(level);
                                buffer_count++;
                            }
                            ++it;
                        }
                    }
                    ++it_a;
                } else {
                    for(auto& buffer: it_a->second.second) {
                        if (buffer.second.first != NULL)
                            delete buffer.second.first;
                        if(buffer.second.second != NULL)
                            delete buffer.second.second;
                    }
                    it_a = buffers.erase(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));
                ImGui::Selectable(str.c_str(), &lod_levels[i]);
                ImGui::NextColumn();
            }
            ImGui::Columns(1);
        }
        if (ImGui::Button("Flush buffers")) {
            //TODO: prefer unique_ptr
            for(auto& buffer: buffers) {
                for(auto& val: buffer.second.second) {
                    delete val.second.first;
                    delete val.second.second;
                }
            }
            buffers.clear();
        }
    }

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

            for (size_t i = 0; i < surrounding.size(); i++) {
                auto &edits = surrounding[i]->getEdits();
                auto offset = glm::ivec3(surrounding::g_corner_offsets[i]) * 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(p.x, p.y, p.z, 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, world::materials::textures_map.data(), world::materials::roughness.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::getModels(draw_call out, const std::optional<geometry::Frustum> &frustum, const glm::llvec3& offset, int density, bool solid) {
        const auto scaling = glm::scale(glm::mat4(1), glm::vec3(1.f / density));
        for (const auto [_, area] : buffers) {
        for (const auto [pos, buf] : area.second) {
            const auto vPos = glm::multiply(pos);
            const glm::vec3 fPos = (glm::vec3(std::get<0>(area.first).raw_as_long() + vPos - offset * glm::llvec3(density)) + std::get<0>(area.first).offset) / glm::vec3(density);
            const auto buffer = solid ? buf.first : buf.second;
            if (buffer != NULL && (!frustum.has_value() || frustum.value().contains(geometry::Box::fromMin(fPos, glm::vec3(CHUNK_LENGTH / (float)density)))))
                out(glm::translate(scaling, fPos * (float)density), buffer, area.first, vPos);
        }}
    }

    void FlatDualMC::getModels(draw_call out, const glm::ifvec3& from, float far_dist, const std::vector<glm::vec3> &occlusion, const glm::llvec3 &offset, int density, bool solid) {
        const auto scaling = glm::scale(glm::mat4(1), glm::vec3(1.f / density));
        const auto start = glm::ifvec3(glm::divide(from.as_voxel(density)));
        const auto dist = far_dist * density / CHUNK_LENGTH;
        for (const auto [_, area] : buffers) {
            const auto area_offset = glm::divide(std::get<0>(area.first).as_voxel());
            robin_hood::unordered_set<chunk_pos> done;
            for (const auto& occ: occlusion) {
                geometry::Ray::iterator it(geometry::Ray(start, occ, dist));
                glm::llvec3 point;
                while (it.next(point)) {
                    auto it = area.second.find(glm::lvec3(point) - area_offset);
                    const auto buffer = solid ? it->second.first : it->second.second;
                    if(it != area.second.end() && buffer != NULL && done.insert(it->first).second) {
                        const auto vPos = glm::multiply(it->first);
                        const glm::vec3 fPos = glm::vec3(std::get<0>(area.first).raw_as_long() + vPos - offset * glm::llvec3(density)) + std::get<0>(area.first).offset;
                        out(glm::translate(scaling, fPos), buffer, area.first, vPos);
                        break;
                    }
                }
            }
        }
    }

    render::Model* FlatDualMC::getEntityModel(size_t id) {
        if (auto ptr = entities.get(id))
            return ptr->get();
        return nullptr;
    }
}
