#include "./Universe.hpp"

#include <Tracy.hpp>
#include <common/TracySystem.hpp>
#include <filesystem>
#include <random>
#include "./Serializer.hpp"
#include "./SharedChunk.hpp"
#include "modules/core/PlanetGenerator.hpp"
#include "core/world/iterators.hpp"

#undef LOG_PREFIX
#define LOG_PREFIX "Server: "

using namespace world::server;

Universe::Universe(const Universe::options&o, Serializer* serializer, world::server_handle *const h):
    handle(h), opts(o),
    chunkFactory(h ? (ChunkFactory*)new SharedChunkFactory() : new StandaloneChunkFactory())
{
    running = true;

    std::filesystem::create_directories(opts.folderPath);
    loadIndex();

    // Workers
    for (size_t i = 0; i < std::max<size_t>(1, std::thread::hardware_concurrency() / 2 - 1); i++) {
        workers.emplace_back([&, serializer] {
#if TRACY_ENABLE
            tracy::SetThreadName("Chunks");
#endif
            const auto read_ctx = serializer->Dicts().make_reader();
            const auto write_ctx = serializer->Dicts().make_writer();
            while (running) {
                if(save_task_t task; saveAreaQueue.pop(task)) {
                    //NOTE: must always save before load to avoid chunk regen
                    //MAYBE: queue.take to avoid concurent write or duplicated work on fast move
                    ZoneScopedN("ProcessSaveArea");
                    std::ostringstream out;
                    if (!task.second.second->isModified()) {
                        out.setstate(std::ios_base::badbit);
                    }
                    const auto average = task.second.second->write(out);
                    const auto rcPos = glm::split(task.second.first);
                    const auto reg = task.first.second->getRegion(opts.folderPath, area_region_ref(task.first.first.val.index(), rcPos.first));
                    reg->write(rcPos.second, write_ctx, out.str(), std::make_optional(average));
                } else if (std::pair<part_id, std::shared_ptr<Part>> task; loadPartQueue.pop(task)) {
                    ZoneScopedN("ProcessLoadPart");
                    Region::Read(opts.folderPath, task.first.val.index(), read_ctx, [&] (const region_chunk_pos& pos, const std::vector<char>& data) {
                        ZoneScopedN("ProcessRead");
                        io::vec_istream idata(data);
                        std::istream iss(&idata);
                        loadedQueue.push({node_chunk_pos{task.first.val, pos}, chunkFactory->create(iss)});
                    });
                } else if(std::pair<part_id, std::shared_ptr<Part>> task; savePartQueue.pop(task)) {
                    //NOTE: must always fully load parts before save to avoid lost chunks
                    ZoneScopedN("ProcessSavePart");
                    const auto SIZE = task.second->chunks.size();
                    const auto alive = std::count_if(task.second->chunks.begin(), task.second->chunks.end(),
                        [](const std::shared_ptr<world::EdittableChunk> &ptr) -> bool { return ptr != nullptr; });
                    size_t i = 0;
                    Region::Write(opts.folderPath, task.first.val.index(), write_ctx, alive, [&] {
                        while (i < SIZE && !task.second->chunks[i]) {
                            i++;
                        }
                        assert(i < SIZE);
                        std::ostringstream out;
                        std::static_pointer_cast<Chunk>(task.second->chunks[i])->write(out);
                        const auto pos = task.second->getPos(i);
                        i++;
                        return std::make_pair(pos, out.str());
                    });
                } else if (std::pair<area_chunk_pos, std::shared_ptr<Area>> task; loadAreaQueue.pop(task)) {
                    //MAYBE: loadQueue.take to avoid duplicated work on fast move
                    ZoneScopedN("ProcessLoadArea");
                    const auto &pos = task.first;
                    const auto rcPos = glm::split(pos.chunk);
                    const auto reg = task.second->getRegion(opts.folderPath, area_region_ref(pos.area.val.index(), rcPos.first));
                    std::vector<char> data;
                    if(reg->read(rcPos.second, read_ctx, data)) {
                        ZoneScopedN("ProcessRead");
                        io::vec_istream idata(data);
                        std::istream iss(&idata);
                        loadedQueue.push({node_chunk_pos{pos.area.val, pos.chunk}, chunkFactory->create(iss)});
                    } else {
                        ZoneScopedN("ProcessGenerate");
                        loadedQueue.push({node_chunk_pos{pos.area.val, pos.chunk}, chunkFactory->create(pos.chunk, task.second->getGenerator())});
                    }
                } else {
                    loadAreaQueue.wait();
                }
            }
        });
    }
#ifndef STANDALONE_SERVER
    if (handle)
        handle->elements = elements.make_ref<world::Elements>();
#endif
}
Universe::~Universe() {
    {
        const auto elements = getElements();
        for(auto& ref: elements->areas) {
            const auto id = elements->withFlag(ref);
            const auto area = elements->findArea(id.val)->get();
            auto &chunks = area->setChunks();
            for (auto it_c = chunks.begin(); it_c != chunks.end(); it_c = chunks.erase(it_c)) {
                saveAreaQueue.emplace(std::make_pair(id.val, area), std::make_pair(it_c->first, std::static_pointer_cast<Chunk>(it_c->second)));
            }
        }
        for(auto& ref: elements->parts) {
            const auto id = elements->withFlag(ref);
            const auto part = elements->findPart(id.val)->get();
            if (!part->chunks.empty())
                savePartQueue.emplace(id.val, part);
        }
    }

    if (auto size = saveAreaQueue.size() + savePartQueue.size(); size > 0) {
        const auto SAVE_CHECK_TIME = 1000;
        do {
            loadAreaQueue.notify_all();
            LOG_I("Saving " << size << " chunks");
            std::this_thread::sleep_for(std::chrono::milliseconds(SAVE_CHECK_TIME));
            size = saveAreaQueue.size() + savePartQueue.size();
        } while (size > 0);
    }
    writeIndex();

    running = false;
    loadAreaQueue.notify_all();

    for (auto &worker: workers) {
        if (worker.joinable())
            worker.join();
    }
    LOG_D("Universe disappeared");
}

