#ifndef ATOMIC_DATA_STORE_HPP_
#define ATOMIC_DATA_STORE_HPP_

#include <MurmurHash3/MurmurHash3.hpp>
#include <Version/Version.hpp>
#include <atomic>
#include <cstddef>
#include <cstdint>
#include <string>
#include <unordered_map>
#include <vector>
#include <shared_mutex>
#include <FileIO/FileIO.hpp>
#include <SHA256/SHA256.hpp>

/**
 * @brief Manages access to files with an atomic key-value store.
 */
class AtomicDataStore
{
private:
    static constexpr size_t numHashTableEntries = 512;
    static constexpr uint64_t tempPosBit = UINT64_C(1) << 63;

    struct MainHeader
    {
        char magic[4];              // "\0VVW"
    };

    struct Region
    {
        uint64_t basePos;
        uint64_t size;
    };

    struct IndexBlock
    {
        uint64_t rootHashTablePos;
        uint32_t rootHashTableSize;
        uint32_t resv;
        uint64_t freeRegionTablePos;
        uint64_t numFreeRegions;
        uint64_t generation;
        SHA256::Digest hash;

        void updateHash();
        bool isValid() const;
    };

    enum class HashTableEntryType : uint8_t
    {
        Unmapped = 0,
        NestedHashTable = 1,
        FileDescriptor = 2,
    };

    struct HashTableEntry
    {
        HashTableEntryType type;
        uint8_t resv;
        uint16_t fileNameLen;
        uint32_t dataSize;
        uint64_t pos;
    };

    struct HashTable
    {
        HashTableEntry entries[numHashTableEntries];
    };

    FileIO file;
    std::string error;
    size_t nextIndexBlockIndex;
    uint64_t nextIndexGeneration;
    std::vector<Region> freeRegionTable;
    std::vector<Region> regionsToRelease;
    uint64_t rootHashTablePos;
    uint32_t rootHashTableSize;
    uint64_t nextTempHashTableIndex = tempPosBit;
    bool headerNeedsWriting = false;

    /**
     * @brief Maps the CURRENT positions of hash tables ("IDs") to their mutable form.
     */
    std::unordered_map<uint64_t, HashTable> pendingHashTables;

    std::shared_mutex mutex;
    std::atomic<uint64_t> lastFinishedGeneration;
    std::atomic<uint64_t> lastWriteTimeUS;

    bool readIndexBlock(IndexBlock &output, size_t &outNextIndex);
    uint64_t allocStorageRegion(size_t size);
    HashTable& getMutableHashTable(uint64_t id, uint64_t compressedSize);
    void encodeFileRecur(
        MurmurHash3 &hasher,
        const std::string &filename,
        const std::vector<uint8_t> &payload,
        uint64_t hashtabPos,
        uint64_t hashtabCompressedSize,
        int depth
    );
    void writeHashTable(uint64_t &pos, uint32_t &size);
    void flushHashTables();
    void encodeFileDescriptor(
        HashTableEntry &entry,
        const std::string &filename,
        const std::vector<uint8_t> &payload
    );
    HashTable& createMutableHashTable(uint64_t &pos);
    void scanHashTable(std::vector<std::string> &paths, uint64_t pos, uint32_t size);

public:
    typedef std::unordered_map<std::string, std::vector<uint8_t>> Plan;

    /**
     * @brief Create and open a new store at the specified path.
     *
     * The store will be opened by this instance, and the signature is not written
     * until you call `update()`. You can therefore ensure that the store will be in
     * a specific state before it is recognized as a store.
     */
    bool create(const std::string &path);

    /**
     * @brief Try to open an atomic data store at the specified path.
     *
     * This will perform all necessary checks to ensure the file is valid.
     * Returns true if the file was opened successfully.
     */
    bool open(const std::string &path);

    const std::string& getError() const;

    /**
     * @brief Apply the specified plan to this world.
     * 
     * This can be safely called from multiple threads.
     * 
     * @param plan Maps file names to new content.
     * @param generation The generation of this plan.
     */
    void update(const Plan &plan, uint64_t generation);

    /**
     * @brief Read a file into the vector.
     * 
     * @param path The path to the file (relative to the world data dir).
     * @param outData Byte vector to which the file data will be assigned.
     * @return true If the read succeeded.
     * @return false If the read failed.
     */
    bool read(std::string path, std::vector<uint8_t> &outData);

    /**
     * @brief Returns true if the specified generation has passed.
     * 
     * This means that all the data from the plan in that generation has
     * been fully written to disk.
     * 
     * This can safely be called from any thread.
     * 
     * @param generation The plan generation.
     */
    bool hasGenerationPassed(uint64_t generation);

    /**
     * @brief Return the last generation number fully written to disk.
     */
    uint64_t getLastFinishedGeneration();

    /**
     * @brief Get the duration of the last write, in microseconds.
     */
    uint64_t getLastWriteTimeUS();

    /**
     * @brief Returns the paths of all files in the store.
     */
    std::vector<std::string> getAllFiles();

    friend class WDumpCommand;
    friend class WUnpackCommand;
    friend class WTestCommand;
    friend class AsyncWorldCloner;
    friend class WRepackCommand;
};

#endif