#include "./Region.hpp"

#include <optional>
#include <fstream>

#undef LOG_PREFIX
#define LOG_PREFIX "Server: "

using namespace world::server;

#define REMOVE_CORRUPTED 1
Region::save_throttler::save_throttler(): nextSaveExponent(2) { }

Region::Region(const std::string &folderPath, const area_region_ref &id) {
    path = folderPath + '/' + std::to_string(id.area.val) + '.' + std::to_string(id.region.x) + '.' +
        std::to_string(id.region.y) + '.' + std::to_string(id.region.z) + ".map";

    load();
}
Region::~Region() {
    if(!content.empty())
        save(saveThrottler.isUnsaved());
}

size_t Region::Read(const std::string &path, const std::function<void(const region_chunk_pos &, const std::optional<Voxel> &, std::vector<char> &&)>& out) {
    std::ifstream file;
    file.open(path, std::ios::in | std::ios::binary);
    if(!file.good()) {
        return 0;
    }

    // Read header
    uint16_t chunkCount; //NOTE: pretty useless
    file.read(reinterpret_cast<char *>(&chunkCount), sizeof(chunkCount));

    while (!file.eof()) {
        // Read pos
        region_chunk_pos pos;
        file.read(reinterpret_cast<char *>(&pos.x), sizeof(region_chunk_pos::value_type));
        file.read(reinterpret_cast<char *>(&pos.y), sizeof(region_chunk_pos::value_type));
        file.read(reinterpret_cast<char *>(&pos.z), sizeof(region_chunk_pos::value_type));

        // Read average if present
        std::optional<Voxel> avg;
        Flags flags = Flags::ZERO;
        file.read(reinterpret_cast<char *>(&flags), 1);
        if (flags & Flags::HAS_AVERAGE) {
            Voxel v;
            file.read(reinterpret_cast<char *>(&v), sizeof(v));
            avg = v;
        }

        // Read size
        uint16_t size = 0;
        if (!(flags & Flags::EMPTY)) {
            file.read(reinterpret_cast<char *>(&size), sizeof(size));
        }

        // Read content
        std::vector<char> data;
        if (size > 0) {
            data.resize(size);
            file.read(data.data(), data.size());
        }
        out(pos, avg, std::move(data));
        file.peek();
    }

    if(file.bad()) {
        LOG_E("Region corrupted read " << path);
    }
    file.close();
    return chunkCount;
}

void Region::Read(const std::string& folder, const part_ref& id, const zstd::read_ctx &ctx, const std::function<void(const region_chunk_pos&, const std::vector<char>&)>& out) {
    const auto path = folder + '/' + std::to_string(id.val) + ".map";
    std::vector<char> buffer;
    Read(path, [&](const region_chunk_pos &pos, const std::optional<Voxel> &, const std::vector<char> &data) {
        if(auto err = ctx.decompress(data, buffer)) {
            LOG_E("Corrupted region chunk: " << path << ":" << (int)pos.x << "." << (int)pos.y << "." << (int)pos.z << " "
                << err.value());
        } else {
            out(pos, buffer);
        }
    });
}
void Region::Write(const std::string& folder, const part_ref& id, const zstd::write_ctx &ctx, size_t size, const std::function<std::pair<region_chunk_pos, std::string>()>& next) {
    const auto path = folder + '/' + std::to_string(id.val) + ".map";

    std::ofstream file(path, std::ios::out | std::ios::binary);
    if (!file.good()) {
        LOG_E("Corrupted region path: " << path);
        return;
    }

    { // Write header
        uint16_t size = (uint16_t)size;
        file.write(reinterpret_cast<char *>(&size), sizeof(size));
    }

    std::vector<char> buffer;
    for (size_t i = 0; i < size; i++) {
        const auto out = next();
        if (out.second.empty()) {
            buffer.clear();
        } else {
            if (auto err = ctx.compress(out.second, buffer)) {
                LOG_E("Corrupted chunk save: " << path << ":" << (int)out.first.x << "." << (int)out.first.y << "." << (int)out.first.z << " "
                    << err.value());
                continue;
            }
        }

        { // Write pos
            region_chunk_pos pos = out.first;
            file.write(reinterpret_cast<char *>(&pos.x), sizeof(region_chunk_pos::value_type));
            file.write(reinterpret_cast<char *>(&pos.y), sizeof(region_chunk_pos::value_type));
            file.write(reinterpret_cast<char *>(&pos.z), sizeof(region_chunk_pos::value_type));
        }

        // Write flags
        Flags flags = Flags::ZERO;
        if (buffer.empty())
            flags = (Flags)(flags | Flags::EMPTY);
        file.write(reinterpret_cast<const char *>(&flags), sizeof(flags));

        if (!(flags & Flags::EMPTY)) {
            assert(buffer.size() < USHRT_MAX);
            auto size = (uint16_t)buffer.size();

            // Write size
            file.write(reinterpret_cast<char *>(&size), sizeof(size));

            // Write content
            file.write(buffer.data(), size);
        }
    }

    if (!file.good()) {
        LOG_E("Region corrupted write " << path);
    }
    file.close();
}