void Universe::update(float deltaTime) {
    ZoneScopedN("Update");

    const auto lock = setElements();
    auto &elements = *lock;

    elements.hierarchy.for_each([&](Elements::hierarchy_entry &entry) {
        //vel += acc * deltaTime
        entry.relative.position += entry.vel.position * deltaTime;
        entry.relative.rotation = glm::slerp(identity_pivot, entry.vel.rotation, deltaTime) * entry.relative.rotation;

        const auto node = elements.nodes.directly_at(entry.self);
        node->absolute = elements.computeAbsolute(entry.parent, entry.relative);
        //TODO: full physics
    });
}

void Universe::upgrade() {
    ZoneScopedN("Upgrade");

    const auto lock = setElements();
    auto &elements = *lock;

    { // Random tick
        const auto size = elements.nodes.size();
        auto rng = std::mt19937(std::random_device()());
        const auto randomTick = opts.randomTick;
        auto dist = std::uniform_int_distribution<generational::id::idx_t>(0, randomTick);
        for (size_t part = 0; part <= size / randomTick; part++) {
            const node_ref ref = part * randomTick + dist(rng);
            if (const auto id = elements.withFlag(ref)) {
                elements.nodes.set_flag(Elements::Set<Elements::Flag::Moved>(id));
            }
        }
    }

    std::vector<world_pos> movedLoaders;
    elements.nodes.iter([&](const node_id &id, const Node &node) {
        if (Elements::Has<Elements::Flag::ChunkLoader>(id) && (
                Elements::Has<Elements::Flag::Moved>(id) || Elements::Has<Elements::Flag::Added>(id)))
            movedLoaders.push_back(node.absolute.position);
    });

    { // Update areas && parts
        ZoneScopedN("World");
#if TRACY_ENABLE
        size_t chunk_count = 0;
        size_t region_count = 0;
        bool allLazy = true;
#endif
        const auto keepDist2 = glm::pow2<int>(opts.keepDistance);
        const auto regionKeepDist2 = glm::pow2(opts.keepDistance + REGION_LENGTH * 2);
        const bool queuesEmpty = loadAreaQueue.empty() && saveAreaQueue.empty();
        for (const auto& ref: elements.areas) {
            ZoneScopedN("Area");
            const auto id = elements.withFlag(ref.val);
            const auto a_node = elements.findArea(id.val);
            assert(a_node);
            const auto area = a_node->get();
            const auto areaRange = glm::pow2(opts.keepDistance + area->getChunks().getRadius());

            //MAYBE: use hierarchy
            std::vector<chunk_pos> inRangeLoaders;
            elements.nodes.iter([&](const node_id& id, const Node& cur) {
                if (Elements::Has<Elements::Flag::ChunkLoader>(id)) {
                    const chunk_pos dist = glm::divide(cur.absolute.position - a_node->absolute.position);
                    if (glm::length2(dist) <= areaRange) //MAYBE: unique
                        inRangeLoaders.push_back(dist);
                }
            });

            auto &chunks = area->setChunks();
            bool lazyArea = queuesEmpty;
            if (inRangeLoaders.empty()) {
                if (!chunks.empty()) {
                    for (auto it_c = chunks.begin(); it_c != chunks.end();) {
                        saveAreaQueue.emplace(std::make_pair(id.val, area), std::make_pair(it_c->first, std::static_pointer_cast<Chunk>(it_c->second)));
                        lazyArea = false;
                        it_c = chunks.erase(it_c);
                    }
                }
            } else {
                { // Check alive chunks
                    ZoneScopedN("Alive");
                    for (auto it_c = chunks.begin(); it_c != chunks.end();) {
                        if (std::any_of(inRangeLoaders.begin(), inRangeLoaders.end(), [&] (const chunk_pos& dist) {
                            return glm::length2(dist - chunk_pos((glm::qua<double>)a_node->absolute.rotation * glm::dvec3(it_c->first))) <= keepDist2;
                        })) {
                            ++it_c;
#if TRACY_ENABLE
                            chunk_count++;
#endif
                        } else {
                            saveAreaQueue.emplace(std::make_pair(id.val, area), std::make_pair(it_c->first, std::static_pointer_cast<Chunk>(it_c->second)));
                            lazyArea = false;
                            it_c = chunks.erase(it_c);
                        }
                    }
                }
                { // Enqueue missing chunks
                    ZoneScopedN("Missing");
                    std::optional<load_queue_t::inserter_t> handle;
                    const lpivot inv_rot = glm::conjugate(a_node->absolute.rotation);
                    for (const auto& to: movedLoaders) {
                        const chunk_pos dist = glm::divide(inv_rot * (to - a_node->absolute.position));
                        if (glm::length2(dist) > areaRange)
                            continue;

                        //TODO: need dist so no easy sphere fill
                        const auto loadDistance = opts.loadDistance;
                        for (int x = -loadDistance; x <= loadDistance; x++) {
                        for (int y = -loadDistance; y <= loadDistance; y++) {
                        for (int z = -loadDistance; z <= loadDistance; z++) {
                            const auto dist2 = x * x + y * y + z * z;
                            if (dist2 <= loadDistance * loadDistance) {
                                const auto p = dist + chunk_pos(x, y, z);
                                if (chunks.inRange(p) && chunks.find(p) == chunks.end()) {
                                    if (!handle) { handle.emplace(loadAreaQueue.inserter()); }
                                    handle.value().first(area_chunk_pos{id.val, p}, area, -dist2);
                                    lazyArea = false;
                                }
                            }
                        }}}
                    }
                }
            }
#if TRACY_ENABLE
            allLazy &= lazyArea;
#endif
            if (lazyArea) { // Clear un-used regions
                ZoneScopedN("Region");
                const auto unique = area->getRegions(); // MAYBE: shared then unique
#if TRACY_ENABLE
                region_count += unique->size();
#endif
                for (auto it_r = unique->begin(); it_r != unique->end(); ++it_r) {
                    if(std::none_of(inRangeLoaders.begin(), inRangeLoaders.end(), [&] (const chunk_pos& dist) {
                        return glm::length2(dist - chunk_pos(it_r->first) * chunk_pos(REGION_LENGTH)) <= regionKeepDist2;
                    })) {
                        unique->erase(it_r); //FIXME: may wait for os file access (long)
                        break; //NOTE: save one only max per update
                    }
                }
            } else {
                loadAreaQueue.notify_all();
            }
        }
        // PARTS
        for (const auto& ref: elements.parts) {
            ZoneScopedN("Part");
            const auto id = elements.withFlag(ref.val);
            const auto r_node = elements.findPart(id.val);
            assert(r_node);
            const auto part = r_node->get();

            //FIXME: pivot point
            const auto partOffset = part->chunkSize() / chunk_pos(2);

            //MAYBE: use hierarchy
            const auto inRangeLoaders = elements.nodes.any([&](const node_id& id, const Node& cur) {
                if (Elements::Has<Elements::Flag::ChunkLoader>(id)) {
                    const chunk_pos dist = glm::divide(cur.absolute.position - r_node->absolute.position);
                    return glm::length2(dist + partOffset) <= keepDist2;
                }
                return false;
            });

            if (!inRangeLoaders) {
                if (!part->chunks.empty()) {
                    savePartQueue.emplace(id.val, part);
                    r_node->content = std::make_shared<Part>(part->size);
                }
            } else {
                if (part->chunks.empty()) {
                    loadPartQueue.emplace(id.val, part);
                    part->allocate();
                }
            }
        }
#if TRACY_ENABLE
        TracyPlot("ChunkCount", static_cast<int64_t>(chunk_count));
        if(allLazy) {
            TracyPlot("Region", static_cast<int64_t>(region_count));
        }
        TracyPlot("AreaLoad", static_cast<int64_t>(loadAreaQueue.size()));
        TracyPlot("AreaUnload", static_cast<int64_t>(saveAreaQueue.size()));
        TracyPlot("PartLoad", static_cast<int64_t>(loadPartQueue.size()));
        TracyPlot("PartUnload", static_cast<int64_t>(savePartQueue.size()));
#endif
        loadAreaQueue.notify_all();
    }
    /* FIXME: { // Update entities
        ZoneScopedN("Entities");
        size_t item_count = 0;
        entities.remove([&](entity_id type, Entity &val) {
            //TODO: find mass center
            const auto OFFSET = glm::vec3(val.size) * val.scale / 2.f;
            const auto DIST2 = glm::pow2(loadDistance); //NOTE: must be < keepDistance + glm::length(OFFSET) to avoid collision miss
            val.instances.remove([&](entity_id, Entity::Instance &inst) {
                if (type == PLAYER_ENTITY_ID) {
                    //MAYBE: update players ?
                    item_count++;
                    return false;
                }

                //TODO: partition space
                auto pos = inst.pos + OFFSET;
                for(const auto& area: areas) {
                    //NOTE: consider entity mass negligible
                    inst.velocity += area.second->getGravity((pos - area.second->getOffset()).as_cell()) * deltaTime;
                }
                //MAYBE: friction
                if (move(pos, inst.velocity, 1, glm::max_axis(OFFSET))) {
                    inst.velocity = glm::vec3(0);
                }
                inst.pos = pos - OFFSET;
                if (entities.at(PLAYER_ENTITY_ID).instances.contains([&](const entity_id&, const Entity::Instance& player) {
                    return glm::length2(glm::divide((player.pos - inst.pos).as_cell())) <= DIST2;
                })) {
                    item_count++;
                    return false;
                }
                //TODO: Save to files if !temporary
                return true;
            });
            return !val.permanant && val.instances.empty();
        });
        TracyPlot("EntityCount", static_cast<int64_t>(item_count));

    }*/

    { // Store loaded chunks
        ZoneScopedN("Load");
        loadedQueue.extractor([&](std::pair<node_chunk_pos, std::shared_ptr<Chunk>> &loaded) {
            if (Elements::Is<Elements::Type::Area>(loaded.first.node)) {
                if (const auto area = elements.findArea(loaded.first.node.val))
                    area->get()->setChunks().emplace(loaded.first.chunk, std::move(loaded.second));
            } else {
                if (const auto part = elements.findPart(loaded.first.node.val))
                    part->get()->emplace(loaded.first.chunk, std::move(loaded.second));
            }
        });
    }
}

