#pragma once
#include <map>
#include <set>
#include <string>
#include <vector>

namespace whetstone {

class LSPCapabilityNegotiator {
public:
    void registerSource(const std::string& source,
                        const std::vector<std::string>& caps) {
        caps_[source] = {caps.begin(), caps.end()};
    }

    bool supports(const std::string& source,
                  const std::string& capability) const {
        auto it = caps_.find(source);
        if (it == caps_.end()) return false;
        return it->second.count(capability) > 0;
    }

    // Returns true if the source supports the capability required by this method.
    // Unknown methods pass through (return true).
    bool shouldRoute(const std::string& source,
                     const std::string& method) const {
        auto it = methodToCap_.find(method);
        if (it == methodToCap_.end()) return true;  // unknown method: pass through
        return supports(source, it->second);
    }

    std::vector<std::string> capabilitiesFor(const std::string& source) const {
        auto it = caps_.find(source);
        if (it == caps_.end()) return {};
        return {it->second.begin(), it->second.end()};
    }

    std::vector<std::string> sourcesSupporting(const std::string& cap) const {
        std::vector<std::string> result;
        for (const auto& [src, capSet] : caps_)
            if (capSet.count(cap)) result.push_back(src);
        return result;
    }

private:
    std::map<std::string, std::set<std::string>> caps_;

    const std::map<std::string, std::string> methodToCap_ = {
        {"textDocument/hover",       "hoverProvider"},
        {"textDocument/definition",  "definitionProvider"},
        {"textDocument/references",  "referencesProvider"},
        {"textDocument/completion",  "completionProvider"},
        {"textDocument/rename",      "renameProvider"},
        {"textDocument/formatting",  "documentFormattingProvider"},
    };
};

} // namespace whetstone