void Region::load() {
    std::unique_lock lock(mutex);
    [[maybe_unused]]
    const auto chunkCount = Read(path, [&](const region_chunk_pos &pos, const std::optional<Voxel> &avg, std::vector<char> &&data) {
        if(!content.emplace(pos, node(avg, std::move(data))).second) {
            LOG_E("Duplicated chunk: " << path << ":" << (int)pos.x << "." << (int)pos.y << "." << (int)pos.z);
        }
    });
    if (content.size() != chunkCount)
        LOG_E("Probably lost some chunks " << content.size() << "/" << chunkCount);
}
bool Region::read(const region_chunk_pos& pos, const zstd::read_ctx& ctx, std::vector<char>& out) {
    std::shared_lock lock(mutex);

    const auto it = content.find(pos);
    if (it == content.end() || it->second.data.empty())
        return false;

    if(auto err = ctx.decompress(it->second.data, out)) {
        LOG_E("Corrupted region chunk: " << path << ":" << (int)pos.x << "." << (int)pos.y << "." << (int)pos.z << " "
            << err.value());
#ifdef REMOVE_CORRUPTED
            LOG_W("Removing");
            lock.unlock();
            {
                std::unique_lock ulock(mutex);
                content.erase(it);
            }
            save(true);
#endif
        return false;
    }
    return true;
}
void Region::write(const region_chunk_pos& pos, const zstd::write_ctx& ctx, const std::string_view& in, const std::optional<world::Voxel>& avg) {
    std::vector<char> buffer;
    if (!in.empty()) {
        if (auto err = ctx.compress(in, buffer)) {
            LOG_E("Corrupted chunk save: " << path << ":" << (int)pos.x << "." << (int)pos.y << "." << (int)pos.z << " "
                << err.value());
            return;
        }
    }
    {
        std::unique_lock lock(mutex);

        // Save buffer
        const auto it = content.find(pos);
        if (it != content.end()) {
            it->second.average = avg;
            it->second.data = std::move(buffer);
        } else {
            content.emplace(pos, node(avg, std::move(buffer)));
        }
        saveThrottler.change();
    }
    save(false);
}

void Region::save(bool force) {
    if(!(force || saveThrottler.mustSave()))
        return;

    std::unique_lock lock(mutex);

    std::ofstream file(path, std::ios::out | std::ios::binary);
    if (!file.good()) {
        LOG_E("Corrupted region path: " << path);
        return;
    }

    { // Write header
        uint16_t size = (uint16_t)content.size();
        file.write(reinterpret_cast<char *>(&size), sizeof(size));
    }

    for(const auto& chunk: content) {
        { // Write pos
            region_chunk_pos pos = chunk.first;
            file.write(reinterpret_cast<char *>(&pos.x), sizeof(region_chunk_pos::value_type));
            file.write(reinterpret_cast<char *>(&pos.y), sizeof(region_chunk_pos::value_type));
            file.write(reinterpret_cast<char *>(&pos.z), sizeof(region_chunk_pos::value_type));
        }

        // Write average if present
        Flags flags = Flags::ZERO;
        if (chunk.second.average.has_value())
            flags = (Flags)(flags | Flags::HAS_AVERAGE);
        if (chunk.second.data.empty())
            flags = (Flags)(flags | Flags::EMPTY);
        file.write(reinterpret_cast<const char *>(&flags), sizeof(flags));
        if (flags & Flags::HAS_AVERAGE) {
            Voxel v = chunk.second.average.value();
            file.write(reinterpret_cast<char *>(&v), sizeof(v));
        }

        if (!(flags & Flags::EMPTY)) {
            assert(chunk.second.data.size() < USHRT_MAX);
            auto size = (uint16_t)chunk.second.data.size();
            const auto out = chunk.second.data.data();

            // Write size
            file.write(reinterpret_cast<char *>(&size), sizeof(size));

            // Write content
            file.write(out, size);
        }
    }

    file.sync_with_stdio(true);
    file.flush();
    if (!file.good()) {
        LOG_E("Region corrupted write " << path);
        file.close();
        return;
    }
    file.close();
    saveThrottler.save();
}