constexpr auto INDEX_FILE = "/nodes.idx";
void Universe::loadIndex() {
    const auto elements = setElements();
    assert(elements->nodes.empty());
    io::ifstream index(opts.folderPath + INDEX_FILE);
    if(index.good()) {
        { // Spawn
            node_ref spawnParent;
            index.read_cast(spawnParent);
            spawnPoint.parent = generational::id(spawnParent.val, 0);
            index.read_cast(spawnPoint.position);
        }
        elements->readIndex(index);
    } else {
        LOG_E("No index file!!! Probably a new world...");
        spawnPoint.position = world_pos(100, 0, 0);
        const auto p = elements->createInstance(Elements::start_point(generational::id(), transform(), velocity(offset_pos(), glm::angleAxis(.1f, glm::vec3(0, 0, 1)))), generational::id(0));
        elements->createInstance(Elements::start_point(generational::id(), transform(), velocity(offset_pos(1))), generational::id(0));
        elements->createInstance(Elements::start_point(p.val, transform(glm::vec3(1000, 0, 0))), generational::id(0));
        //TODO: generate universe
        {
            const auto radius = 1 << 6;
            const auto surface = radius * CHUNK_LENGTH * 3 / 4;
            generator::RoundPlanet::Params opts(surface);
            Area::params params{radius, 0, std::vector<char>(sizeof(opts))};
            memcpy(params.params.data(), &opts, params.params.size());
            elements->createArea(Elements::start_point(generational::id(), transform(world_pos(-surface, 0, 0)), velocity(offset_pos(), glm::angleAxis(.01f, glm::vec3(0, 1, 0)))), params);
        }
    }
    index.close();
}
void Universe::writeIndex() {
    io::ofstream index(opts.folderPath + INDEX_FILE, std::ios::out | std::ios::binary);
    if(index.good()) {
        index.write_cast(spawnPoint.parent.val.index());
        index.write_cast(spawnPoint.position);
        getElements()->writeIndex(index);
        if(!index.good())
            LOG_E("Index file write error");
    } else {
        LOG_E("Failed to open index file");
    }

    index.close();
}

