#include "DistantUniverse.hpp"

#include <Tracy.hpp>

#include "Area.hpp"
#include "../contouring/Abstract.hpp"
#include "../../core/world/raycast.hpp"
#include "../../core/world/iterators.hpp"
#include "../../core/net/io.hpp"
#include "../../core/utils/logger.hpp"
#include "Chunk.hpp"
#include <random>
#include <algorithm>

using namespace world::client;

DistantUniverse::DistantUniverse(const connection& ct, const Universe::options& opt, const std::string& contouring):
    Universe(contouring), options(opt), serverDistance(0),
    peer(ct, [&](const data::out_view& buf, net::PacketFlags flags){ return onPacket(buf, flags); }) { }
DistantUniverse::~DistantUniverse() { }

void DistantUniverse::update(voxel_pos pos, float deltaTime) {
    const auto cur_chunk = glm::divide(pos);
    const auto chunkChange = cur_chunk != last_chunk;
    last_chunk = cur_chunk;

    { //NOTE: triggers onPacket
        ZoneScopedN("Pull");
        peer.pull();
    }

    { // Update alive areas
        ZoneScopedN("World");
        auto rng = std::mt19937(std::rand());
        const auto contouringThreshold = UINT32_MAX / (1 + contouring->getQueueSize());
        for (auto& area: areas) {
            ZoneScopedN("Area");
            const bool chunkChangeArea = (false && area.second->move(glm::vec3(deltaTime))) || chunkChange; // TODO: area.velocity
            const chunk_pos diff = glm::divide(pos - area.second->getOffset().as_voxel());
            auto &chunks = area.second->setChunks();
            if (glm::length2(diff) <= glm::pow2(options.keepDistance + area.second->getChunks().getRadius())) {
                ZoneScopedN("Alive");
                auto it_c = chunks.begin();
                while (it_c != chunks.end()) {
                    if (glm::length2(diff - it_c->first) > glm::pow2(options.keepDistance)) {
                        it_c = chunks.erase(it_c);
                    } else {
                        if(const auto neighbors = std::dynamic_pointer_cast<Chunk>(it_c->second)->update(deltaTime, rng() < contouringThreshold)) {
                            contouring->onUpdate(std::make_pair(area.first, it_c->first), diff, chunks, neighbors.value());
                        }
                        ++it_c;
                    }
                }
                //FIXME: Produce important duplicated (chunkChangeArea)
                //MAYBE: if dist to unloaded chunk < X
                if (chunkChangeArea || mayQueryChunks) { // Request missing chunks
                    ZoneScopedN("Missing");
                    std::vector<chunk_pos> missingChunks;
                    std::vector<region_pos> missingRegions;
                    std::vector<long> missingDist;
                    const auto cl_area = std::dynamic_pointer_cast<Area>(area.second);
                    //TODO: use easy sphere fill
                    const auto queryDistance = std::min<int>(options.loadDistance, serverDistance);
                    for (int x = -queryDistance; x <= queryDistance; x++) {
                    for (int y = -queryDistance; y <= queryDistance; y++) {
                    for (int z = -queryDistance; z <= queryDistance; z++) {
                        const auto dist2 = x * x + y * y + z * z;
                        const auto p = diff + chunk_pos(x, y, z);
                        if (dist2 <= queryDistance * queryDistance && chunks.inRange(p)) {
                            if (auto it = chunks.find(p); it != chunks.end() &&
                                std::dynamic_pointer_cast<Chunk>(it->second)->isTrusted(options.trustMajorant))
                                //TODO: config accept perfect average
                            {
                                contouring->onNotify(std::make_pair(area.first, p), diff, chunks);
                                continue;
                            }

                            const auto rcPos = glm::split(p);
                            if(auto it_r = cl_area->regionCache.find(rcPos.first); it_r != cl_area->regionCache.end()) {
                                if (auto it_rc = it_r->second.find(rcPos.second); it_rc != it_r->second.end() && (it_rc->second.swap() || options.useAverages)) {
                                    auto ck = std::make_shared<Chunk>(it_rc->second);
                                    ck->invalidate(geometry::Faces::All);
                                    chunks.emplace(p, std::dynamic_pointer_cast<world::Chunk>(ck));
                                    contouring->onNotify(std::make_pair(area.first, p), diff, chunks);
                                    if (ck->isTrusted(options.trustMajorant))
                                        continue;
                                }
                            }

                            const region_pos rPos = glm::divide(p);
                            if (!cl_area->regionCache.contains(rPos)) {
                                cl_area->regionCache.emplace(rPos, 0);
                                missingRegions.push_back(rPos);
                            }
                            if (missingDist.size() >= net::MAX_PENDING_CHUNK_COUNT) {
                                if (dist2 > missingDist.front())
                                    continue;
                                missingDist.erase(missingDist.begin());
                                missingChunks.erase(missingChunks.begin());
                            }

                            auto it = missingDist.begin();
                            auto itv = missingChunks.begin();
                            while (it != missingDist.end()) {
                                if (dist2 > *it)
                                    break;
                                ++it;
                                ++itv;
                            }
                            missingDist.insert(it, dist2);
                            missingChunks.insert(itv, p);
                        }
                    }}}
                    if(!missingChunks.empty()) {
                        auto packet = net::PacketWriter(net::client_packet_type::MISSING_CHUNKS, sizeof(area_id) + missingChunks.size() * sizeof(chunk_pos));
                        packet.write(area.first);
                        packet.write(missingChunks.data(), missingChunks.size() * sizeof(chunk_pos));
                        peer.send(packet.finish());
                    }
                    if(!missingRegions.empty()) {
                        auto packet = net::PacketWriter(net::client_packet_type::MISSING_REGIONS, sizeof(area_id) + missingRegions.size() * sizeof(region_pos));
                        packet.write(area.first);
                        packet.write(missingRegions.data(), missingRegions.size() * sizeof(region_pos));
                        peer.send(packet.finish());
                    }
                    mayQueryChunks = false;
                }
            }
        }
    }
    if (options.movePrediction) {
        entities.remove([&](size_t eId, Entity &entity) {
            entity.instances.remove([&](size_t iId, Universe::Entity::Instance &inst) {
                if (eId == PLAYER_ENTITY_ID.index && iId == playerId)
                    return false;

                // MAYBE: leap after entities packet
                inst.pos += inst.velocity * deltaTime;
                return false;
            });
            return false;
        });
    }

    contouring->update(pos, areas);
}

