package rewriter_test

import (
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
	parser "github.com/milvus-io/milvus/internal/parser/planparserv2"
	"github.com/milvus-io/milvus/internal/parser/planparserv2/rewriter"
	"github.com/milvus-io/milvus/pkg/v3/proto/planpb"
	"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)

func buildSchemaHelperWithJSON(t *testing.T) *typeutil.SchemaHelper {
	fields := []*schemapb.FieldSchema{
		{FieldID: 101, Name: "Int64Field", DataType: schemapb.DataType_Int64},
		{FieldID: 102, Name: "JSONField", DataType: schemapb.DataType_JSON},
		{FieldID: 103, Name: "$meta", DataType: schemapb.DataType_JSON, IsDynamic: true},
		{FieldID: 104, Name: "NullableJSONField", DataType: schemapb.DataType_JSON, Nullable: true},
	}
	schema := &schemapb.CollectionSchema{
		Name:               "rewrite_json_test",
		AutoID:             false,
		Fields:             fields,
		EnableDynamicField: true,
	}
	helper, err := typeutil.CreateSchemaHelper(schema)
	require.NoError(t, err)
	return helper
}

func TestRewrite_JSON_NestedPath_NullableColumnInfo(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)

	expr, err := parser.ParseExpr(helper, `NullableJSONField["age"] in [1, 2]`, nil)
	require.NoError(t, err)
	require.NotNil(t, expr)
	term := expr.GetTermExpr()
	require.NotNil(t, term)
	require.True(t, term.GetColumnInfo().GetNullable())
	require.Equal(t, []string{"age"}, term.GetColumnInfo().GetNestedPath())
}

func TestRewrite_JSON_NestedPath_Nullable_TautologyKeepsPredicate(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)

	expr, err := parser.ParseExpr(helper, `NullableJSONField["age"] in [1, 2] or NullableJSONField["age"] != 1`, nil)
	require.NoError(t, err)
	require.NotNil(t, expr)
	require.False(t, rewriter.IsAlwaysTrueExpr(expr))
	require.NotNil(t, expr.GetBinaryExpr())
}

func TestRewrite_JSON_NestedPath_Nullable_ContradictionsKeepPredicate(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)

	for _, exprStr := range []string{
		`NullableJSONField["age"] in [1] and NullableJSONField["age"] == 2`,
		`NullableJSONField["age"] > 100 and NullableJSONField["age"] < 50`,
		`not (NullableJSONField["age"] in [1] and NullableJSONField["age"] == 2)`,
		`not (NullableJSONField["age"] > 100 and NullableJSONField["age"] < 50)`,
	} {
		expr, err := parser.ParseExpr(helper, exprStr, nil)
		require.NoError(t, err, exprStr)
		require.NotNil(t, expr, exprStr)
		require.False(t, rewriter.IsAlwaysFalseExpr(expr), "nullable JSON path must not fold to valid false: %s", exprStr)
		require.False(t, rewriter.IsAlwaysTrueExpr(expr), "nullable JSON path under NOT must not fold to valid true: %s", exprStr)
	}
}

func TestRewrite_JSON_NestedPath_NonNullable_ContradictionsKeepPredicate(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)

	for _, exprStr := range []string{
		`JSONField["age"] in [1] and JSONField["age"] == 2`,
		`JSONField["age"] > 100 and JSONField["age"] < 50`,
		`not (JSONField["age"] in [1] and JSONField["age"] == 2)`,
		`not (JSONField["age"] > 100 and JSONField["age"] < 50)`,
	} {
		expr, err := parser.ParseExpr(helper, exprStr, nil)
		require.NoError(t, err, exprStr)
		require.NotNil(t, expr, exprStr)
		require.False(t, rewriter.IsAlwaysFalseExpr(expr), "JSON path must not fold to valid false when the path can be missing: %s", exprStr)
		require.False(t, rewriter.IsAlwaysTrueExpr(expr), "JSON path under NOT must not fold to valid true when the path can be missing: %s", exprStr)
	}
}

