#include <catch2/catch_test_macros.hpp>
#include <memory>
#include <vector>

#include "fold/ir/expr.hpp"
#include "fold/ir/expr_json.hpp"
#include "fold/symbolic/ai_pre_pass.hpp"

using namespace fold::ir;
using namespace fold::symbolic;

namespace {

// Helper to build a simple RustExprNode tree.
struct NodeBuilder {
    RustNodeArena& arena;

    explicit NodeBuilder(RustNodeArena& a) : arena(a) {}

    RustExprNode* integer(std::int64_t v, std::uint32_t width = 32, bool is_signed = true) {
        auto* n = arena.alloc();
        n->kind = Expr::Kind::IntegerLiteral;
        n->int_value = v;
        n->int_width = width;
        n->int_signed = is_signed;
        return n;
    }

    RustExprNode* binary(BinaryOpcode op, RustExprNode* lhs, RustExprNode* rhs,
                         std::uint32_t width = 32, bool is_signed = true) {
        auto* n = arena.alloc();
        n->kind = Expr::Kind::BinaryOp;
        n->binary_op = op;
        n->binary_lhs = lhs;
        n->binary_rhs = rhs;
        n->int_width = width;
        n->int_signed = is_signed;
        return n;
    }

    RustExprNode* unary(UnaryOpcode op, RustExprNode* sub) {
        auto* n = arena.alloc();
        n->kind = Expr::Kind::UnaryOp;
        n->unary_op = op;
        n->unary_sub = sub;
        return n;
    }

    RustExprNode* declref(const std::string& name, std::uint32_t width = 32,
                          bool is_signed = true) {
        auto* n = arena.alloc();
        n->kind = Expr::Kind::DeclRef;
        n->declref_symbol = name;
        n->int_width = width;
        n->int_signed = is_signed;
        return n;
    }
};

} // namespace

TEST_CASE("AI pre-pass: constant add within range is safe", "[ai_pre_pass]") {
    RustNodeArena arena;
    NodeBuilder b(arena);

    // 10 + 20 in i32 → safe (30 fits in i32).
    auto* root = b.binary(BinaryOpcode::Add, b.integer(10, 32), b.integer(20, 32), 32);

    AiPrePass pp;
    pp.analyze(root);

    REQUIRE(pp.totalArithmeticNodes() == 1);
    REQUIRE(pp.nodesProvedSafe() == 1);
    REQUIRE(pp.nodesNeedingZ3() == 0);
    REQUIRE(pp.isArithmeticSafe(root, 32, true));
}

TEST_CASE("AI pre-pass: constant mul within range is safe", "[ai_pre_pass]") {
    RustNodeArena arena;
    NodeBuilder b(arena);

    // 10 * 20 in i32 → safe (200 fits in i32).
    auto* root = b.binary(BinaryOpcode::Mul, b.integer(10, 32), b.integer(20, 32), 32);

    AiPrePass pp;
    pp.analyze(root);

    REQUIRE(pp.totalArithmeticNodes() == 1);
    REQUIRE(pp.nodesProvedSafe() == 1);
    REQUIRE(pp.isArithmeticSafe(root, 32, true));
}

TEST_CASE("AI pre-pass: constant sub within range is safe", "[ai_pre_pass]") {
    RustNodeArena arena;
    NodeBuilder b(arena);

    // 100 - 50 in i32 → safe (50 fits in i32).
    auto* root = b.binary(BinaryOpcode::Sub, b.integer(100, 32), b.integer(50, 32), 32);

    AiPrePass pp;
    pp.analyze(root);

    REQUIRE(pp.totalArithmeticNodes() == 1);
    REQUIRE(pp.nodesProvedSafe() == 1);
    REQUIRE(pp.isArithmeticSafe(root, 32, true));
}