bool DistantUniverse::onPacket(const data::out_view& buf, net::PacketFlags) {
    using namespace net;

    auto packet = PacketReader(buf);
    server_packet_type type;
    if(!packet.read(type)) {
        LOG_D("Empty packet");
        return true;
    }
    switch (type) {
        case server_packet_type::QUIT: {
            bool is_app = false;
            is_app &= packet.read(is_app);
            uint16_t reason = 0;
            packet.read(reason);
            if (is_app) {
                LOG_E("Disconnected from server (" << reason << ')');
            } else {
                LOG_D("Connection lost (" << reason << ')');
            }
            break;
        }

        case server_packet_type::CAPABILITIES: {
            packet.read(serverDistance);
            auto predictable = false;
            packet.read(predictable);
            if (predictable) {
                packet.read(floodfillLimit);
            } else {
                options.editHandling = false;
                options.editPrediction = false;
            }
            auto packet = PacketWriter(client_packet_type::CAPABILITIES, sizeof(bool));
            packet.write(options.editHandling);
            peer.send(packet.finish());
            break;
        }

        case server_packet_type::COMPRESSION: {
            const auto remain = packet.readRemaining();
            dict.emplace(remain.data(), remain.size());
            LOG_T("Compression dictionnary loaded");
            mayQueryChunks = true;
            break;
        }

        case server_packet_type::TELEPORT: {
            voxel_pos pos;
            if (packet.read(playerId) && packet.read(pos))
                onTeleport(pos);
            break;
        }

        case server_packet_type::MESSAGE: {
            const auto remain = packet.readRemaining();
            onMessage(std::string((const char*)remain.data(), remain.size()));
            break;
        }

        case server_packet_type::AREAS: {
            while(!packet.isDone()) {
                area_id id;
                world::Area::params p;
                if(!packet.read(id) || !packet.read(p))
                    break;

                if (auto it = areas.find(id); it != areas.end()) {
                    std::dynamic_pointer_cast<Area>(it->second)->update(p);
                } else {
                    areas.emplace(id, std::make_shared<Area>(p));
                }
            }
            break;
        }

        case server_packet_type::REGION: {
            ZoneScopedN("Region");
            area_<region_pos> pos;
            if(!packet.read(pos))
                break;

            auto it = areas.find(pos.first);
            if(it == areas.end()) {
                LOG_W("Region area not found " << pos.first.index);
                break;
            }
            auto area = std::dynamic_pointer_cast<world::client::Area>(it->second);
            auto it_r = area->regionCache.find(pos.second);
            if (it_r == area->regionCache.end()) {
                it_r = area->regionCache.emplace(pos.second, 0).first;
            }
            // MAYBE: use Voxel swag bit to flag full void/air chunks to avoiding MISSING_CHUNKS
            // MAYBE: then create virtual or non chunk of single material
            uint16_t full = 0;
            uint16_t total = 0;
            while (!packet.isDone()) {
                region_chunk_pos cpos;
                Voxel voxel;
                if (!packet.read(cpos) || !packet.read(voxel))
                    break;

                it_r->second.insert_or_assign(cpos, voxel);
                total++;
                if (voxel.swap())
                    full++;
            }
            break;
        }

        case server_packet_type::CHUNK: {
            ZoneScopedN("Chunk");
            if (!dict.has_value())
                break;

            if (packet.isDone()) {
                mayQueryChunks = true;
                break;
            }

            area_<chunk_pos> pos;
            if(!packet.read(pos))
                break;

            auto it = areas.find(pos.first);
            if(it == areas.end()) {
                LOG_W("Chunk area not found " << pos.first.index);
                break;
            }
            if(!it->second->getChunks().inRange(pos.second)) {
                LOG_W("Chunk out of area " << pos.first.index);
                break;
            }

            std::vector<char> buffer;
            if(auto err = dict.value().decompress(packet.readRemaining(), buffer)) {
                LOG_E("Corrupted chunk packet " << err.value());
                break;
            }
            data::vec_istream idata(buffer);
            std::istream iss(&idata);
            auto ck = std::make_shared<Chunk>(iss);
            ck->invalidate(geometry::Faces::All);
            auto ptr = std::dynamic_pointer_cast<world::Chunk>(ck);
            auto &chunks = it->second->setChunks();
            if (auto it_c = chunks.find(pos.second); it_c != chunks.end()) {
                it_c->second = ptr;
            } else {
                chunks.emplace(pos.second, ptr);
            }
            break;
        }

        case server_packet_type::EDITS: {
            ZoneScopedN("Edits");
            if (!options.editHandling) {
                LOG_W("Unhandled edit type");
                break;
            }

            const auto fill = packet.read<action::FillShape>();
            if (!fill)
                break;

            world::iterator::ApplySplit<Chunk>(areas, *fill, floodfillLimit, [](std::shared_ptr<Chunk> &ck, chunk_pos, chunk_voxel_idx idx,
                Voxel, Voxel next, float delay) { ck->apply(Chunk::Edit{next, delay, idx}); });

            break;
        }
        case server_packet_type::RAW_EDITS: {
            ZoneScopedN("Raw Edits");
            area_id id;
            if(!packet.read(id))
                break;

            auto it = areas.find(id);
            if(it == areas.end()) {
                LOG_W("Edit area not found " << id.index);
                break;
            }

            while(!packet.isDone()) {
                chunk_pos pos = chunk_pos(INT_MAX);
                packet.read(pos);
                chunk_voxel_idx count = 0;
                packet.read(count);
                if (auto ptr = it->second->setChunks().findInRange(pos)) {
                    auto chunk = std::dynamic_pointer_cast<Chunk>(ptr.value());
                    for (auto i = 0; i < count && !packet.isDone(); i++) {
                        Chunk::Edit edit;
                        packet.read(edit);
                        chunk->apply(edit);
                    }
                } else {
                    packet.skip(count * sizeof(Chunk::Edit));
                }
            }
            break;
        }

        case server_packet_type::ENTITY_TYPES: {
            entities.clear();
            while(!packet.isDone()) {
                size_t index;
                glm::usvec3 size;
                glm::vec3 scale;
                uint8_t flags;
                if(packet.read(index) && packet.read(size) && packet.read(scale) && packet.read(flags)) {
                    const bool permanant = (flags & (1 << 0)) != 0;
                    const bool area = (flags & (1 << 1)) != 0;
                    entities.set(index, Entity(size, scale, permanant, area));
                    //MAYBE: just update
                }
            }
            break;
        }

        case server_packet_type::ENTITIES: {
            entities.remove([](size_t, Entity &entity) { entity.instances.clear(); return false; });
            while(!packet.isDone()) {
                size_t entity_idx;
                size_t count;
                if (!(packet.read(entity_idx) && packet.read(count)))
                    break;

                if (auto entity = entities.get(entity_idx)) {
                    if (count == 0)
                        continue;
                    for (size_t i = 0; i < count && !packet.isDone(); i++) {
                        size_t idx;
                        glm::ifvec3 pos;
                        glm::vec3 vel;
                        if (packet.read(idx) && packet.read(pos) && packet.read(vel))
                            entity->instances.set_emplace(idx, Universe::Entity::Instance{pos, vel});
                    }
                    if (dict.has_value()) {
                        if (auto area = std::get_if<Universe::Entity::area_t>(&entity->shape)) {
                            if (area->capacity())
                                continue;

                            const auto scale = glm::lvec3(1) + glm::divide(entity->size);
                            area->reserve(scale.x * scale.y * scale.z);
                        } else if (auto model = std::get_if<Universe::Entity::model_t>(&entity->shape)) {
                            if (model->capacity())
                                continue;

                            model->reserve(8);
                        }
                        peer.send(PacketWriter::Of(client_packet_type::MISSING_ENTITY, entity_idx));
                    }
                } else {
                    LOG_W("Unknown entity type " << entity_idx);
                    packet.skip(count * (sizeof(size_t) + sizeof(glm::ifvec3) + sizeof(glm::vec3)));
                }
            }
            entities.remove([&](size_t id, const Entity &entity) {
                if (!entity.permanant && entity.instances.empty()) {
                    contouring->onEntityUnload(id);
                    return true;
                }
                return false;
            });
            break;
        }

        case server_packet_type::ENTITY_SHAPE: {
            assert(dict.has_value());

            size_t id;
            bool is_area;
            if (!packet.read(id) || !packet.read(is_area))
                break;

            if (auto it = entities.get(id)) {
                if (is_area) {
                    auto area = std::get_if<Universe::Entity::area_t>(&it->shape);
                    assert(area);
                    area->clear();
                    while (!packet.isDone()) {
                        std::vector<char> buffer;
                        if(auto err = dict.value().decompress(packet.readSizedPart(), buffer)) {
                            LOG_E("Corrupted entity chunk packet " << err.value());
                            break;
                        }
                        data::vec_istream idata(buffer);
                        std::istream iss(&idata);
                        area->push_back(std::make_shared<Chunk>(iss));
                    }
                    auto scale = glm::lvec3(1) + glm::divide(it->size);
                    if ((long)area->size() != scale.x * scale.y * scale.z) {
                        LOG_E("Corrupted entity area " << area->size() << "!=" << (scale.x * scale.y * scale.z));
                        break;
                    }
                    contouring->onEntityLoad(id, scale, *area);
                } else {
                    auto model = std::get_if<Universe::Entity::model_t>(&it->shape);
                    assert(model);
                    auto remain = packet.readRemaining();
                    model->resize(remain.size());
                    //MAYBE: avoid storing model
                    memcpy(model->data(), remain.data(), remain.size());
                    data::vec_istream idata(*model);
                    std::istream iss(&idata);
                    contouring->onEntityLoad(id, iss);
                }
            } else {
                LOG_W("Shape for unknown entity type " << id);
            }
            break;
        }

        default:
            LOG_W("Bad packet from server");
            break;
    }
    return true;
}

