#include "Memory.hpp"

#include <optional>

using namespace world::server;

#define REMOVE_CORRUPTED 1
#define LAZYNESS 8

MemoryRegion::MemoryRegion(const std::string &folderPath, const area_<region_pos> &pos) {
    path = folderPath + '/' + std::to_string(pos.first.index) + '.' + std::to_string(pos.second.x) + '.' +
        std::to_string(pos.second.y) + '.' + std::to_string(pos.second.z) + ".map";

    load();
}
MemoryRegion::~MemoryRegion() {
    if(!content.empty())
        save(changed);
}

void MemoryRegion::load() {
    std::unique_lock lock(mutex);

    std::ifstream file;
    file.open(path, std::ios::in | std::ios::binary);
    if(!file.good()) {
        return;
    }

    // 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::unique_ptr<std::vector<char>> data = nullptr;
        if (size > 0) {
            data = std::make_unique<std::vector<char>>();
            data->resize(size);
            file.read(data->data(), data->size());
        }
        if(!content.emplace(pos, node(avg, std::move(data))).second) {
            LOG_E("Duplicated chunk: " << path << ":" << (int)pos.x << "." << (int)pos.y << "." << (int)pos.z);
        }
        file.peek();
    }

    if(file.bad()) {
        LOG_E("Region corrupted read " << path);
    }
    assert(content.size() == chunkCount);
    file.close();
}
bool MemoryRegion::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 == nullptr)
        return false;

    if(auto err = ctx.decompress(*it->second.data.get(), 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 MemoryRegion::write(const region_chunk_pos& pos, const zstd::write_ctx& ctx, const std::string_view& in, const std::optional<world::Voxel>& avg) {
    std::unique_ptr<std::vector<char>> buffer = nullptr;
    if (!in.empty()) {
        buffer = std::make_unique<std::vector<char>>();
        if (auto err = ctx.compress(in, *buffer.get())) {
            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()) {
            if (buffer == nullptr) {
                buffer = std::move(it->second.data);
            }
            content.erase(it);
        }
        content.emplace(pos, node(avg, std::move(buffer)));
        changed = true;
    }
    save(false);
}

void MemoryRegion::save(bool force) {
    if(!force && rand() % LAZYNESS == 0)
        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 == nullptr)
            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);
        }
    }

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