#ifndef PROXYSERVER_CONN_MAP_H
#define PROXYSERVER_CONN_MAP_H

#include <iostream>
#include <map>
#include <shared_mutex>
#include <mutex>
#include <memory>


template <typename Key, typename Value>
class SafeMap {
public:

    void add(const Key& key, std::unique_ptr<Value> &value) {
        std::unique_lock<std::shared_timed_mutex> lock(mutex_);
        map_[key] = std::move(value);
    }

    Value* get(const Key& key) const {
        std::shared_lock<std::shared_timed_mutex> lock(mutex_);
        auto it = map_.find(key);
        if (it != map_.end()) {
            return it->second.get();
        }
        return nullptr;
    }

    void remove(const Key& key) {
        std::unique_lock<std::shared_timed_mutex> lock(mutex_);
        map_.erase(key);
    }

    void clear() {
        std::unique_lock<std::shared_timed_mutex> lock(mutex_);
        map_.clear();
    }
    
    typename std::map<Key, std::unique_ptr<Value>>::iterator begin() {
        return map_.begin();
    }

    typename std::map<Key, std::unique_ptr<Value>>::iterator end() {
        return map_.end();
    }
    
protected:
    mutable std::shared_timed_mutex mutex_;
    std::map<Key, std::unique_ptr<Value>> map_;
};

template <typename Key, typename Value>
class TcpConnMap : public SafeMap<Key, Value> {
public:
    static TcpConnMap& getInstance() {
        static TcpConnMap instance; // 在首次使用时创建
        return instance;
    }

    TcpConnMap(TcpConnMap const&) = delete;
    void operator=(TcpConnMap const&) = delete;

private:
    TcpConnMap() = default;
};

template <typename Key, typename Value>
class UdpConnMap : public SafeMap<Key, Value> {
public:
    static UdpConnMap& getInstance() {
        static UdpConnMap instance; // 在首次使用时创建
        return instance;
    }

    UdpConnMap(UdpConnMap const&) = delete;
    void operator=(UdpConnMap const&) = delete;

private:
    UdpConnMap() = default;
};

#endif //PROXYSERVER_CONN_MAP_H