func TestRewrite_JSON_RootValue_DynamicTypeKeepsThreeValuedPredicates(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)

	for _, exprStr := range []string{
		`JSONField in [1, 2] or JSONField != 1`,
		`JSONField in [1, 2] and JSONField == 3`,
		`JSONField > 100 and JSONField < 50`,
		`not (JSONField > 100 and JSONField < 50)`,
	} {
		expr, err := parser.ParseExpr(helper, exprStr, nil)
		require.NoError(t, err, exprStr)
		require.NotNil(t, expr, exprStr)
		require.False(t, rewriter.IsAlwaysFalseExpr(expr),
			"JSON type mismatch can be UNKNOWN, so the predicate must not fold to false: %s", exprStr)
		require.False(t, rewriter.IsAlwaysTrueExpr(expr),
			"JSON type mismatch can be UNKNOWN, so the predicate must not fold to true: %s", exprStr)
	}
}

func TestRewrite_JSON_InWithRangeUsesLiteralType(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)

	expr, err := parser.ParseExpr(helper,
		`JSONField["v"] in [1, 2, 3] and JSONField["v"] >= 2`, nil)
	require.NoError(t, err)
	term := expr.GetTermExpr()
	require.NotNil(t, term)
	require.Len(t, term.GetValues(), 2)
	require.Equal(t, []int64{2, 3}, []int64{
		term.GetValues()[0].GetInt64Val(),
		term.GetValues()[1].GetInt64Val(),
	})
}

func TestRewrite_JSON_InWithRangePreservesLargeInt64Precision(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)

	expr, err := parser.ParseExpr(helper,
		`JSONField["v"] in [9007199254740992, 9007199254740993] and JSONField["v"] >= 9007199254740993`, nil)
	require.NoError(t, err)
	term := expr.GetTermExpr()
	require.NotNil(t, term)
	require.Len(t, term.GetValues(), 1)
	require.Equal(t, int64(9007199254740993), term.GetValues()[0].GetInt64Val())
}

func TestRewrite_JSON_InWithRangeTypeMismatchSkipsOptimization(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)

	expr, err := parser.ParseExpr(helper,
		`JSONField["v"] in [1, 2] and JSONField["v"] >= "2"`, nil)
	require.NoError(t, err)
	require.NotNil(t, expr.GetBinaryExpr())
	require.NotNil(t, findTermExpr(expr))
	require.NotNil(t, findUnaryRangeExpr(expr, planpb.OpType_GreaterEqual))
}

func TestRewrite_JSON_NestedPath_ScalarNotEqualRewriteAllowed(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)

	for _, exprStr := range []string{
		`not (JSONField["age"] == 1)`,
		`JSONField["age"] not in [1]`,
	} {
		expr, err := parser.ParseExpr(helper, exprStr, nil)
		require.NoError(t, err, exprStr)
		require.NotNil(t, expr, exprStr)
		unaryRange := expr.GetUnaryRangeExpr()
		require.NotNil(t, unaryRange, "scalar JSON NOT equality should become !=: %s", exprStr)
		require.Equal(t, planpb.OpType_NotEqual, unaryRange.GetOp(), exprStr)
	}
}

func TestRewrite_JSON_NestedPath_ArrayNotEqualRewriteAllowed(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)

	for _, exprStr := range []string{
		`not (JSONField["arr"] == [1, 2])`,
		`JSONField["arr"] not in [[1, 2]]`,
	} {
		expr, err := parser.ParseExpr(helper, exprStr, nil)
		require.NoError(t, err, exprStr)
		require.NotNil(t, expr, exprStr)
		unaryRange := expr.GetUnaryRangeExpr()
		require.NotNil(t, unaryRange, "array-valued JSON NOT equality should become !=: %s", exprStr)
		require.Equal(t, planpb.OpType_NotEqual, unaryRange.GetOp(), exprStr)
	}
}

func TestRewrite_JSON_NestedPath_ArrayNotEqualsKeepPredicates(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)

	expr, err := parser.ParseExpr(helper, `JSONField["arr"] != [1, 2] and JSONField["arr"] != [3, 4]`, nil)
	require.NoError(t, err)
	require.NotNil(t, expr)
	require.NotNil(t, expr.GetBinaryExpr(), "array-valued JSON != chain must not become NOT(IN)")
}

