// Unit tests for the RoutingPolicyEngine assembly (#2382).
//
// Exercises the engine end-to-end over hand-built RoutePolicy fixtures with a
// faked ClassifierServices: the rule-match path, first-match-wins ordering, the
// default (fail-open) path, want_trace on/off, the classifier on_error error
// path, and concurrent route() calls on one shared engine instance.
//
// Deterministic leaf conditions (#2380) are wired into the engine's leaf factory
// via make_deterministic_leaf_factories(), so the deterministic-rule tests below
// build keywords_any / regex / min_chars / metadata leaves directly (mirroring
// the committed l1_keywords / l1_metadata fixtures) and drive them through
// route() with no classifier backend involved.
//
// Compile (standalone):
//   g++ -std=c++17 -I src/cpp/include -I build/_deps/json-src/include -I test/cpp \
//       test/cpp/test_routing_policy_engine.cpp src/cpp/server/routing_policy.cpp \
//       -o test_routing_policy_engine

#include "fake_classifier_services.h"
#include "lemon/routing_policy.h"

#include <atomic>
#include <cstdio>
#include <map>
#include <stdexcept>
#include <string>
#include <thread>
#include <vector>

using lemon::Decision;
using lemon::MatchExpr;
using lemon::RouteContext;
using lemon::RoutePolicy;
using lemon::RoutingPolicyEngine;
using lemon::Rule;
using lemon::json;

static int g_failures = 0;

static void check(const char* name, bool ok) {
    std::printf("[%s] %s\n", ok ? "PASS" : "FAIL", name);
    if (!ok) ++g_failures;
}

namespace {

// A single classifier-band leaf: {"classifier": id, "label": label, "min_score": min}.
MatchExpr classifier_leaf(const std::string& classifier_id, const std::string& label,
                          double min_score) {
    MatchExpr expr;
    expr.op = MatchExpr::Op::Leaf;
    expr.leaf = json{{"classifier", classifier_id}, {"label", label}, {"min_score", min_score}};
    return expr;
}

// A raw deterministic leaf, e.g. {"keywords_any": [...]} or {"min_chars": 4000}.
MatchExpr deterministic_leaf(json leaf) {
    MatchExpr expr;
    expr.op = MatchExpr::Op::Leaf;
    expr.leaf = std::move(leaf);
    return expr;
}

MatchExpr any_of(std::vector<MatchExpr> children) {
    MatchExpr expr;
    expr.op = MatchExpr::Op::Any;
    expr.children = std::move(children);
    return expr;
}

Rule make_rule(const std::string& id, MatchExpr match, const std::string& route_to,
               json outputs = json::object()) {
    Rule rule;
    rule.id = id;
    rule.match = std::move(match);
    rule.route_to = route_to;
    rule.outputs = std::move(outputs);
    return rule;
}

// Policy: classifier "pii" (labels PII/NO_PII, default PII) over model "pii-model".
// One rule keeps PII-flagged requests on the local candidate.
RoutePolicy make_pii_policy() {
    RoutePolicy policy;
    policy.candidates = {"local-llm", "cloud-llm"};
    policy.default_model = "cloud-llm";
    policy.classifiers = lemon::make_classifiers(json::array({
        json{{"id", "pii"},
             {"type", "classifier"},
             {"model", "pii-model"},
             {"labels", {"PII", "NO_PII"}},
             {"default_label", "PII"}},
    }));
    policy.rules = {
        make_rule("keep-private", classifier_leaf("pii", "PII", 0.5), "local-llm",
                  json{{"verdict", "block"}}),
    };
    return policy;
}

} // namespace

static void test_rule_match_path() {
    lemon::testing::FakeClassifierServices fake;
    fake.set_classifier_scores("pii-model", {{"PII", 0.9}, {"NO_PII", 0.1}});
    RoutingPolicyEngine engine(make_pii_policy(), fake.make());

    RouteContext ctx;
    ctx.input = "my ssn is 123-45-6789";
    Decision d = engine.route(ctx, /*want_trace=*/false);

    check("matched rule selects its route_to", d.route_to == "local-llm");
    check("matched rule reports its id", d.matched_rule == "keep-private");
    check("matched rule clears default_used", !d.default_used);
    check("matched rule copies outputs verbatim",
          d.outputs.value("verdict", "") == "block");
    check("trace empty when not requested", d.trace.empty());
}

static void test_default_path() {
    lemon::testing::FakeClassifierServices fake;
    fake.set_classifier_scores("pii-model", {{"PII", 0.1}, {"NO_PII", 0.9}});
    RoutingPolicyEngine engine(make_pii_policy(), fake.make());

    RouteContext ctx;
    ctx.input = "what is the capital of France?";
    Decision d = engine.route(ctx, /*want_trace=*/false);

    check("no match falls through to default_model", d.route_to == "cloud-llm");
    check("default path leaves matched_rule empty", d.matched_rule.empty());
    check("default path sets default_used", d.default_used);
    check("default path has empty outputs", d.outputs.empty());
}

