#include <z3++.h>

#include <catch2/catch_test_macros.hpp>

#include "fold/symbolic/source_path_state.hpp"

namespace {

z3::check_result checkPath(const z3::expr& path, const z3::expr& extra) {
    z3::solver solver(path.ctx(), "QF_BV");
    solver.add(path && extra);
    return solver.check();
}

} // namespace

TEST_CASE("SourcePathState starts with true path", "[source-path-state]") {
    auto state = fold::SourcePathState::create();
    z3::solver solver(state.ctx(), "QF_BV");
    solver.add(!state.currentPath());
    REQUIRE(solver.check() == z3::unsat);
}

TEST_CASE("SourcePathState assume creates truthy and falsy path copies",
          "[source-path-state][assume]") {
    auto state = fold::SourcePathState::create();
    z3::expr cond = state.ctx().bool_const("cond");

    auto branches = state.assume(cond);

    REQUIRE(branches.truthy.has_value());
    REQUIRE(branches.falsy.has_value());
    CHECK(checkPath(branches.truthy->currentPath(), !cond) == z3::unsat);
    CHECK(checkPath(branches.falsy->currentPath(), cond) == z3::unsat);
    CHECK(checkPath(state.currentPath(), cond) == z3::sat);
    CHECK(checkPath(state.currentPath(), !cond) == z3::sat);
}

TEST_CASE("SourcePathState assume pruning is opt-in", "[source-path-state][assume][prune]") {
    auto state = fold::SourcePathState::create();
    z3::expr impossible = state.ctx().bool_val(false);

    auto branches = state.assume(impossible);

    REQUIRE(branches.truthy.has_value());
    REQUIRE(branches.falsy.has_value());
    CHECK(checkPath(branches.truthy->currentPath(), state.ctx().bool_val(true)) == z3::unsat);
    CHECK(checkPath(branches.falsy->currentPath(), state.ctx().bool_val(true)) == z3::sat);
}

TEST_CASE("SourcePathState assume prunes only children proved infeasible",
          "[source-path-state][assume][prune]") {
    auto state = fold::SourcePathState::create();
    z3::expr impossible = state.ctx().bool_val(false);

    auto branches = state.assume(impossible, fold::SourcePathAssumeOptions{
                                                 .prune_infeasible = true,
                                                 .timeout_ms = 100,
                                             });

    CHECK_FALSE(branches.truthy.has_value());
    REQUIRE(branches.falsy.has_value());
    CHECK(checkPath(branches.falsy->currentPath(), impossible) == z3::unsat);
}

TEST_CASE("SourcePathState assume pruning respects current path",
          "[source-path-state][assume][prune]") {
    auto state = fold::SourcePathState::create();
    z3::expr cond = state.ctx().bool_const("cond");
    auto constrained = state.withPath(cond);

    auto branches = constrained.assume(cond, fold::SourcePathAssumeOptions{
                                                 .prune_infeasible = true,
                                                 .timeout_ms = 100,
                                             });

    REQUIRE(branches.truthy.has_value());
    CHECK_FALSE(branches.falsy.has_value());
    CHECK(checkPath(branches.truthy->currentPath(), !cond) == z3::unsat);
}

TEST_CASE("SourcePathState withPath leaves the original state unchanged",
          "[source-path-state][copy]") {
    auto state = fold::SourcePathState::create();
    z3::expr cond = state.ctx().bool_const("cond");

    auto constrained = state.withPath(cond);

    CHECK(checkPath(constrained.currentPath(), !cond) == z3::unsat);
    CHECK(checkPath(state.currentPath(), cond) == z3::sat);
    CHECK(checkPath(state.currentPath(), !cond) == z3::sat);
}

TEST_CASE("ScopedSourcePath restores raw SymbolicEnv paths", "[source-path-state][scope]") {
    auto env = fold::SymbolicEnv::create();
    z3::expr cond = env.ctx().bool_const("cond");

    {
        auto scope = fold::scopedAssumption(env, cond);
        CHECK(checkPath(env.currentPath(), !cond) == z3::unsat);
    }

    CHECK(checkPath(env.currentPath(), cond) == z3::sat);
    CHECK(checkPath(env.currentPath(), !cond) == z3::sat);
}

TEST_CASE("ScopedSourcePath can be restored explicitly", "[source-path-state][scope]") {
    auto env = fold::SymbolicEnv::create();
    z3::expr cond = env.ctx().bool_const("cond");

    auto scope = fold::scopedAssumption(env, cond, false);
    CHECK(checkPath(env.currentPath(), cond) == z3::unsat);
    scope.restore();

    CHECK(checkPath(env.currentPath(), cond) == z3::sat);
    CHECK(checkPath(env.currentPath(), !cond) == z3::sat);
}