// Test JSON field with int comparison - AND tightening
func TestRewrite_JSON_Int_AND_Strengthen(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)
	expr, err := parser.ParseExpr(helper, `JSONField["price"] > 10 and JSONField["price"] > 20`, nil)
	require.NoError(t, err)
	require.NotNil(t, expr)
	ure := expr.GetUnaryRangeExpr()
	require.NotNil(t, ure, "should merge to single unary range")
	require.Equal(t, planpb.OpType_GreaterThan, ure.GetOp())
	require.Equal(t, int64(20), ure.GetValue().GetInt64Val())
	// Verify column info
	require.Equal(t, schemapb.DataType_JSON, ure.GetColumnInfo().GetDataType())
	require.Equal(t, []string{"price"}, ure.GetColumnInfo().GetNestedPath())
}

func TestRewrite_JSON_IntRangePreservesLargeInt64Precision(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)

	for _, exprStr := range []string{
		`JSONField["v"] >= 9007199254740992 and JSONField["v"] >= 9007199254740993`,
		`JSONField["v"] >= 9007199254740993 and JSONField["v"] >= 9007199254740992`,
		`JSONField["v"] >= 9007199254740992.0 and JSONField["v"] >= 9007199254740993`,
	} {
		expr, err := parser.ParseExpr(helper, exprStr, nil)
		require.NoError(t, err, exprStr)
		unaryRange := expr.GetUnaryRangeExpr()
		require.NotNil(t, unaryRange, exprStr)
		require.Equal(t, planpb.OpType_GreaterEqual, unaryRange.GetOp(), exprStr)
		require.Equal(t, int64(9007199254740993), unaryRange.GetValue().GetInt64Val(), exprStr)
	}
}

func TestRewrite_JSON_IntOrRangePreservesLargeInt64Precision(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)

	for _, exprStr := range []string{
		`JSONField["v"] >= 9007199254740993 or JSONField["v"] >= 9007199254740992`,
		`JSONField["v"] >= 9007199254740992 or JSONField["v"] >= 9007199254740993`,
	} {
		expr, err := parser.ParseExpr(helper, exprStr, nil)
		require.NoError(t, err, exprStr)
		unaryRange := expr.GetUnaryRangeExpr()
		require.NotNil(t, unaryRange, exprStr)
		require.Equal(t, planpb.OpType_GreaterEqual, unaryRange.GetOp(), exprStr)
		require.Equal(t, int64(9007199254740992), unaryRange.GetValue().GetInt64Val(), exprStr)
	}
}

func TestRewrite_JSON_AndBinaryRangesPreservesLargeInt64Precision(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)
	expr, err := parser.ParseExpr(helper,
		`(JSONField["v"] >= 9007199254740992 and JSONField["v"] <= 9007199254740996) and `+
			`(JSONField["v"] >= 9007199254740993 and JSONField["v"] <= 9007199254740995)`, nil)
	require.NoError(t, err)

	binaryRange := expr.GetBinaryRangeExpr()
	require.NotNil(t, binaryRange)
	require.Equal(t, int64(9007199254740993), binaryRange.GetLowerValue().GetInt64Val())
	require.Equal(t, int64(9007199254740995), binaryRange.GetUpperValue().GetInt64Val())
}

func TestRewrite_JSON_OrBinaryRangesPreservesLargeInt64Precision(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)
	expr, err := parser.ParseExpr(helper,
		`(JSONField["v"] >= 9007199254740993 and JSONField["v"] <= 9007199254740995) or `+
			`(JSONField["v"] >= 9007199254740992 and JSONField["v"] <= 9007199254740996)`, nil)
	require.NoError(t, err)

	binaryRange := expr.GetBinaryRangeExpr()
	require.NotNil(t, binaryRange)
	require.Equal(t, int64(9007199254740992), binaryRange.GetLowerValue().GetInt64Val())
	require.Equal(t, int64(9007199254740996), binaryRange.GetUpperValue().GetInt64Val())
}

