#include <catch2/catch_test_macros.hpp>
#include <chrono>

#include "fold/symbolic/batch_safety_observer.hpp"

namespace {

class RecordingObserver final : public fold::SafetyObserver {
  public:
    void onDivision(const z3::expr&, const z3::expr&, const fold::ir::Expr&, const fold::ir::Expr&,
                    const z3::expr&, const fold::ir::SourceLoc&, bool, unsigned,
                    const fold::DefUseContext*, const fold::SymbolicEnv*) override {
        ++divisions;
    }

    void onFree(const z3::expr&, const fold::ir::Expr&, const z3::expr&, const fold::ir::SourceLoc&,
                bool, bool) override {
        ++frees;
    }

    int divisions = 0;
    int frees = 0;
};

} // namespace

TEST_CASE("BatchSafetyObserver defers expensive hazard hooks", "[symbolic][observer]") {
    z3::context context;
    const z3::expr lhs = context.bv_val(12, 32);
    const z3::expr rhs = context.bv_val(3, 32);
    const z3::expr path = context.bool_val(true);
    const fold::ir::Expr expression;
    const fold::ir::SourceLoc location;
    RecordingObserver recording;
    fold::BatchSafetyObserver batch(&recording);

    batch.onDivision(lhs, rhs, expression, expression, path, location, true, 32);

    CHECK(recording.divisions == 0);
    CHECK(batch.pending() == 1);
    batch.flush();
    CHECK(recording.divisions == 1);
    CHECK(batch.pending() == 0);
}

TEST_CASE("BatchSafetyObserver preserves immediate stateful hooks", "[symbolic][observer]") {
    z3::context context;
    const z3::expr pointer = context.bv_val(1, 64);
    const z3::expr path = context.bool_val(true);
    RecordingObserver recording;
    fold::BatchSafetyObserver batch(&recording);

    batch.onFree(pointer, {}, path, {}, true, false);

    CHECK(recording.frees == 1);
    CHECK(batch.pending() == 0);
}

TEST_CASE("BatchSafetyObserver drops pending work after its deadline", "[symbolic][observer]") {
    z3::context context;
    const z3::expr value = context.bv_val(1, 32);
    const z3::expr path = context.bool_val(true);
    RecordingObserver recording;
    fold::BatchSafetyObserver batch(&recording);
    batch.onDivision(value, value, {}, {}, path, {}, true, 32);
    batch.setDeadline(std::chrono::steady_clock::now() - std::chrono::milliseconds(1));

    batch.flush();

    CHECK(recording.divisions == 0);
    CHECK(batch.pending() == 0);
    CHECK(batch.deadlineHit());
}