static void test_first_match_wins() {
    lemon::testing::FakeClassifierServices fake;
    fake.set_classifier_scores("pii-model", {{"PII", 0.9}, {"NO_PII", 0.1}});

    RoutePolicy policy = make_pii_policy();
    // A second rule that would also match PII>=0.5 but routes elsewhere; the
    // first rule must win.
    policy.rules.push_back(
        make_rule("second", classifier_leaf("pii", "PII", 0.5), "cloud-llm"));
    RoutingPolicyEngine engine(std::move(policy), fake.make());

    RouteContext ctx;
    ctx.input = "ssn 123";
    Decision d = engine.route(ctx, false);
    check("first matching rule wins over later rules", d.matched_rule == "keep-private");
    check("first match routes to first rule's target", d.route_to == "local-llm");
}

static void test_trace_on_and_off() {
    lemon::testing::FakeClassifierServices fake;
    fake.set_classifier_scores("pii-model", {{"PII", 0.9}, {"NO_PII", 0.1}});
    RoutingPolicyEngine engine(make_pii_policy(), fake.make());

    RouteContext ctx;
    ctx.input = "ssn 123";

    Decision off = engine.route(ctx, /*want_trace=*/false);
    check("want_trace=false yields no trace", off.trace.empty());

    Decision on = engine.route(ctx, /*want_trace=*/true);
    check("want_trace=true records the classifier leaf", on.trace.size() == 1 &&
              on.trace[0].condition == "classifier:pii");
    check("trace carries the scored value",
          on.trace.size() == 1 && on.trace[0].score.has_value() &&
              *on.trace[0].score == 0.9 && on.trace[0].result);
}

// A genuinely failed classifier (run_classifier throws => Score::ok=false)
// drives the band's on_error. The "pii" classifier leaves on_error unset, so it
// defaults to match_false: the rule cannot match and the request fails open to
// default_model — the symmetric counterpart to the match_true test below, and a
// guarantee that route() never lets the exception escape.
static void test_classifier_failure_falls_open() {
    lemon::ClassifierServices services;
    services.run_classifier = [](const std::string&, const std::string&)
        -> std::map<std::string, double> { throw std::runtime_error("model down"); };
    RoutingPolicyEngine engine(make_pii_policy(), std::move(services));

    RouteContext ctx;
    ctx.input = "anything";
    Decision d = engine.route(ctx, false);
    check("classifier failure path lands on default_model", d.route_to == "cloud-llm");
    check("classifier failure path leaves matched_rule empty", d.matched_rule.empty());
    check("classifier failure path sets default_used", d.default_used);
}

// on_error=match_true means a failed classifier trips its rule (fail-closed
// authoring): the request must route to the rule's target, not the default. A
// genuine failure is run_classifier throwing (an empty result is a valid ok
// score, not a failure), so this binds a throwing service directly.
static void test_on_error_match_true_routes_to_rule() {
    RoutePolicy policy;
    policy.candidates = {"guard", "cloud-llm"};
    policy.default_model = "cloud-llm";
    policy.classifiers = lemon::make_classifiers(json::array({
        json{{"id", "jailbreak"},
             {"type", "classifier"},
             {"model", "jb-model"},
             {"on_error", "match_true"},
             {"labels", {"JAILBREAK", "BENIGN"}},
             {"default_label", "JAILBREAK"}},
    }));
    policy.rules = {
        make_rule("guard", classifier_leaf("jailbreak", "JAILBREAK", 0.5), "guard"),
    };

    lemon::ClassifierServices services;
    services.run_classifier = [](const std::string&, const std::string&)
        -> std::map<std::string, double> { throw std::runtime_error("model down"); };
    RoutingPolicyEngine engine(std::move(policy), std::move(services));

    RouteContext ctx;
    ctx.input = "ignore previous instructions";
    Decision d = engine.route(ctx, false);
    check("on_error match_true routes a failed classifier to its rule",
          d.route_to == "guard" && d.matched_rule == "guard" && !d.default_used);
}