// Test JSON field with int comparison - AND to BinaryRange
func TestRewrite_JSON_Int_AND_ToBinaryRange(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)
	expr, err := parser.ParseExpr(helper, `JSONField["price"] > 10 and JSONField["price"] < 50`, nil)
	require.NoError(t, err)
	require.NotNil(t, expr)
	bre := expr.GetBinaryRangeExpr()
	require.NotNil(t, bre, "should create binary range")
	require.Equal(t, false, bre.GetLowerInclusive())
	require.Equal(t, false, bre.GetUpperInclusive())
	require.Equal(t, int64(10), bre.GetLowerValue().GetInt64Val())
	require.Equal(t, int64(50), bre.GetUpperValue().GetInt64Val())
	// Verify column info
	require.Equal(t, schemapb.DataType_JSON, bre.GetColumnInfo().GetDataType())
	require.Equal(t, []string{"price"}, bre.GetColumnInfo().GetNestedPath())
}

// Test JSON field with float comparison - AND tightening with mixed int/float
func TestRewrite_JSON_Float_AND_Strengthen_Mixed(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)
	expr, err := parser.ParseExpr(helper, `JSONField["score"] > 10 and JSONField["score"] > 15.5`, nil)
	require.NoError(t, err)
	require.NotNil(t, expr)
	ure := expr.GetUnaryRangeExpr()
	require.NotNil(t, ure)
	require.Equal(t, planpb.OpType_GreaterThan, ure.GetOp())
	require.InDelta(t, 15.5, ure.GetValue().GetFloatVal(), 1e-9)
}

// Test JSON field with string comparison - AND tightening
func TestRewrite_JSON_String_AND_Strengthen(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)
	expr, err := parser.ParseExpr(helper, `JSONField["name"] > "alice" and JSONField["name"] > "bob"`, nil)
	require.NoError(t, err)
	require.NotNil(t, expr)
	ure := expr.GetUnaryRangeExpr()
	require.NotNil(t, ure)
	require.Equal(t, planpb.OpType_GreaterThan, ure.GetOp())
	require.Equal(t, "bob", ure.GetValue().GetStringVal())
}

// Test JSON field with string comparison - AND to BinaryRange
func TestRewrite_JSON_String_AND_ToBinaryRange(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)
	expr, err := parser.ParseExpr(helper, `JSONField["name"] >= "alice" and JSONField["name"] <= "zebra"`, nil)
	require.NoError(t, err)
	require.NotNil(t, expr)
	bre := expr.GetBinaryRangeExpr()
	require.NotNil(t, bre)
	require.Equal(t, true, bre.GetLowerInclusive())
	require.Equal(t, true, bre.GetUpperInclusive())
	require.Equal(t, "alice", bre.GetLowerValue().GetStringVal())
	require.Equal(t, "zebra", bre.GetUpperValue().GetStringVal())
}

// Test JSON field with OR - weakening lower bounds
func TestRewrite_JSON_Int_OR_Weaken(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)
	expr, err := parser.ParseExpr(helper, `JSONField["age"] > 10 or JSONField["age"] > 20`, nil)
	require.NoError(t, err)
	require.NotNil(t, expr)
	ure := expr.GetUnaryRangeExpr()
	require.NotNil(t, ure)
	require.Equal(t, planpb.OpType_GreaterThan, ure.GetOp())
	require.Equal(t, int64(10), ure.GetValue().GetInt64Val())
}

// Test JSON field with OR - weakening upper bounds
func TestRewrite_JSON_String_OR_Weaken_Upper(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)
	expr, err := parser.ParseExpr(helper, `JSONField["category"] < "electronics" or JSONField["category"] < "sports"`, nil)
	require.NoError(t, err)
	require.NotNil(t, expr)
	ure := expr.GetUnaryRangeExpr()
	require.NotNil(t, ure)
	require.Equal(t, planpb.OpType_LessThan, ure.GetOp())
	require.Equal(t, "sports", ure.GetValue().GetStringVal())
}

