#include <rocksdb/db.h>
#include <rocksdb/options.h>
#include <iostream>
#include <string>
#include <unordered_map>
#include <vector>
#include <regex>
#include "tser/tser.hpp"

// Wrapper class for RocksDB
class RocksDBWrapper
{
public:
    RocksDBWrapper(const std::string &db_path, const std::string &index_path);

    ~RocksDBWrapper();

    template <typename T>
    void store(const std::string &key, const T &value)
    {
        tser::BinaryArchive archive(0);
        archive.save(value);

        rocksdb::Status status = db_->Put(rocksdb::WriteOptions(), serializeKey(key), archive.get_buffer().data());
        if (!status.ok())
        {
            std::cerr << "Error storing value in RocksDB database: " << status.ToString() << std::endl;
        }

        // Index the struct's members
        index_members(serializeKey(key), value);
    }

    template <typename T>
    bool get(const std::string &key, T &value)
    {
        std::string serialized_value;
        rocksdb::Status status = db_->Get(rocksdb::ReadOptions(), serializeKey(key), &serialized_value);
        if (!status.ok())
        {
            std::cerr << "Error getting value from RocksDB database: " << status.ToString() << std::endl;
            return false;
        }

        tser::BinaryArchive archive(0);
        archive.initialize(serialized_value);
        value = archive.load<T>();
        return true;
    }

    template <typename T>
    std::vector<std::string> search(const std::string &member_name, const std::string &member_value)
    {
        std::vector<std::string> keys;
        std::string prefix = member_name + ":" + serializeKey(member_value) + ":";
        rocksdb::Iterator *it = index_db_->NewIterator(rocksdb::ReadOptions());
        for (it->Seek(prefix); it->Valid() && it->key().ToString().find(prefix) == 0; it->Next())
        {
            keys.push_back(it->value().ToString());
        }
        delete it;
        return keys;
    }

    inline std::vector<std::string> split(const std::string &s, const std::string &delimiter)
    {
        size_t last = 0;
        size_t next = 0;
        std::vector<std::string> results;
        while ((next = s.find(delimiter, last)) != std::string::npos)
        {
            results.emplace_back(s.substr(last, next - last));
            last = next + delimiter.length();
        }
        results.emplace_back(s.substr(last));
        return results;
    }

    template <typename T>
    void search_text(const std::string &search_value, std::vector<T> &res)
    {
        std::vector<std::string> keys;
        std::vector<std::string> tmp;
        rocksdb::Iterator *it = index_db_->NewIterator(rocksdb::ReadOptions());

        for (it->SeekToFirst(); it->Valid(); it->Next())
        {
            std::string key = it->key().ToString();
            auto _split = this->split(key, ":");
            if (this->StringMatch(_split[1], search_value))
            {
                keys.push_back(_split[2]);
            }
        }

        std::vector<rocksdb::Slice> slices;
        for (const auto &key : keys)
        {
            slices.push_back(rocksdb::Slice(key));
        }

        std::vector<rocksdb::Status> statuses = db_->MultiGet(rocksdb::ReadOptions(), slices, &tmp);

        for (const auto &_v : tmp)
        {
            T _val;
            this->DeSerialize(_v, _val);
            res.emplace_back(_val);
        }
        delete it;
    }

    // New keyExists method
    bool keyExists(const std::string &key);

    // New conditional search method
    template <typename T>
    void search_conditional(std::vector<T> &res, std::function<bool(const T &)> condition)
    {
        rocksdb::Iterator *it = db_->NewIterator(rocksdb::ReadOptions());

        for (it->SeekToFirst(); it->Valid(); it->Next())
        {
            std::string key = it->key().ToString();
            std::string serialized_value = it->value().ToString();

            T _val;
            this->DeSerialize(serialized_value, _val);

            if (condition(_val))
            {
                res.push_back(_val);
            }
        }
        delete it;
    }

private:
    rocksdb::DB *db_;
    rocksdb::DB *index_db_;

    bool StringMatch(const std::string &text, const std::string &pattern);

    template <typename T>
    inline void DeSerialize(const std::string &data, T &value)
    {
        tser::BinaryArchive archive(0);
        archive.initialize(data);
        value = archive.load<T>();
    }

    template <typename T>
    inline void Serialize(const T &value, std::string &data)
    {
        tser::BinaryArchive archive(0);
        archive.save(value);
        data = archive.get_buffer().data();
    }

    template <typename T>
    void index_single_member(const std::string &key, const std::string &member_name, const T &member_value)
    {
        if constexpr (std::is_arithmetic_v<std::decay_t<T>>)
        {
            std::string member_value_str = std::to_string(member_value);
            std::string index_key = member_name + ":" + serializeKey(member_value_str) + ":" + key;
            rocksdb::Status status = index_db_->Put(rocksdb::WriteOptions(), index_key, key);
            if (!status.ok())
            {
                std::cerr << "Error indexing member in RocksDB index database: " << status.ToString() << std::endl;
            }
        }
        else if constexpr (std::is_same_v<std::decay_t<T>, std::string>)
        {
            std::string index_key = member_name + ":" + serializeKey(member_value) + ":" + key;
            rocksdb::Status status = index_db_->Put(rocksdb::WriteOptions(), index_key, key);
            if (!status.ok())
            {
                std::cerr << "Error indexing member in RocksDB index database: " << status.ToString() << std::endl;
            }
        }
        else if constexpr (tser::is_container_v<std::decay_t<T>>)
        {
            for (const auto &elem : member_value)
            {
                index_single_member(key, member_name, elem);
            }
        }
        else
        {
            // Handle other types if necessary
        }
    }

    template <typename T>
    void index_members(const std::string &key, const T &value)
    {
        if constexpr (tser::is_tser_t_v<T>)
        {
            auto member_names = T::_memberNames;
            size_t index = 0;
            std::apply([&](auto &&...args)
                       { ((index_single_member(key, member_names[index++], args)), ...); },
                       value.members());
        }
    }

    std::string serializeKey(const std::string &key)
    {
        // Implement a method to serialize keys to handle special characters
        std::string serialized_key;
        for (const char c : key)
        {
            if (c == ':')
            {
                serialized_key += "%3A"; // Percent-encoding for ':'
            }
            else
            {
                serialized_key += c;
            }
        }
        return serialized_key;
    }
};