void DistantUniverse::emit(const world::action::packet &action) {
    if(const auto move = std::get_if<world::action::Move>(&action)) {
        peer.send(net::PacketWriter::Of(net::client_packet_type::MOVE, move->pos), net::client::queue::MOVE, 1);
    } else if(const auto message = std::get_if<world::action::Message>(&action)) {
        peer.send(net::PacketWriter::Of(net::client_packet_type::MESSAGE, message->text.data(), message->text.size()));
    } else if(const auto fill = std::get_if<world::action::FillShape>(&action)) {
        peer.send(net::PacketWriter::Of(net::client_packet_type::FILL_SHAPE, *fill));
        if (options.editPrediction) {
            ZoneScopedN("Fill");
            const auto keepDelay = 5 + (peer.getRTT() / 20000.f); // 5s + 50RTT
            world::iterator::ApplySplit<Chunk>(areas, *fill, floodfillLimit, [&](std::shared_ptr<Chunk> &ck, chunk_pos, chunk_voxel_idx idx,
                Voxel, Voxel next, float delay) { ck->addFutureEdit(Chunk::Edit{next, keepDelay - delay * 2, idx}, delay); });
        }
    } else {
        LOG_W("Bad action " << action.index());
    }
}

Universe::ray_result DistantUniverse::raycast(const geometry::Ray &ray) const {
    return Raycast(ray, areas);
}