// Test JSON field with numeric types (int and float) - SHOULD merge
func TestRewrite_JSON_NumericTypes_Merge(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)
	// JSONField["value"] > 10.5 (float) and JSONField["value"] > 20 (int)
	// Numeric types should merge - both treated as Double
	expr, err := parser.ParseExpr(helper, `JSONField["value"] > 10.5 and JSONField["value"] > 20`, nil)
	require.NoError(t, err)
	require.NotNil(t, expr)
	ure := expr.GetUnaryRangeExpr()
	require.NotNil(t, ure, "numeric types should merge to single predicate")
	require.Equal(t, planpb.OpType_GreaterThan, ure.GetOp())
	// Should pick the stronger bound (20)
	switch ure.GetValue().GetVal().(type) {
	case *planpb.GenericValue_Int64Val:
		require.Equal(t, int64(20), ure.GetValue().GetInt64Val())
	case *planpb.GenericValue_FloatVal:
		require.InDelta(t, 20.0, ure.GetValue().GetFloatVal(), 1e-9)
	default:
		t.Fatalf("unexpected value type")
	}
}

// Test JSON field with mixed type categories - should NOT merge
// Note: Numeric types (int and float) ARE compatible, but numeric and string are not
func TestRewrite_JSON_MixedTypes_NoMerge(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)
	// JSONField["value"] > 10 (numeric) and JSONField["value"] > "hello" (string)
	// These should remain separate as they have different type categories
	expr, err := parser.ParseExpr(helper, `JSONField["value"] > 10 and JSONField["value"] > "hello"`, nil)
	require.NoError(t, err)
	require.NotNil(t, expr)
	be := expr.GetBinaryExpr()
	require.NotNil(t, be, "should remain as binary expr (AND)")
	require.Equal(t, planpb.BinaryExpr_LogicalAnd, be.GetOp())
	// Both sides should be UnaryRangeExpr
	require.NotNil(t, be.GetLeft().GetUnaryRangeExpr())
	require.NotNil(t, be.GetRight().GetUnaryRangeExpr())
}

// Test JSON field with mixed type categories in OR - should NOT merge
func TestRewrite_JSON_MixedTypes_OR_NoMerge(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)
	expr, err := parser.ParseExpr(helper, `JSONField["data"] > 100 or JSONField["data"] > "text"`, nil)
	require.NoError(t, err)
	require.NotNil(t, expr)
	be := expr.GetBinaryExpr()
	require.NotNil(t, be, "should remain as binary expr (OR)")
	require.Equal(t, planpb.BinaryExpr_LogicalOr, be.GetOp())
}

// Test dynamic field ($meta) - same as JSON field
func TestRewrite_DynamicField_Int_AND_Strengthen(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)
	// $meta is the dynamic field
	expr, err := parser.ParseExpr(helper, `$meta["count"] > 5 and $meta["count"] > 15`, nil)
	require.NoError(t, err)
	require.NotNil(t, expr)
	ure := expr.GetUnaryRangeExpr()
	require.NotNil(t, ure)
	require.Equal(t, planpb.OpType_GreaterThan, ure.GetOp())
	require.Equal(t, int64(15), ure.GetValue().GetInt64Val())
}

// Test dynamic field - AND to BinaryRange
func TestRewrite_DynamicField_Float_AND_ToBinaryRange(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)
	expr, err := parser.ParseExpr(helper, `$meta["rating"] >= 1.0 and $meta["rating"] <= 5.0`, nil)
	require.NoError(t, err)
	require.NotNil(t, expr)
	bre := expr.GetBinaryRangeExpr()
	require.NotNil(t, bre)
	require.Equal(t, true, bre.GetLowerInclusive())
	require.Equal(t, true, bre.GetUpperInclusive())
	require.InDelta(t, 1.0, bre.GetLowerValue().GetFloatVal(), 1e-9)
	require.InDelta(t, 5.0, bre.GetUpperValue().GetFloatVal(), 1e-9)
}

// Test different nested paths - should NOT merge
func TestRewrite_JSON_DifferentPaths_NoMerge(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)
	expr, err := parser.ParseExpr(helper, `JSONField["a"] > 10 and JSONField["b"] > 20`, nil)
	require.NoError(t, err)
	require.NotNil(t, expr)
	be := expr.GetBinaryExpr()
	require.NotNil(t, be, "different paths should not merge")
	require.Equal(t, planpb.BinaryExpr_LogicalAnd, be.GetOp())
}