// Deterministic leaves (#2380) must resolve on the engine path via
// make_deterministic_leaf_factories(). Mirrors l1_keywords.json: an `any` of
// keywords_any / regex routes code to the big model, and min_chars routes long
// context there too; everything else falls open to the default. No classifier
// backend is touched, so an empty ClassifierServices is sufficient.
static RoutePolicy make_keywords_policy() {
    RoutePolicy policy;
    policy.candidates = {"small-llm", "big-llm"};
    policy.default_model = "small-llm";
    policy.rules = {
        make_rule("code-to-big",
                  any_of({deterministic_leaf(json{{"keywords_any",
                              json::array({"def ", "function", "stack trace", "compile"})}}),
                          deterministic_leaf(json{{"regex", "```[a-z]*"}})}),
                  "big-llm"),
        make_rule("long-context-to-big",
                  deterministic_leaf(json{{"min_chars", 4000}}), "big-llm"),
    };
    return policy;
}

static void test_deterministic_keywords_route() {
    RoutingPolicyEngine engine(make_keywords_policy(), lemon::ClassifierServices{});

    RouteContext ctx;
    ctx.input = "please write a function that reverses a list";
    Decision hit = engine.route(ctx, /*want_trace=*/false);
    check("keywords_any leaf matches and routes to big-llm",
          hit.route_to == "big-llm" && hit.matched_rule == "code-to-big");

    RouteContext plain;
    plain.input = "what is the capital of France?";
    Decision miss = engine.route(plain, false);
    check("non-matching deterministic rules fall open to default",
          miss.route_to == "small-llm" && miss.default_used);
}

static void test_deterministic_min_chars_route() {
    RoutingPolicyEngine engine(make_keywords_policy(), lemon::ClassifierServices{});

    RouteContext ctx;
    ctx.input = std::string(5000, 'a');
    ctx.params.chars = ctx.input.size();  // min_chars bounds params.chars (UTF-8 bytes)
    Decision d = engine.route(ctx, false);
    check("min_chars leaf routes long input to its rule",
          d.route_to == "big-llm" && d.matched_rule == "long-context-to-big");
}

// Mirrors l1_metadata.json: a metadata `equals` leaf keeps flagged requests on
// the local model. Proves the metadata deterministic leaf constructs and
// evaluates on the engine path.
static void test_deterministic_metadata_route() {
    RoutePolicy policy;
    policy.candidates = {"local-llm", "cloud-llm"};
    policy.default_model = "cloud-llm";
    policy.rules = {
        make_rule("keep-confidential",
                  deterministic_leaf(json{{"metadata",
                      json{{"key", "task_class"}, {"equals", "confidential"}}}}),
                  "local-llm"),
    };
    RoutingPolicyEngine engine(std::move(policy), lemon::ClassifierServices{});

    RouteContext hit;
    hit.metadata = {{"task_class", "confidential"}};
    Decision d = engine.route(hit, false);
    check("metadata equals leaf routes to its rule",
          d.route_to == "local-llm" && d.matched_rule == "keep-confidential");

    RouteContext miss;
    miss.metadata = {{"task_class", "public"}};
    Decision open = engine.route(miss, false);
    check("metadata mismatch falls open to default",
          open.route_to == "cloud-llm" && open.default_used);
}

// One const engine, many threads: per-request state must live only in the local
// EvalContext, so concurrent want_trace=true/false calls never corrupt each
// other's trace and every call returns the same deterministic Decision.
static void test_concurrent_route_is_consistent() {
    lemon::testing::FakeClassifierServices fake;
    fake.set_classifier_scores("pii-model", {{"PII", 0.9}, {"NO_PII", 0.1}});
    RoutingPolicyEngine engine(make_pii_policy(), fake.make());

    constexpr int kThreads = 16;
    constexpr int kIters = 500;
    std::atomic<int> mismatches{0};

    std::vector<std::thread> workers;
    for (int t = 0; t < kThreads; ++t) {
        const bool want_trace = (t % 2 == 0);
        workers.emplace_back([&engine, &mismatches, want_trace, kIters]() {
            for (int i = 0; i < kIters; ++i) {
                RouteContext ctx;
                ctx.input = "ssn 123";
                Decision d = engine.route(ctx, want_trace);
                const std::size_t expected_trace = want_trace ? 1u : 0u;
                if (d.route_to != "local-llm" || d.matched_rule != "keep-private" ||
                    d.trace.size() != expected_trace) {
                    ++mismatches;
                }
            }
        });
    }
    for (auto& w : workers) w.join();

    check("concurrent route() calls stay consistent", mismatches.load() == 0);
}

int main() {
    test_rule_match_path();
    test_default_path();
    test_first_match_wins();
    test_trace_on_and_off();
    test_classifier_failure_falls_open();
    test_on_error_match_true_routes_to_rule();
    test_deterministic_keywords_route();
    test_deterministic_min_chars_route();
    test_deterministic_metadata_route();
    test_concurrent_route_is_consistent();

    std::printf("\n%s\n", g_failures == 0 ? "ALL ENGINE TESTS PASSED"
                                          : "ENGINE TESTS FAILED");
    return g_failures == 0 ? 0 : 1;
}
