#include <MurmurHash3/MurmurHash3.hpp>
#include <cstdint>
#include <cstdio>
#include <chrono>
#include <vector>
#include <zlib.h>
#include <AtomicDataStore/AtomicDataStore.hpp>
#include <SHA256/SHA256.hpp>
#include <Log/Log.hpp>
#include <Utilities/Utilities.hpp>

#ifdef DEBUG_BUILD
#include <filesystem>
#include <fstream>
#endif

using namespace std::chrono_literals;

void AtomicDataStore::IndexBlock::updateHash()
{
    SHA256 hasher;
    hasher.feed(this, offsetof(AtomicDataStore::IndexBlock, hash));
    hasher.digest(hash);
}

bool AtomicDataStore::IndexBlock::isValid() const
{
    SHA256 hasher;
    hasher.feed(this, offsetof(AtomicDataStore::IndexBlock, hash));

    SHA256::Digest correctHash;
    hasher.digest(correctHash);

    if (memcmp(hash.bytes, correctHash.bytes, sizeof(hash)) != 0)
    {
        return false;
    }

    return true;
}

bool AtomicDataStore::readIndexBlock(IndexBlock &output, size_t &outNextIndex)
{
    IndexBlock indexBlocks[2];
    if (!file.read(sizeof(MainHeader), indexBlocks, sizeof(indexBlocks)))
    {
        return false;
    }

    bool aValid = indexBlocks[0].isValid();
    bool bValid = indexBlocks[1].isValid();

    if (!aValid && !bValid)
    {
        return false;
    }

    if (!aValid)
    {
        output = indexBlocks[1];
        outNextIndex = 0;
        return true;
    }

    if (!bValid)
    {
        output = indexBlocks[0];
        outNextIndex = 1;
        return true;
    }

    if (indexBlocks[0].generation - indexBlocks[1].generation == 1)
    {
        output = indexBlocks[0];
        outNextIndex = 1;
        return true;
    }

    if (indexBlocks[1].generation - indexBlocks[0].generation == 1)
    {
        output = indexBlocks[1];
        outNextIndex = 0;
        return true;
    }

    return false;
}

uint64_t AtomicDataStore::allocStorageRegion(size_t size)
{
    for (size_t i = 0; i < freeRegionTable.size(); i++)
    {
        auto &region = freeRegionTable[i];
        if (region.size >= size)
        {
            uint64_t result = region.basePos;
            region.basePos += size;
            region.size -= size;

            if (region.size == 0)
            {
                freeRegionTable.erase(freeRegionTable.begin() + i);
            }

            return result;
        }
    }

    crash("Failed to find any storage region in this world file!");
}

AtomicDataStore::HashTable& AtomicDataStore::getMutableHashTable(uint64_t pos, uint64_t compressedSize)
{
    auto it = pendingHashTables.find(pos);
    if (it != pendingHashTables.end())
    {
        return it->second;
    }

    HashTable &newTable = pendingHashTables[pos];

    std::vector<uint8_t> compressedHashTable(compressedSize);
    if (!file.read(pos, compressedHashTable.data(), compressedHashTable.size()))
    {
        crash(formatString("Failed to read a hash table, pos=0x$+$+", std::hex, pos));
    }

    std::vector<uint8_t> uncompressedHashTable;
    if (!uncompressBytes(compressedHashTable, uncompressedHashTable))
    {
        crash("Failed to uncompress a hash table.");
    }

    if (uncompressedHashTable.size() != sizeof(HashTable))
    {
        crash("An uncompressed hash table is the wrong size.");
    }
    
    memcpy(&newTable, uncompressedHashTable.data(), uncompressedHashTable.size());

    return newTable;
}