TEST_CASE("AI pre-pass: overflow add is NOT proved safe", "[ai_pre_pass]") {
    RustNodeArena arena;
    NodeBuilder b(arena);

    // INT32_MAX + 1 → overflow, should NOT be proved safe.
    auto* root = b.binary(BinaryOpcode::Add, b.integer(2147483647LL, 32), b.integer(1, 32), 32);

    AiPrePass pp;
    pp.analyze(root);

    REQUIRE(pp.totalArithmeticNodes() == 1);
    REQUIRE(pp.nodesProvedSafe() == 0);
    REQUIRE(pp.nodesNeedingZ3() == 1);
    REQUIRE_FALSE(pp.isArithmeticSafe(root, 32, true));
}

TEST_CASE("AI pre-pass: variable operand returns unknown (needs Z3)", "[ai_pre_pass]") {
    RustNodeArena arena;
    NodeBuilder b(arena);

    // x + 10 where x is a variable → cannot bound x → NOT proved safe.
    auto* root = b.binary(BinaryOpcode::Add, b.declref("x"), b.integer(10, 32), 32);

    AiPrePass pp;
    pp.analyze(root);

    REQUIRE(pp.totalArithmeticNodes() == 1);
    REQUIRE(pp.nodesProvedSafe() == 0);
    REQUIRE(pp.nodesNeedingZ3() == 1);
}

TEST_CASE("AI pre-pass: unsigned add within range is safe", "[ai_pre_pass]") {
    RustNodeArena arena;
    NodeBuilder b(arena);

    // u32: 100 + 200 → safe (300 < UINT32_MAX).
    auto* root = b.binary(BinaryOpcode::Add, b.integer(100, 32, false), b.integer(200, 32, false),
                          32, false);

    AiPrePass pp;
    pp.analyze(root);

    REQUIRE(pp.totalArithmeticNodes() == 1);
    REQUIRE(pp.nodesProvedSafe() == 1);
    REQUIRE(pp.isArithmeticSafe(root, 32, false));
}

TEST_CASE("AI pre-pass: negative sub overflow is NOT proved safe", "[ai_pre_pass]") {
    RustNodeArena arena;
    NodeBuilder b(arena);

    // INT32_MIN - 1 → overflow.
    auto* root = b.binary(BinaryOpcode::Sub, b.integer(-2147483648LL, 32), b.integer(1, 32), 32);

    AiPrePass pp;
    pp.analyze(root);

    REQUIRE(pp.totalArithmeticNodes() == 1);
    REQUIRE_FALSE(pp.isArithmeticSafe(root, 32, true));
}

TEST_CASE("AI pre-pass: nested arithmetic both safe", "[ai_pre_pass]") {
    RustNodeArena arena;
    NodeBuilder b(arena);

    // (10 + 20) * (30 - 5) → 30 * 25 = 750, fits in i32.
    auto* lhs = b.binary(BinaryOpcode::Add, b.integer(10, 32), b.integer(20, 32), 32);
    auto* rhs = b.binary(BinaryOpcode::Sub, b.integer(30, 32), b.integer(5, 32), 32);
    auto* root = b.binary(BinaryOpcode::Mul, lhs, rhs, 32);

    AiPrePass pp;
    pp.analyze(root);

    REQUIRE(pp.totalArithmeticNodes() == 3);
    REQUIRE(pp.nodesProvedSafe() == 3);
    REQUIRE(pp.isArithmeticSafe(root, 32, true));
    REQUIRE(pp.isArithmeticSafe(lhs, 32, true));
    REQUIRE(pp.isArithmeticSafe(rhs, 32, true));
}

TEST_CASE("AI pre-pass: mixed precision i8 add within range", "[ai_pre_pass]") {
    RustNodeArena arena;
    NodeBuilder b(arena);

    // 50 + 50 in i8 → safe (100 < 127).
    auto* root = b.binary(BinaryOpcode::Add, b.integer(50, 8), b.integer(50, 8), 8);

    AiPrePass pp;
    pp.analyze(root);

    REQUIRE(pp.totalArithmeticNodes() == 1);
    REQUIRE(pp.nodesProvedSafe() == 1);
    REQUIRE(pp.isArithmeticSafe(root, 8, true));
}