std::shared_ptr<Chunk> Universe::findChunk(const node_chunk_pos& p) const {
    const auto lock = getElements();
    if (Elements::Is<Elements::Type::Area>(p.node)) {
        if (auto area = lock->findArea(p.node.val)) {
            return std::static_pointer_cast<Chunk>(area->get()->getChunks().findInRange(p.chunk));
        }
    } else {
        if (auto part = lock->findPart(p.node.val)) {
            return std::static_pointer_cast<Chunk>(part->get()->findInRange(p.chunk));
        }
    }
    return nullptr;
}
const world::Node* Universe::getPlayer(node_id id) const {
    //FIXME: assert(IsPlayer(id))
    return getElements()->findNode(id);
}
world::Node* Universe::setPlayer(node_id id) {
    //FIXME: assert(IsPlayer(id))
    return setElements()->findNode(id);
}

void Universe::fill(const action::FillShape& fill, fill_edits *const edits) {
    const auto lock = setElements();
    ZoneScopedN("Fill");
    //TODO: ItemList inventory;
    //TODO: multipart
    const auto anchor_id = lock->withFlag(fill.pos.node);
    const auto node = lock->findNode(anchor_id);
    if (!node)
        return;

    const auto applyCB = [&](std::shared_ptr<Chunk>& ck, chunk_pos ck_pos, chunk_voxel_idx idx, Voxel/*prev*/, Voxel next, float delay) {
        if (edits)
            (*edits)[ck_pos].push_back(ChunkEdits::Edit{next, delay, idx});
        //TODO: apply break table
        //TODO: inventory
        ck->replace(idx, next, delay);
    };
    const auto splitCBer = [&](const Elements::start_point& rtf, const auto& getter) {
        return [&] (const robin_hood::unordered_set<voxel_pos> &voxels) {
            //MAYBE: make it void after air propagation setup
            const auto cleanVoxel = Voxel(0, Voxel::DENSITY_MAX);
            voxel_pos min = voxel_pos(INT64_MAX);
            voxel_pos max = voxel_pos(INT64_MIN);
            for (const auto& full: voxels) {
                min = glm::min<3, long long>(min, full);
                max = glm::max<3, long long>(max, full);
            }
            auto start = rtf;
            start.relative.position += min;
            const auto part_id = lock->createPart(start, max - min);
            const auto node = lock->findPart(part_id);
            auto& part = *node->get();
            {
                const auto chunk_size = part.chunkSize();
                const auto chunk_count = chunk_size.x * chunk_size.y * chunk_size.z;
                part.chunks.reserve(chunk_count);
                for (long i = 0; i < chunk_count; i++)
                    part.chunks.push_back(chunkFactory->create());
            }
            world::iterator::last_chunk<Chunk> src_ck;
            world::iterator::last_chunk<Chunk> dst_ck{nullptr, chunk_pos(-1)};
            Voxel cur;
            for (const auto& full: voxels) {
                const auto src_split = glm::splitIdx(full);
                if (getter(src_split, cur, src_ck) && cur.is_solid()) {
                    if (edits)
                        (*edits)[src_ck.pos].push_back(ChunkEdits::Edit{cleanVoxel, fill.radius * world::iterator::FILL_DELAY, src_split.second});

                    src_ck.ptr->replace(src_split.second, cleanVoxel);
                    const auto dst_split = glm::splitIdx(full - min);
                    if (dst_split.first != dst_ck.pos) {
                        dst_ck.ptr = std::static_pointer_cast<Chunk>(part.findInRange(dst_split.first));
                        assert(dst_ck.ptr);
                        dst_ck.pos = dst_split.first;
                    }
                    dst_ck.ptr->set(dst_split.second, cur);
                }
            }
        };
    };
    switch (Elements::GetType(anchor_id)) {
        case Elements::Type::Area: {
            const auto& chunks = NodeOf<Area>::Make(node)->get()->getChunks();
            const auto rtf = Elements::start_point(anchor_id, transform());
            world::iterator::Apply<Chunk>(chunks, fill, applyCB);
            world::iterator::Split<Chunk>(chunks, fill, opts.floodFillLimit, splitCBer(rtf,
                [&](const std::pair<glm::lvec3, glm::idx>& split, Voxel& out, world::iterator::last_chunk<Chunk>& ck) {
                    return world::iterator::GetVoxel<Chunk>(chunks, split, out, ck);
                }
            ));
            break;
        }
        /* TODO: case Elements::Type::Part: {
            const auto& chunks = *NodeOf<Part>::Make(node)->get()->chunks;
            const auto entry = lock->findHierarchy(anchor_id);
            const auto rtf = start_point(entry->parent, entry->relative, entry->velocity);
            world::iterator::Apply<Chunk>(chunks, fill, applyCB);
            world::iterator::Split<Chunk>(chunks, fill, opts.floodFillLimit, splitCBer(rtf,
                [&](const std::pair<glm::lvec3, glm::idx>& split, Voxel& out, world::iterator::last_chunk<Chunk>& ck) {
                    return world::iterator::GetVoxel<Chunk>(chunks, split, out, ck);
                }
            ));
            break;
        }*/
        default:
            LOG_E("Unhandled fill target type " << (int)Elements::GetType(anchor_id));
            break;
    }
}

world::node_id Universe::addPlayer(std::optional<relative_pos> spawn) {
    const auto &sp = spawn.value_or(spawnPoint);
    const auto CHUNK_LOADER = Elements::Set<Elements::Flag::ChunkLoader>(0);
    return setElements()->createInstance(Elements::start_point(sp), generational::id(0, 0), CHUNK_LOADER).val;
}
void Universe::removePlayer(node_id id) {
    //FIXME: save player infos ?
    setElements()->remove(id);
}
bool Universe::movePlayer(node_id id, const relative_pos& to) {
    if (auto player = getPlayer(id)) {
        //auto target = elements.computeAbsolute(elements.findParent(id), world::transform(to.position));
        //TODO: check dist + collision from a to b
        setElements()->move(id, to);
        return true;
    }
    return false;
}