void AtomicDataStore::encodeFileRecur(
    MurmurHash3 &hasher,
    const std::string &filename,
    const std::vector<uint8_t> &payload,
    uint64_t hashtabPos,
    uint64_t hashtabCompressedSize,
    int depth
)
{
    HashTable &hashtab = getMutableHashTable(hashtabPos, hashtabCompressedSize);
    hasher.feed(filename.data(), filename.size());

    HashTableEntry &entry = hashtab.entries[hasher.digest() % numHashTableEntries];

    if (entry.type == HashTableEntryType::Unmapped)
    {
        encodeFileDescriptor(entry, filename, payload);
    }
    else if (entry.type == HashTableEntryType::NestedHashTable)
    {
        encodeFileRecur(hasher, filename, payload, entry.pos, entry.dataSize, depth + 1);
    }
    else if (entry.type == HashTableEntryType::FileDescriptor)
    {
        // Check if it's the same file name.
        std::vector<char> fileNameChars;
        fileNameChars.resize(entry.fileNameLen);

        if (!file.read(entry.pos, fileNameChars.data(), entry.fileNameLen))
        {
            crash("Failed to read the world file.");
        }

        std::string existingFileName(fileNameChars.data(), fileNameChars.size());
        if (existingFileName == filename)
        {
            // File already exists; set new contents for it.
            regionsToRelease.push_back({entry.pos, entry.fileNameLen + entry.dataSize});
            encodeFileDescriptor(entry, filename, payload);
            return;
        }

        // Conflict; nest deeper, or ignore if we have no content.
        if (payload.empty())
        {
            return;
        }

        MurmurHash3 nestedHasher;
        for (int i = 0; i <= depth + 1; i++)
        {
            nestedHasher.feed(existingFileName.data(), existingFileName.size());
        }

        uint64_t newHashTableId;
        HashTable &nestedHashTable = createMutableHashTable(newHashTableId);

        nestedHashTable.entries[nestedHasher.digest() % numHashTableEntries] = entry;

        entry = {};
        entry.type = HashTableEntryType::NestedHashTable;
        entry.pos = newHashTableId;

        encodeFileRecur(hasher, filename, payload, newHashTableId, 0, depth + 1);
    }
    else
    {
        crash("Unexpected hash table entry type.");
    }
}

void AtomicDataStore::writeHashTable(uint64_t &pos, uint32_t &size)
{
    if (!pendingHashTables.contains(pos))
    {
        // No updates.
        return;
    }

    HashTable &table = getMutableHashTable(pos, size);

    for (int i = 0; i < numHashTableEntries; i++)
    {
        HashTableEntry &entry = table.entries[i];
        if (entry.type != HashTableEntryType::NestedHashTable)
        {
            continue;
        }

        writeHashTable(entry.pos, entry.dataSize);
    }

    if ((pos & tempPosBit) == 0)
    {
        if (size == 0)
        {
            crash("AtomicDataStore::writeHashTable: original size is 0");
        }

        regionsToRelease.push_back({pos, size});
    }

    std::vector<uint8_t> uncompressedData((uint8_t*) &table, (uint8_t*) &table + sizeof(table));
    std::vector<uint8_t> compressedData;
    compressBytes(uncompressedData, compressedData, 9);

    if (compressedData.empty())
    {
        crash("Sanity check failed: size is 0 for compressed data");
    }

    size = compressedData.size();
    pos = allocStorageRegion(size);
    
    if (!file.write(pos, compressedData.data(), size))
    {
        crash("Failed to write a hash table to the world file.");
    }
}

void AtomicDataStore::flushHashTables()
{
    writeHashTable(rootHashTablePos, rootHashTableSize);
    pendingHashTables.clear();
}

void AtomicDataStore::encodeFileDescriptor(
    HashTableEntry &entry,
    const std::string &filename,
    const std::vector<uint8_t> &payload
)
{
    if (payload.empty())
    {
        entry = {};
        return;
    }

    size_t descSize = filename.size() + payload.size();
    uint64_t descPos = allocStorageRegion(descSize);

    if (!file.write(descPos, filename.data(), filename.size()))
    {
        crash(
            formatString(
                "Failed to write a file descriptor header: descPos=0x$+$+, filename=$+",
                std::hex, descPos,
                filename
            )
        );
    }

    if (!file.write(descPos + filename.size(), payload.data(), payload.size()))
    {
        crash(
            formatString(
                "Failed to write a file descriptor payload: descPos=0x$+$+, filename=$+, size=$+",
                std::hex, descPos, filename, payload.size()
            )
        );
    }

    entry.type = HashTableEntryType::FileDescriptor;
    entry.resv = 0;
    entry.fileNameLen = filename.size();
    entry.dataSize = payload.size();
    entry.pos = descPos;
}

AtomicDataStore::HashTable& AtomicDataStore::createMutableHashTable(uint64_t &pos)
{
    pos = nextTempHashTableIndex++;
    return pendingHashTables[pos];
}

