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

namespace whetstone {

struct DocumentInfo {
    std::string uri;
    std::string languageId;
    std::string content;
    int         version = 0;
};

class LSPWorkspaceIndex {
public:
    void open(const std::string& uri, const std::string& languageId,
              const std::string& content, int version) {
        docs_[uri] = DocumentInfo{uri, languageId, content, version};
    }

    void update(const std::string& uri, const std::string& content, int version) {
        auto it = docs_.find(uri);
        if (it == docs_.end()) return;
        it->second.content = content;
        it->second.version = version;
    }

    void close(const std::string& uri) {
        docs_.erase(uri);
    }

    bool contains(const std::string& uri) const {
        return docs_.count(uri) > 0;
    }

    std::string getLanguage(const std::string& uri) const {
        auto it = docs_.find(uri);
        if (it == docs_.end()) return "";
        return it->second.languageId;
    }

    DocumentInfo getDocument(const std::string& uri) const {
        auto it = docs_.find(uri);
        if (it == docs_.end()) return {};
        return it->second;
    }

    std::vector<std::string> listUris(const std::string& languageId) const {
        std::vector<std::string> result;
        for (const auto& [uri, info] : docs_)
            if (info.languageId == languageId)
                result.push_back(uri);
        return result;
    }

    std::vector<std::string> allLanguages() const {
        std::set<std::string> langs;
        for (const auto& [uri, info] : docs_)
            langs.insert(info.languageId);
        return {langs.begin(), langs.end()};
    }

    int size() const { return static_cast<int>(docs_.size()); }

private:
    std::map<std::string, DocumentInfo> docs_;
};

} // namespace whetstone