#include <functional>
#include <iostream>
#include <map>
#include <random>
#include <string>
#include <vector>

// Hash函数，可以根据需求更改为其他函数，这里使用了std::hash
typedef std::function<size_t(const std::string &)> HashFunction;
// 示例哈希函数
size_t hashFunction(const std::string &key) { return std::hash<std::string>{}(key); }


// 一致性哈希类
class ConsistentHash {
  private:
    // 虚拟节点数量，可根据需求调整
    const size_t virtualNodes;
    // 哈希函数
    HashFunction hashFunction;
    // 圆环
    std::map<size_t, std::string> hashRing;

    // 物理节点负载信息，存储节点上的 IP 数量
    std::map<std::string, std::vector<std::string>> nodeIPs;

  public:
    // 需要服务的IP地址
    std::vector<std::string> ipDataset; // 将ipDataset作为类成员
    // 构造函数
    ConsistentHash(size_t vNodes, HashFunction func) : virtualNodes(vNodes), hashFunction(func) {}

    void init(std::vector<std::string> nodes) {
        for (const std::string &node : nodes) {
            addNode(node);
        }
        for (auto &ip : ipDataset) {
            getNode(ip);
        }
    }

    // 生成模拟的 IP 地址数据集
    void generateIPDataset(int dataSize) {
        std::random_device rd;
        std::mt19937 gen(rd());
        std::uniform_int_distribution<> dis(1, 255);

        for (int i = 0; i < dataSize; ++i) {
            std::string ip = std::to_string(dis(gen)) + "." + std::to_string(dis(gen)) + "." +
                             std::to_string(dis(gen)) + "." + std::to_string(dis(gen));
            ipDataset.push_back(ip);
        }
    }

    // 添加节点到哈希环
    void addNode(const std::string &node) {
        for (size_t i = 0; i < virtualNodes; ++i) {
            // 为每个节点创建多个虚拟节点，用不同的哈希值
            std::string virtualNodeName = node + "_" + std::to_string(i);
            size_t hash = hashFunction(virtualNodeName);
            hashRing[hash] = node; // 虚拟节点映射到实际节点
        }
    }

    void resetRing() {
        nodeIPs.clear(); // 清除负载记录
        // 重新分配 IP 地址到虚拟节点
        for (auto &ip : ipDataset) {
            std::string selectedNode = getNode(ip);
        }
    }

    // 删除节点及其对应的所有虚拟节点，并重新分配 IP 地址
    void removeNode(const std::string &node) {
        std::vector<size_t> virtualNodeHashesToRemove;

        // 找到需要移除的节点对应的虚拟节点
        for (auto it = hashRing.begin(); it != hashRing.end();) {
            if (it->second == node) {
                virtualNodeHashesToRemove.push_back(it->first);
                it = hashRing.erase(it);
            } else {
                ++it;
            }
        }

        // 从节点负载信息中移除节点
        nodeIPs.erase(node);

        // 重新分配需要移除节点上的虚拟节点对应的 IP 地址到其他节点
        for (size_t hash : virtualNodeHashesToRemove) {
            for (auto &pair : nodeIPs) {
                std::string newNode = pair.first;
                if (newNode != node) {
                    size_t newHash = hashFunction(newNode + "_" + std::to_string(hash));
                    hashRing[newHash] = newNode;
                }
            }
        }
    }

    // 根据给定的键查找对应的节点
    std::string getNode(const std::string &key) {
        size_t hash = hashFunction(key);
        auto it = hashRing.lower_bound(hash);
        // std::cout << it->first << std::endl;
        if (it == hashRing.end()) {
            // 如果超出了范围，返回第一个节点（环形结构）
            it = hashRing.begin();
        }
        // 将 IP 添加到所选节点
        std::string selectedNode = it->second;
        nodeIPs[selectedNode].push_back(key);
        return selectedNode;
    }

    // 获取节点负载信息，按照 IP 数量计算负载百分比
    void statisticPerf() const {
        std::map<std::string, double> loadPercentages;
        size_t totalKeys = 0;

        // Calculate total number of keys (IPs) on the entire hash ring
        for (const auto &nodeIP : nodeIPs) {
            totalKeys += nodeIP.second.size();
        }

        // Calculate load percentage for each node
        for (const auto &nodeIP : nodeIPs) {
            double percentage = (static_cast<double>(nodeIP.second.size()) / totalKeys) * 100.0;
            loadPercentages[nodeIP.first] = percentage;
        }

        // Output node load percentages
        std::cout << "Node Loads (%):\n";
        for (const auto &load : loadPercentages) {
            std::cout << load.first << ": " << load.second << "%" << std::endl;
        }
    }
};

int main() {
    // 创建一致性哈希对象并传入MurmurHash3作为哈希函数
    ConsistentHash consistentHash(1000, hashFunction); // 构造哈希环，设置虚拟节点个数
    // 生成模拟的 IP 地址数据集
    consistentHash.generateIPDataset(100000); // 生成模拟 IP 地址

    // 初始化存储节点，分配IP到存储节点
    std::vector<std::string> nodeList = {"node1", "node2", "node3", "node4"};
    consistentHash.init(nodeList);
    consistentHash.statisticPerf();

    // 增加存储节点，重新分配所有IP地址
    consistentHash.addNode("node5");
    consistentHash.resetRing();
    consistentHash.statisticPerf();
  
    // 删除存储节点，重新分配该节点上的IP
    consistentHash.removeNode("node3");
    consistentHash.statisticPerf();

    return 0;
}