void AtomicDataStore::scanHashTable(std::vector<std::string> &paths, uint64_t pos, uint32_t size)
{
    HashTable table;

    std::vector<uint8_t> compressedData(size);
    if (!file.read(pos, compressedData.data(), size))
    {
        crash("Failed to read a hash table.");
    }

    std::vector<uint8_t> uncompressedData;
    if (!uncompressBytes(compressedData, uncompressedData))
    {
        crash("Failed to uncompress a hash table.");
    }

    if (uncompressedData.size() != sizeof(table))
    {
        crash("An uncompressed hash table is the wrong size.");
    }

    memcpy(&table, uncompressedData.data(), uncompressedData.size());

    for (int i = 0; i < numHashTableEntries; i++)
    {
        HashTableEntry &entry = table.entries[i];
        if (entry.type == HashTableEntryType::FileDescriptor)
        {
            std::vector<char> fileNameChars;
            fileNameChars.resize(entry.fileNameLen);

            if (!file.read(entry.pos, fileNameChars.data(), entry.fileNameLen))
            {
                crash("Failed to read the world file.");
            }

            std::string fileName(fileNameChars.data(), fileNameChars.size());
            paths.push_back(fileName);
        }
        else if (entry.type == HashTableEntryType::NestedHashTable)
        {
            scanHashTable(paths, entry.pos, entry.dataSize);
        }
    }
}

bool AtomicDataStore::create(const std::string &path)
{
    if (!file.create(path))
    {
        error = "Failed to create the world file.";
        return false;
    }

    if (!file.lock())
    {
        file.close();
        error = "Failed to lock the world file.";
        return false;
    }

    createMutableHashTable(rootHashTablePos);
    rootHashTableSize = 0;

    nextIndexBlockIndex = 0;
    nextIndexGeneration = 0;

    Region initialRegion{};
    initialRegion.basePos = sizeof(MainHeader) + 2 * sizeof(IndexBlock);
    initialRegion.size = ~initialRegion.basePos + 1;

    freeRegionTable.push_back(initialRegion);
    headerNeedsWriting = true;

    return true;
}

bool AtomicDataStore::open(const std::string &path)
{
    if (!file.open(path))
    {
        error = "Failed to open the world file.";
        return false;
    }

    if (!file.lock())
    {
        file.close();
        error = "The world is open in another program.";
        return false;
    }

    MainHeader mainHeader;
    if (!file.read(0, &mainHeader, sizeof(mainHeader)))
    {
        file.close();
        error = "This is not a Voxelverse world.";
        return false;
    }

    if (memcmp(mainHeader.magic, "\0VVW", 4) != 0)
    {
        file.close();
        error = "This is not a Voxelverse world.";
        return false;
    }

    IndexBlock currentIndexBlock;
    if (!readIndexBlock(currentIndexBlock, nextIndexBlockIndex))
    {
        file.close();
        error = "Failed to read the index block.";
        return false;
    }

    if (currentIndexBlock.numFreeRegions == 0)
    {
        file.close();
        error = "There are no empty regions in this file.";
        return false;
    }

    nextIndexGeneration = currentIndexBlock.generation + 1;

    freeRegionTable.resize(currentIndexBlock.numFreeRegions);
    if (!file.read(currentIndexBlock.freeRegionTablePos, freeRegionTable.data(), sizeof(Region) * freeRegionTable.size()))
    {
        file.close();
        error = "Failed to read the free region table.";
        return false;
    }

    regionsToRelease.push_back({currentIndexBlock.freeRegionTablePos, sizeof(Region) * freeRegionTable.size()});

    while (freeRegionTable.back().size == 0)
    {
        freeRegionTable.resize(freeRegionTable.size() - 1);
    }

    rootHashTablePos = currentIndexBlock.rootHashTablePos;
    rootHashTableSize = currentIndexBlock.rootHashTableSize;

    return true;
}

const std::string& AtomicDataStore::getError() const
{
    return error;
}

