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

namespace whetstone {

struct GuardResult {
    bool                     passed = true;
    std::vector<std::string> regressions;   // approved langs whose output changed
    std::vector<std::string> newLanguages;  // unapproved langs present in results
};

class ParityRegressionGuard {
public:
    // Mark (astNodeId, language) as approved with an expected output.
    void approve(const std::string& astNodeId, const std::string& language,
                 const std::string& expectedOutput) {
        baselines_[{astNodeId, language}] = expectedOutput;
    }

    void updateBaseline(const std::string& astNodeId, const std::string& language,
                        const std::string& newOutput) {
        baselines_[{astNodeId, language}] = newOutput;
    }

    bool isApproved(const std::string& astNodeId, const std::string& language) const {
        return baselines_.count({astNodeId, language}) > 0;
    }

    int approvedCount(const std::string& astNodeId) const {
        int n = 0;
        for (const auto& [key, _] : baselines_)
            if (key.first == astNodeId) ++n;
        return n;
    }

    // Check results against baselines.
    // regressions: approved language whose actual output != baseline.
    // newLanguages: language in results that has no baseline.
    GuardResult guard(const std::string& astNodeId,
                      const std::vector<RunResult>& results) const {
        GuardResult gr;
        for (const auto& r : results) {
            if (!r.ran) continue;
            auto key = std::make_pair(astNodeId, r.language);
            auto it  = baselines_.find(key);
            if (it == baselines_.end()) {
                gr.newLanguages.push_back(r.language);
            } else if (r.output != it->second) {
                gr.regressions.push_back(r.language);
            }
        }
        gr.passed = gr.regressions.empty();
        return gr;
    }

private:
    std::map<std::pair<std::string, std::string>, std::string> baselines_;
};

} // namespace whetstone