// Test nested JSON path
func TestRewrite_JSON_NestedPath_AND_Strengthen(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)
	expr, err := parser.ParseExpr(helper, `JSONField["user"]["age"] > 18 and JSONField["user"]["age"] > 21`, nil)
	require.NoError(t, err)
	require.NotNil(t, expr)
	ure := expr.GetUnaryRangeExpr()
	require.NotNil(t, ure)
	require.Equal(t, planpb.OpType_GreaterThan, ure.GetOp())
	require.Equal(t, int64(21), ure.GetValue().GetInt64Val())
	require.Equal(t, []string{"user", "age"}, ure.GetColumnInfo().GetNestedPath())
}

// Test inclusive vs exclusive bounds
func TestRewrite_JSON_AND_EquivalentBounds(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)
	// JSONField["x"] >= 10 AND JSONField["x"] > 10 → JSONField["x"] > 10 (exclusive is stronger)
	expr, err := parser.ParseExpr(helper, `JSONField["x"] >= 10 and JSONField["x"] > 10`, nil)
	require.NoError(t, err)
	require.NotNil(t, expr)
	ure := expr.GetUnaryRangeExpr()
	require.NotNil(t, ure)
	require.Equal(t, planpb.OpType_GreaterThan, ure.GetOp())
	require.Equal(t, int64(10), ure.GetValue().GetInt64Val())
}

// Test OR with inclusive bounds preference
func TestRewrite_JSON_OR_EquivalentBounds(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)
	// JSONField["x"] >= 10 OR JSONField["x"] > 10 → JSONField["x"] >= 10 (inclusive is weaker)
	expr, err := parser.ParseExpr(helper, `JSONField["x"] >= 10 or JSONField["x"] > 10`, nil)
	require.NoError(t, err)
	require.NotNil(t, expr)
	ure := expr.GetUnaryRangeExpr()
	require.NotNil(t, ure)
	require.Equal(t, planpb.OpType_GreaterEqual, ure.GetOp())
	require.Equal(t, int64(10), ure.GetValue().GetInt64Val())
}

// Test that scalar field and JSON field don't interfere
func TestRewrite_JSON_And_Scalar_Independent(t *testing.T) {
	helper := buildSchemaHelperWithJSON(t)
	expr, err := parser.ParseExpr(helper, `Int64Field > 10 and Int64Field > 20 and JSONField["price"] > 5 and JSONField["price"] > 15`, nil)
	require.NoError(t, err)
	require.NotNil(t, expr)
	// Should result in AND of two optimized predicates
	be := expr.GetBinaryExpr()
	require.NotNil(t, be)
	require.Equal(t, planpb.BinaryExpr_LogicalAnd, be.GetOp())

	// Collect all predicates
	var predicates []*planpb.Expr
	var collect func(*planpb.Expr)
	collect = func(e *planpb.Expr) {
		if be := e.GetBinaryExpr(); be != nil && be.GetOp() == planpb.BinaryExpr_LogicalAnd {
			collect(be.GetLeft())
			collect(be.GetRight())
		} else {
			predicates = append(predicates, e)
		}
	}
	collect(expr)

	require.Equal(t, 2, len(predicates), "should have two optimized predicates")

	// Check that both are optimized to the stronger bounds
	var hasInt64, hasJSON bool
	for _, p := range predicates {
		if ure := p.GetUnaryRangeExpr(); ure != nil {
			col := ure.GetColumnInfo()
			if col.GetDataType() == schemapb.DataType_Int64 {
				hasInt64 = true
				require.Equal(t, int64(20), ure.GetValue().GetInt64Val())
			} else if col.GetDataType() == schemapb.DataType_JSON {
				hasJSON = true
				require.Equal(t, int64(15), ure.GetValue().GetInt64Val())
			}
		}
	}
	require.True(t, hasInt64, "should have optimized Int64Field predicate")
	require.True(t, hasJSON, "should have optimized JSONField predicate")
}