void AtomicDataStore::update(const std::unordered_map<std::string, std::vector<uint8_t>> &plan, uint64_t generation)
{
    // No need to do anything if the plan is empty, unless we don't have a header yet.
    if (plan.empty() && !headerNeedsWriting)
    {
        lastFinishedGeneration = generation;
        return;
    }

    std::chrono::steady_clock::time_point writeStart = std::chrono::steady_clock::now();

    // Lock the access mutex.
    std::unique_lock<std::shared_mutex> lock(mutex);

    // Encode the files.
    for (const auto &[filename, data] : plan)
    {
        std::vector<uint8_t> payload;
        if (!data.empty())
        {
            compressBytes(data, payload, 9);
        }

        MurmurHash3 hasher;
        encodeFileRecur(hasher, filename, payload, rootHashTablePos, rootHashTableSize, 0);
    }

    flushHashTables();

    // Make and self-embed the new region table.
    std::vector<Region> updatedRegionTable = freeRegionTable;
    updatedRegionTable.insert(updatedRegionTable.end(), regionsToRelease.begin(), regionsToRelease.end());
    regionsToRelease.clear();

    // Sort the regions by address.
    std::sort(
        updatedRegionTable.begin(),
        updatedRegionTable.end(),
        [](const Region& a, const Region& b)
        {
            return a.basePos < b.basePos;
        }
    );

    // Merge neighboring regions.
    for (size_t i = 0; i < updatedRegionTable.size() - 1; i++)
    {
        Region &thisRegion = updatedRegionTable[i];
        Region &nextRegion = updatedRegionTable[i + 1];

        if (thisRegion.basePos + thisRegion.size == nextRegion.basePos)
        {
            thisRegion.size += nextRegion.size;
            updatedRegionTable.erase(updatedRegionTable.begin() + i + 1);
            i--;
        }
    }

    updatedRegionTable.push_back({0, 0});

    size_t regionTableSize = sizeof(Region) * updatedRegionTable.size();
    uint64_t newRegionTablePos = allocStorageRegion(regionTableSize);

    regionsToRelease.push_back({newRegionTablePos, regionTableSize});

    // Mark the new region as non-free.
    for (auto &region : updatedRegionTable)
    {
        if (region.size == 0)
        {
            continue;
        }
    
        if (region.basePos == newRegionTablePos)
        {
            region.basePos += regionTableSize;
            region.size -= regionTableSize;

            if (region.size == 0)
            {
                region.basePos = 0;
            }

            break;
        }
        else if (region.basePos + region.size == newRegionTablePos + regionTableSize)
        {
            region.size -= regionTableSize;
            break;
        }
        else if (
            region.basePos < newRegionTablePos &&
            (
                region.basePos + region.size > newRegionTablePos + regionTableSize ||
                region.basePos + region.size == 0
            )
        )
        {
            // Break up the block where we overlap.
            auto &lastRegion = updatedRegionTable.back();
            lastRegion.basePos = newRegionTablePos + regionTableSize;
            lastRegion.size = region.basePos + region.size - lastRegion.basePos;

            region.size = newRegionTablePos - region.basePos;
            break;
        }
    }

    std::sort(
        updatedRegionTable.begin(),
        updatedRegionTable.end(),
        [](const Region& a, const Region& b)
        {
            // Place zero size at the end
            if (a.size == 0) return false;
            if (b.size == 0) return true;

            // Otherwise sort by size ascending
            return a.size < b.size;
        }
    );

    freeRegionTable = updatedRegionTable;
    while (freeRegionTable.back().size == 0)
    {
        freeRegionTable.resize(freeRegionTable.size() - 1);
    }

#ifdef DEBUG_BUILD
    for (const Region &region : freeRegionTable)
    {
        if (region.basePos != 0 && region.size == 0)
        {
            std::filesystem::create_directories("plandump");
            for (const auto &[filename, data] : plan)
            {
                std::filesystem::path fullpath = formatString("plandump/$+", filename);
                std::filesystem::create_directories(fullpath.parent_path());

                std::ofstream ofs(fullpath, std::ios::binary);
                ofs.write((char*) data.data(), data.size());
                ofs.close();
            }
            crash("AtomicDataStore: We were about to write a free region table with an invalid region");
        }
    }
#endif

    if (!file.write(newRegionTablePos, updatedRegionTable.data(), regionTableSize))
    {
        crash("Failed to write the new free region table.");
    }

    // Sync the new data and write a new index block.
    if (!file.sync())
    {
        crash("Failed to sync the world data.");
    }

    IndexBlock newIndexBlock{};
    newIndexBlock.rootHashTablePos = rootHashTablePos;
    newIndexBlock.rootHashTableSize = rootHashTableSize;
    newIndexBlock.freeRegionTablePos = newRegionTablePos;
    newIndexBlock.numFreeRegions = updatedRegionTable.size();
    newIndexBlock.generation = nextIndexGeneration++;
    newIndexBlock.updateHash();

    if (!file.write(sizeof(MainHeader) + nextIndexBlockIndex * sizeof(IndexBlock), &newIndexBlock, sizeof(newIndexBlock)))
    {
        crash("Failed to write the new index.");
    }

    if (!file.sync())
    {
        crash("Failed to sync the new index.");
    }

    nextIndexBlockIndex ^= 1;
    memset(&newIndexBlock, 0xFF, sizeof(newIndexBlock));

    if (!file.write(sizeof(MainHeader) + nextIndexBlockIndex * sizeof(IndexBlock), &newIndexBlock, sizeof(newIndexBlock)))
    {
        crash("Failed to wipe the old index.");
    }

    if (headerNeedsWriting)
    {
        if (!file.write(0, "\0VVW", 4))
        {
            crash("Failed to write the main header.");
        }

        headerNeedsWriting = false;
    }

    if (!file.sync())
    {
        crash("Failed to sync the world file.");
    }

    // Mark this plan generation as complete.
    lastFinishedGeneration = generation;

    // Save the time.
    lastWriteTimeUS.store((std::chrono::steady_clock::now() - writeStart) / 1us);
}