bool DistantUniverse::isAreaFree(const area_<voxel_pos> &pos, const geometry::Shape shape, const uint16_t radius) const {
    if (const auto it = areas.find(pos.first); it != areas.end()) {
        const auto center = pos.second + it->second->getOffset().as_voxel();
        return !entities.contains([&](size_t, const Entity &entity) {
            return entity.instances.contains([&](size_t, const Universe::Entity::Instance &inst) {
                return geometry::InShape(shape, center, radius, inst.pos.as_voxel(), entity.size);
            });
        });
    } else
        return false;
}

void DistantUniverse::getEntitiesModels(const std::function<void(size_t, const std::vector<glm::mat4>&)>& call,
    const std::optional<geometry::Frustum>& frustum, const glm::llvec3& offset, int density)
{
    std::vector<glm::mat4> mats;
    entities.iter([&](size_t eId, const Entity &entity) {
        entity.instances.iter([&](size_t iId, const Universe::Entity::Instance &inst) {
            if (eId == PLAYER_ENTITY_ID.index && iId == playerId)
                return;

            glm::mat4 tmp;
            if (contouring::Abstract::CullEntity(tmp, entity.size, entity.scale, inst.pos, frustum, offset, density))
                mats.push_back(tmp);
        });
        if(!mats.empty()) {
            call(eId, mats);
            mats.resize(0);
        }
    });
}