#pragma once
#include "LSPDiagnosticsAggregator.h"
#include <map>
#include <set>
#include <string>
#include <utility>
#include <vector>

namespace whetstone {

class LSPErrorRecovery {
public:
    void markCrashed(const std::string& source) {
        crashed_.insert(source);
    }

    void markRecovered(const std::string& source) {
        crashed_.erase(source);
    }

    // Unknown sources are healthy (optimistic default).
    bool isHealthy(const std::string& source) const {
        return crashed_.count(source) == 0;
    }

    std::vector<std::string> healthySources(
            const std::vector<std::string>& all) const {
        std::vector<std::string> result;
        for (const auto& s : all)
            if (isHealthy(s)) result.push_back(s);
        return result;
    }

    void cacheDiagnostics(const std::string& source, const std::string& uri,
                          const std::vector<Diagnostic>& diags) {
        cache_[{source, uri}] = diags;
    }

    // Returns cached diags if source is crashed, else empty.
    std::vector<Diagnostic> fallbackDiagnostics(const std::string& source,
                                                const std::string& uri) const {
        if (isHealthy(source)) return {};
        auto it = cache_.find({source, uri});
        if (it == cache_.end()) return {};
        return it->second;
    }

    void clearCache(const std::string& source) {
        for (auto it = cache_.begin(); it != cache_.end(); ) {
            if (it->first.first == source) it = cache_.erase(it);
            else ++it;
        }
    }

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

private:
    std::set<std::string> crashed_;
    std::map<std::pair<std::string, std::string>, std::vector<Diagnostic>> cache_;
};

} // namespace whetstone