#include "DistantUniverse.hpp"

#include <Tracy.hpp>

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

#undef LOG_PREFIX
#define LOG_PREFIX "Client: "

using namespace world::client;

DistantUniverse::DistantUniverse(const Universe::login& login, const connection& ct, const Universe::options& opt, const std::string& contouring):
    Universe(contouring), options(opt), serverDistance(0),
    peer(ct, [&](const memory::read_view& buf, net::PacketFlags flags){ return onPacket(buf, flags); })
{
    {
        auto packet = memory::write_buffer(net::client_packet_type::HELLO, sizeof(uint8_t) + (login.name.length()+1) + login.token.length());
        packet.write<uint8_t>(login.token.empty() ? 0 : 64);
        packet.write(login.name.c_str(), login.name.length()+1);
        if (!login.token.empty())
            packet.write(login.token.data(), login.token.size());

        peer.send(packet.finish());
    }
}
DistantUniverse::~DistantUniverse() { }

void DistantUniverse::update(voxel_pos pos, float deltaTime) {
    { //NOTE: triggers onPacket
        ZoneScopedN("Pull");
        peer.pull();
    }

    if (options.movePrediction) {
        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
        });
    }

    const auto notifyTick = contouringNotifier.mustNotify(pos, deltaTime);

    { // Update alive areas
        ZoneScopedN("World");
        auto rng = std::mt19937(std::rand());
        const auto contouringThreshold = UINT32_MAX / (1 + contouring->getQueueSize());
        for (auto& ref: elements.areas) {
            ZoneScopedN("Area");
            const area_id id = elements.withFlag(ref).val;
            const auto node = elements.findArea(id);
            const chunk_pos areaOff = glm::divide((lpivot)glm::conjugate(node->absolute.rotation) * ((world_pos)pos - node->absolute.position));
            auto &chunks = node->get()->setChunks();
            if (glm::length2(areaOff) <= glm::pow2<glm::ll>(options.keepDistance + chunks.getRadius())) {
                ZoneScopedN("Alive");
                for (auto it_c = chunks.begin(); it_c != chunks.end();) {
                    const cell_pos chunkPos = node->absolute.computeChild(glm::multiply(it_c->first));
                    const auto chunkDist = glm::length2(glm::divide(chunkPos - pos));
                    if (chunkDist <= glm::pow2(options.keepDistance)) {
                        if (const auto neighbors = it_c->second->setEdits()->update(deltaTime, rng() < contouringThreshold)) {
                            contouring->onUpdate(area_chunk_pos{id.val, it_c->first}, chunkDist, chunks, neighbors.value());
                        } //MAYBE: if notifyTick
                        ++it_c;
                    } else {
                        it_c = chunks.erase(it_c);
                    }
                }
                //MAYBE: if dist to unloaded chunk < X
                if (mayQueryChunks && notifyTick) { // Request missing chunks
                    ZoneScopedN("Missing");
                    std::vector<chunk_pos> missingChunks;
                    std::vector<region_pos> missingRegions;
                    std::vector<long> missingDist;
                    const auto area = std::static_pointer_cast<Area>(node->get());
                    //TODO: use easy sphere fill
                    const auto queryDistance = getQueryDistance();
                    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 = areaOff + chunk_pos(x, y, z);
                        if (dist2 <= glm::pow2(queryDistance) && chunks.inRange(p)) {
                            if (const auto it = chunks.find(p); it != chunks.end() &&
                                std::static_pointer_cast<Chunk>(it->second)->isTrusted(options.trustMajorant))
                            {
                                contouring->onNotify(area_chunk_pos{id.val, p}, dist2, chunks);
                                continue;
                            }

                            const auto rcPos = glm::split(p);
                            if(auto it_r = area->regionCache.find(rcPos.first); it_r != 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->setEdits().invalidate(geometry::Faces::All);
                                    chunks.emplace(p, ck);
                                    contouring->onNotify(area_chunk_pos{id.val, p}, dist2, chunks);
                                    if (ck->isTrusted(options.trustMajorant))
                                        continue;
                                }
                            }

                            if (!area->regionCache.contains(rcPos.first) &&
                                std::find(missingRegions.begin(), missingRegions.end(), region_pos(rcPos.first)) == missingRegions.end())
                            {
                                missingRegions.push_back(rcPos.first);
                            }
                            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()) {
                        LOG_T("Missing chunks " << missingChunks.size());
                        auto packet = memory::write_buffer(net::client_packet_type::REQUEST_CHUNKS, sizeof(node_id) + missingChunks.size() * sizeof(chunk_pos));
                        packet.write<node_id>(id.val);
                        packet.write(missingChunks.data(), missingChunks.size() * sizeof(chunk_pos));
                        peer.send(packet.finish());
                        mayQueryChunks = false;
                    }
                    if(!missingRegions.empty()) {
                        auto packet = memory::write_buffer(net::client_packet_type::REQUEST_REGIONS, sizeof(area_id) + missingRegions.size() * sizeof(region_pos));
                        packet.write<area_id>(id.val);
                        packet.write(missingRegions.data(), missingRegions.size() * sizeof(region_pos));
                        peer.send(packet.finish());
                    }
                }
            } else if(notifyTick) {
                chunks.clear();
                chunks.rehash(0);
                auto& regions = std::static_pointer_cast<Area>(node->get())->regionCache;
                regions.clear();
                regions.rehash(0);
            }
        }
    }
    //TODO: remove out of range (instances, parts, null node)
    if (notifyTick) {
        const auto queryDist2 = glm::pow2(getQueryDistance());
        for(const auto& ref: elements.parts) {
            const part_id id = elements.withFlag(ref).val;
            const auto node = elements.findPart(id);
            const auto dist = glm::length2(glm::divide((cell_pos)node->absolute.position - pos));
            if (dist <= glm::pow2(options.keepDistance)) {
                const auto part = node->get();
                if (part->loaded()) {
                    contouring->onNotify(id, dist, part->chunkSize(), part->chunks, false);
                } else if (dist <= queryDist2) {
                    const auto fullCount = part->chunkCount();
                    const auto &chunks = part->chunks;
                    size_t size = 0;
                    for (size_t i = 0; i < fullCount; i++) {
                        if (i >= chunks.size() || chunks[i] == nullptr)
                            size++;
                    }
                    LOG_T("Missing chunks in part " << size);
                    auto packet = memory::write_buffer(net::client_packet_type::REQUEST_CHUNKS, sizeof(node_id) + size * sizeof(chunk_pos));
                    packet.write<node_id>(id.val);
                    for (size_t i = 0; i < fullCount; i++) {
                        if (i >= chunks.size() || chunks[i] == nullptr) {
                            packet.write<chunk_pos>(part->getPos(i));
                        }
                    }
                    peer.send(packet.finish());
                }
            } else {
                auto &chunks = node->get()->chunks;
                chunks.clear();
                chunks.shrink_to_fit();
            }
        }
        //TODO: remove temporary on last instance removal
        elements.models.iter([&] (const model_id& id, const Model& model) {
            if (!model.instances.empty()) {
                io::vec_istream idata(model.dt);
                std::istream iss(&idata);
                contouring->onLoad(id, iss);
            }
        });
    }

    contouring->update(pos, elements);
}

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

    //MAYBE: client::Serializer

    auto packet = memory::read_buffer(buf, nullptr);
    server_packet_type type;
    if(!packet.read(type)) {
        LOG_D("Empty packet");
        return true;
    }
    switch (type) {
        case server_packet_type::QUIT: {
            uint16_t reason = 42;
            packet.read(reason);
            LOG_E("Disconnected from server (" << reason << ')');
            peer.disconnect();
            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 = memory::write_buffer(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: {
            relative_transform pos;
            if (packet.read(playerId) && packet.read(pos)) {
                if (!elements.withFlag(playerId)) {
                    LOG_W("FIXME: store if instance not reserved yet");
                }
                onTeleport(pos, getAbsolute(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::REGION: {
            ZoneScopedN("Region");
            area_region_ref pos(area_ref(), region_pos(0));
            if(!packet.read(pos))
                break;

            const auto id = elements.nodes.with_flag(pos.area.val);
            if (!Elements::Is<Elements::Type::Area>(id))
                break;

            const auto node = NodeOf<Area>::Make(elements.nodes.directly_at(id));
            if (!node) {
                LOG_W("Region area not found " << pos.area.val);
                break;
            }
            auto &regions = node->get()->regionCache;
            auto it_r = regions.find(pos.region);
            if (it_r == regions.end()) {
                it_r = regions.emplace(pos.region, 0).first;
            }
            // MAYBE: create virtual Chunk without voxels of single material
            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);
            }
            break;
        }

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

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

            node_chunk_ref pos{node_ref(), chunk_pos(0)};
            if(!packet.read(pos))
                break;

            const auto id = elements.nodes.with_flag(pos.node.val);
            const auto node = elements.nodes.directly_at(id);
            if (!node) {
                LOG_W("node not found " << pos.node.val);
                break;
            }

            const auto readChunk = [&] {
                std::vector<char> buffer;
                if (auto err = dict.value().decompress(packet.readRemaining(), buffer)) {
                    LOG_E("Corrupted chunk packet " << err.value());
                    return std::shared_ptr<Chunk>(nullptr);
                }
                io::vec_istream idata(buffer);
                std::istream iss(&idata);
                return std::make_shared<Chunk>(iss);
            };

            switch (Elements::GetType(id)) {
            case Elements::Type::Area: {
                const auto area = NodeOf<Area>::Make(node)->get();
                if(!area->getChunks().inRange(pos.chunk)) {
                    LOG_W("Chunk out of area " << pos.node.val);
                    break;
                }
                if (auto ck = readChunk()) {
                    area->setChunks().insert_or_assign(pos.chunk, ck);
                }
                break;
            }
            case Elements::Type::Part: {
                const auto part = NodeOf<Part>::Make(node)->get();
                if(!part->inRange(pos.chunk)) {
                    LOG_W("Chunk out of part " << pos.node.val);
                    break;
                }
                if (auto ck = readChunk()) {
                    part->allocate();
                    part->emplace(pos.chunk, ck);
                }
                break;
            }

            default:
                LOG_W("Not chunk holding node type");
                break;
            }
            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;

            const auto anchor_id = elements.withFlag(fill->pos.node);
            const auto node = elements.findNode(anchor_id);
            if (!node)
                break;

            //MAYBE: make it void after air propagation setup
            const auto cleanVoxel = Voxel(0, Voxel::DENSITY_MAX);
            switch (Elements::GetType(anchor_id)) {
                case Elements::Type::Area: {
                    const auto& chunks = NodeOf<Area>::Make(node)->get()->getChunks();
                    world::iterator::ApplyEraseSplitted<Chunk>(chunks, *fill, floodfillLimit, cleanVoxel,
                        [&](std::shared_ptr<Chunk>& ck, const chunk_pos&, chunk_voxel_idx idx, Voxel/*prev*/, Voxel next, float delay) {
                            return ck->setEdits().apply(ChunkEdits::Edit{next, delay, idx});
                    });
                    break;
                }
                /* TODO: case Elements::Type::Part: {
                    const auto& chunks = *NodeOf<Part>::Make(node)->get()->chunks;
                    world::iterator::ApplyEraseSplitted<Chunk>(chunks, *fill, floodfillLimit, cleanVoxel,
                        [&](std::shared_ptr<Chunk>& ck, const chunk_pos&, chunk_voxel_idx idx, Voxel prev, Voxel next, float delay) {
                            return ck->replace(ChunkEdits::Edit{next, delay, idx});
                    });
                    break;
                }*/
                default:
                    LOG_E("Unhandled fill target type " << (int)Elements::GetType(anchor_id));
                    break;
            }
            break;
        }
        case server_packet_type::RAW_EDITS: {
            ZoneScopedN("Raw Edits");
            node_ref ref;
            if(!packet.read(ref))
                break;

            const auto id = elements.nodes.with_flag(ref.val);
            const auto node = elements.nodes.directly_at(id);
            if (!node) {
                LOG_W("node not found " << ref.val);
                break;
            }

            const auto readEdits = [&](std::function<std::shared_ptr<EdittableChunk>(const chunk_pos&)> findInRange) {
                while(!packet.isDone()) {
                    chunk_pos pos = chunk_pos(INT_MAX);
                    packet.read(pos);
                    chunk_voxel_idx count = 0;
                    packet.read(count);
                    if (auto chunk = findInRange(pos)) {
                        const auto edits = chunk->setEdits();
                        for (auto i = 0; i < count && !packet.isDone(); i++) {
                            ChunkEdits::Edit edit;
                            packet.read(edit);
                            edits->apply(edit);
                        }
                    } else {
                        packet.skip(count * sizeof(ChunkEdits::Edit));
                    }
                }
            };

            switch (Elements::GetType(id)) {
            case Elements::Type::Area: {
                const auto area = NodeOf<Area>::Make(node)->get();
                const auto& chunks = area->setChunks();
                readEdits([&](const world::chunk_pos &pos) {
                    return chunks.findInRange(pos);
                });
                break;
            }
            case Elements::Type::Part: {
                const auto part = NodeOf<Part>::Make(node)->get();
                readEdits([&](const world::chunk_pos &pos) {
                    return part->findInRange(pos);
                });
                break;
            }

            default:
                LOG_W("Not chunk holding node type");
                break;
            }
            break;
        }

        case server_packet_type::ENTITIES: {
            while(!packet.isDone()) {
                node_id id;
                if (!packet.read(id))
                    break;

                const auto added = Elements::Has<Elements::Flag::Added>(id);
                if (Elements::Has<Elements::Flag::Moved>(id) || added) {
                    MyElements::start_point state = MyElements::start_point(generational::id(), transform());
                    if (!packet.read(state))
                        break;

                    if (added) {
                        switch (Elements::GetType(id)) {
                            case Elements::Type::Area: {
                                Area::params params;
                                if (packet.read(params)) {
                                    if (elements.withFlag(id)) {
                                        LOG_D("Added entity already exists " << id.val.val);
                                        break;
                                    }
                                    id = elements.createAt(id, state, std::make_shared<Area>(params));
                                    elements.areas.push_back(id);
                                }
                                break;
                            }
                            case Elements::Type::Part: {
                                glm::usvec3 size;
                                if (packet.read(size)) {
                                    if (elements.withFlag(id)) {
                                        LOG_D("Added entity already exists " << id.val.val);
                                        break;
                                    }
                                    id = elements.createAt(id, state, std::make_shared<Part>(size));
                                    elements.parts.push_back(id);
                                }
                                break;
                            }
                            case Elements::Type::Instance: {
                                model_id mid;
                                if (packet.read(mid)) {
                                    const auto model = elements.models.directly_at(mid);
                                    if (!model) {
                                        LOG_E("FIXME: must query model");
                                        break;
                                    }
                                    if (elements.withFlag(id)) {
                                        LOG_D("Added entity already exists " << id.val.val);
                                        break;
                                    }
                                    id = elements.createAt(id, state, model->origin ? model->origin : std::make_shared<Instance>(mid));
                                    model->instances.push_back(id.val.index());
                                }
                                break;
                            }

                            default:
                                LOG_W("Unhandled entity type " << id.val.val);
                                break;
                        }
                    } else if (elements.withFlag(id)) {
                        elements.move(id, state);
                    } else {
                        LOG_W("Must query content " << id.val.val);
                    }
                } else if (Elements::Has<Elements::Flag::Removed>(id)) {
                    elements.remove(id);
                } else {
                    LOG_W("Corrupted entity with id " << id.val.val);
                    break;
                }
            }
            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(memory::read_buffer::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(memory::read_buffer::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(memory::read_buffer::Of(net::client_packet_type::FILL_SHAPE, *fill));
        if (options.editPrediction) {
            ZoneScopedN("Fill");
            const auto keepDelay = 5 + (peer.getRTT() / 20000.f); // 5s + 50RTT
            const auto anchor_id = elements.withFlag(fill->pos.node);
            const auto node = elements.findNode(anchor_id);
            if (!node)
                return;

            //MAYBE: make it void after air propagation setup
            const auto cleanVoxel = Voxel(0, Voxel::DENSITY_MAX);
            switch (Elements::GetType(anchor_id)) {
                case Elements::Type::Area: {
                    const auto& chunks = NodeOf<Area>::Make(node)->get()->getChunks();
                    world::iterator::ApplyEraseSplitted<Chunk>(chunks, *fill, floodfillLimit, cleanVoxel,
                        [&](std::shared_ptr<Chunk>& ck, const chunk_pos&, chunk_voxel_idx idx, Voxel/*prev*/, Voxel next, float delay) {
                            return ck->setEdits().addFutureEdit(ChunkEdits::Edit{next, keepDelay - delay * 2, idx}, delay);
                    });
                    break;
                }
                default:
                    LOG_E("Unhandled fill target type " << (int)Elements::GetType(anchor_id));
                    break;
            }
        }
    } else {
        LOG_W("Bad action " << action.index());
    }
}

DistantUniverse::MyElements::MyElements(): EdittableElements(true) { }

world::node_id DistantUniverse::MyElements::createAt(node_id id, const start_point& sp, const std::shared_ptr<world::Element>& el) {
    return EdittableElements::createAt(id, sp, el);
}