bool AtomicDataStore::read(std::string path, std::vector<uint8_t> &outData)
{
    outData.clear();

    std::shared_lock<std::shared_mutex> lock(mutex);

    if (rootHashTablePos & tempPosBit)
    {
        return false;
    }

    uint64_t hashTablePos = rootHashTablePos;
    uint64_t hashTableSize = rootHashTableSize;
    HashTableEntry entry;
    MurmurHash3 hasher;

    for (;;)
    {
        hasher.feed(path.data(), path.size());
        uint64_t index = hasher.digest() % numHashTableEntries;

        std::vector<uint8_t> compressedHashTable(hashTableSize);
        if (!file.read(hashTablePos, compressedHashTable.data(), compressedHashTable.size()))
        {
            crash(
                formatString(
                    "Failed to read a hash table: hashTablePos=0x$+$+",
                    std::hex, hashTablePos
                )
            );
        }

        std::vector<uint8_t> uncompressedHashTable;
        if (!uncompressBytes(compressedHashTable, uncompressedHashTable))
        {
            crash("Failed to uncompress a hash table.");
        }

        if (uncompressedHashTable.size() != sizeof(HashTable))
        {
            crash("An uncompressed hash table is the wrong size.");
        }

        uint64_t entryPos = sizeof(entry) * index;
        memcpy(&entry, uncompressedHashTable.data() + entryPos, sizeof(entry));

        if (entry.type != HashTableEntryType::NestedHashTable)
        {
            break;
        }
        
        hashTablePos = entry.pos;
        hashTableSize = entry.dataSize;
    }

    if (entry.type != HashTableEntryType::FileDescriptor)
    {
        return false;
    }

    std::vector<char> nameChars;
    nameChars.resize(entry.fileNameLen);

    if (!file.read(entry.pos, nameChars.data(), nameChars.size()))
    {
        crash("Failed to read a file name.");
    }

    std::string nameFound(nameChars.data(), nameChars.size());
    if (nameFound != path)
    {
        return false;
    }

    std::vector<uint8_t> payload;
    payload.resize(entry.dataSize);

    if (!file.read(entry.pos + entry.fileNameLen, payload.data(), payload.size()))
    {
        crash("Failed to read file data.");
    }

    if (!uncompressBytes(payload, outData))
    {
        crash("Failed to uncompress file data.");
    }

    return true;
}

bool AtomicDataStore::hasGenerationPassed(uint64_t generation)
{
    return ((lastFinishedGeneration.load() - generation) & (UINT64_C(1) << 63)) == 0;
}

uint64_t AtomicDataStore::getLastFinishedGeneration()
{
    return lastFinishedGeneration.load();
}

uint64_t AtomicDataStore::getLastWriteTimeUS()
{
    return lastWriteTimeUS.load();
}

std::vector<std::string> AtomicDataStore::getAllFiles()
{
    std::vector<std::string> paths;

    std::shared_lock<std::shared_mutex> lock(mutex);

    if (rootHashTablePos & tempPosBit)
    {
        return paths;
    }

    scanHashTable(paths, rootHashTablePos, rootHashTableSize);

    return paths;
}