TEST_CASE("AI pre-pass: i8 overflow NOT proved safe", "[ai_pre_pass]") {
    RustNodeArena arena;
    NodeBuilder b(arena);

    // 100 + 100 in i8 → 200 > 127, overflow.
    auto* root = b.binary(BinaryOpcode::Add, b.integer(100, 8), b.integer(100, 8), 8);

    AiPrePass pp;
    pp.analyze(root);

    REQUIRE(pp.totalArithmeticNodes() == 1);
    REQUIRE_FALSE(pp.isArithmeticSafe(root, 8, true));
}

TEST_CASE("AI pre-pass: clear resets state", "[ai_pre_pass]") {
    RustNodeArena arena;
    NodeBuilder b(arena);

    auto* root = b.binary(BinaryOpcode::Add, b.integer(10, 32), b.integer(20, 32), 32);

    AiPrePass pp;
    pp.analyze(root);
    REQUIRE(pp.totalArithmeticNodes() == 1);

    pp.clear();
    REQUIRE(pp.totalArithmeticNodes() == 0);
    REQUIRE(pp.nodesProvedSafe() == 0);
}

TEST_CASE("AI pre-pass: non-arithmetic binary ops do not count", "[ai_pre_pass]") {
    RustNodeArena arena;
    NodeBuilder b(arena);

    // 10 == 20 is not arithmetic (comparison, no overflow risk).
    auto* root = b.binary(BinaryOpcode::EQ, b.integer(10, 32), b.integer(20, 32), 32);

    AiPrePass pp;
    pp.analyze(root);

    // Only Add/Sub/Mul count. EQ is ignored.
    REQUIRE(pp.totalArithmeticNodes() == 0);
}

TEST_CASE("AI pre-pass: intervalFor returns correct interval", "[ai_pre_pass]") {
    RustNodeArena arena;
    NodeBuilder b(arena);

    auto* lit42 = b.integer(42, 32);
    auto* lit7 = b.integer(7, 32);
    auto* root = b.binary(BinaryOpcode::Add, lit42, lit7, 32);

    AiPrePass pp;
    pp.analyze(root);

    auto iv42 = pp.intervalFor(lit42);
    REQUIRE_FALSE(iv42.top);
    REQUIRE(iv42.lo == 42);
    REQUIRE(iv42.hi == 42);

    auto iv7 = pp.intervalFor(lit7);
    REQUIRE_FALSE(iv7.top);
    REQUIRE(iv7.lo == 7);
    REQUIRE(iv7.hi == 7);
}

TEST_CASE("AI pre-pass: unary minus on constant interval is safe context", "[ai_pre_pass]") {
    RustNodeArena arena;
    NodeBuilder b(arena);

    // -(10) + 30 → safe.
    auto* neg = b.unary(UnaryOpcode::Minus, b.integer(10, 32));
    auto* root = b.binary(BinaryOpcode::Add, neg, b.integer(30, 32), 32);

    AiPrePass pp;
    pp.analyze(root);

    // The unary minus itself is not counted as arithmetic, but the add is.
    REQUIRE(pp.totalArithmeticNodes() == 1);
    // -10 + 30 = 20, fits in i32 → safe.
    REQUIRE(pp.isArithmeticSafe(root, 32, true));
}

TEST_CASE("AI pre-pass: nullptr node returns safe defaults", "[ai_pre_pass]") {
    AiPrePass pp;
    pp.analyze(nullptr);
    REQUIRE(pp.totalArithmeticNodes() == 0);
    REQUIRE(pp.intervalFor(nullptr).top == true);
    REQUIRE_FALSE(pp.isArithmeticSafe(nullptr, 32, true));
}