// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#include <gtest/gtest.h>
#include "arrow/memory_pool.h"
#include "arrow/status.h"

#include "gandiva/projector.h"
#include "gandiva/tests/test_util.h"
#include "gandiva/tree_expr_builder.h"

namespace gandiva {

using arrow::boolean;
using arrow::int32;

class TestBooleanExpr : public ::testing::Test {
 public:
  void SetUp() { pool_ = arrow::default_memory_pool(); }

 protected:
  arrow::MemoryPool* pool_;
};

TEST_F(TestBooleanExpr, SimpleAnd) {
  // schema for input fields
  auto fielda = field("a", int32());
  auto fieldb = field("b", int32());
  auto schema = arrow::schema({fielda, fieldb});

  // output fields
  auto field_result = field("res", boolean());

  // build expression.
  // (a > 0) && (b > 0)
  auto node_a = TreeExprBuilder::MakeField(fielda);
  auto node_b = TreeExprBuilder::MakeField(fieldb);
  auto literal_0 = TreeExprBuilder::MakeLiteral((int32_t)0);
  auto a_gt_0 =
      TreeExprBuilder::MakeFunction("greater_than", {node_a, literal_0}, boolean());
  auto b_gt_0 =
      TreeExprBuilder::MakeFunction("greater_than", {node_b, literal_0}, boolean());

  auto node_and = TreeExprBuilder::MakeAnd({a_gt_0, b_gt_0});
  auto expr = TreeExprBuilder::MakeExpression(node_and, field_result);

  // Build a projector for the expressions.
  std::shared_ptr<Projector> projector;
  auto status = Projector::Make(schema, {expr}, TestConfiguration(), &projector);
  EXPECT_TRUE(status.ok());

  // FALSE_VALID && ?  => FALSE_VALID
  int num_records = 4;
  auto arraya = MakeArrowArrayInt32({-2, -2, -2, -2}, {true, true, true, true});
  auto arrayb = MakeArrowArrayInt32({-2, -2, 2, 2}, {true, false, true, false});
  auto exp = MakeArrowArrayBool({false, false, false, false}, {true, true, true, true});
  auto in_batch = arrow::RecordBatch::Make(schema, num_records, {arraya, arrayb});

  arrow::ArrayVector outputs;
  status = projector->Evaluate(*in_batch, pool_, &outputs);
  EXPECT_TRUE(status.ok());
  EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0));

  // FALSE_INVALID && ?
  num_records = 4;
  arraya = MakeArrowArrayInt32({-2, -2, -2, -2}, {false, false, false, false});
  arrayb = MakeArrowArrayInt32({-2, -2, 2, 2}, {true, false, true, false});
  exp = MakeArrowArrayBool({false, false, false, false}, {true, false, false, false});
  in_batch = arrow::RecordBatch::Make(schema, num_records, {arraya, arrayb});
  outputs.clear();
  status = projector->Evaluate(*in_batch, pool_, &outputs);
  EXPECT_TRUE(status.ok());
  EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0));

  // TRUE_VALID && ?
  num_records = 4;
  arraya = MakeArrowArrayInt32({2, 2, 2, 2}, {true, true, true, true});
  arrayb = MakeArrowArrayInt32({-2, -2, 2, 2}, {true, false, true, false});
  exp = MakeArrowArrayBool({false, false, true, false}, {true, false, true, false});
  in_batch = arrow::RecordBatch::Make(schema, num_records, {arraya, arrayb});
  outputs.clear();
  status = projector->Evaluate(*in_batch, pool_, &outputs);
  EXPECT_TRUE(status.ok());
  EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0));

  // TRUE_INVALID && ?
  num_records = 4;
  arraya = MakeArrowArrayInt32({2, 2, 2, 2}, {false, false, false, false});
  arrayb = MakeArrowArrayInt32({-2, -2, 2, 2}, {true, false, true, false});
  exp = MakeArrowArrayBool({false, false, false, false}, {true, false, false, false});
  in_batch = arrow::RecordBatch::Make(schema, num_records, {arraya, arrayb});
  outputs.clear();
  status = projector->Evaluate(*in_batch, pool_, &outputs);
  EXPECT_TRUE(status.ok());
  EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0));
}

TEST_F(TestBooleanExpr, IsTrueFalse) {
  // schema for input fields
  auto field_a = field("a", boolean());
  auto schema = arrow::schema({field_a});

  // output fields
  auto field_result_1 = field("is_true", boolean());
  auto field_result_2 = field("is_not_false", boolean());
  auto field_result_3 = field("is_false", boolean());
  auto field_result_4 = field("is_not_true", boolean());

  // build expressions.
  auto is_true = TreeExprBuilder::MakeExpression("istrue", {field_a}, field_result_1);
  auto is_not_false =
      TreeExprBuilder::MakeExpression("isnotfalse", {field_a}, field_result_2);
  auto is_false = TreeExprBuilder::MakeExpression("isfalse", {field_a}, field_result_3);
  auto is_not_true =
      TreeExprBuilder::MakeExpression("isnottrue", {field_a}, field_result_4);

  // Build a projector for the expressions.
  std::shared_ptr<Projector> projector;
  auto status = Projector::Make(schema, {is_true, is_not_false, is_false, is_not_true},
                                TestConfiguration(), &projector);
  EXPECT_TRUE(status.ok());

  int num_records = 6;
  auto in_a = MakeArrowArrayBool({true, false, false, true, true, false},
                                 {true, true, true, true, false, false});

  auto expa = MakeArrowArrayBool({true, false, false, true, false, false},
                                 {true, true, true, true, true, true});
  auto expb = MakeArrowArrayBool({true, false, false, true, true, true},
                                 {true, true, true, true, true, true});
  auto expc = MakeArrowArrayBool({false, true, true, false, false, false},
                                 {true, true, true, true, true, true});
  auto expd = MakeArrowArrayBool({false, true, true, false, true, true},
                                 {true, true, true, true, true, true});

  auto in_batch = arrow::RecordBatch::Make(schema, num_records, {in_a});

  arrow::ArrayVector outputs;
  status = projector->Evaluate(*in_batch, pool_, &outputs);
  EXPECT_TRUE(status.ok());
  EXPECT_ARROW_ARRAY_EQUALS(expa, outputs.at(0));
  EXPECT_ARROW_ARRAY_EQUALS(expb, outputs.at(1));
  EXPECT_ARROW_ARRAY_EQUALS(expc, outputs.at(2));
  EXPECT_ARROW_ARRAY_EQUALS(expd, outputs.at(3));
}

TEST_F(TestBooleanExpr, SimpleOr) {
  // schema for input fields
  auto fielda = field("a", int32());
  auto fieldb = field("b", int32());
  auto schema = arrow::schema({fielda, fieldb});

  // output fields
  auto field_result = field("res", boolean());

  // build expression.
  // (a > 0) || (b > 0)
  auto node_a = TreeExprBuilder::MakeField(fielda);
  auto node_b = TreeExprBuilder::MakeField(fieldb);
  auto literal_0 = TreeExprBuilder::MakeLiteral((int32_t)0);
  auto a_gt_0 =
      TreeExprBuilder::MakeFunction("greater_than", {node_a, literal_0}, boolean());
  auto b_gt_0 =
      TreeExprBuilder::MakeFunction("greater_than", {node_b, literal_0}, boolean());

  auto node_or = TreeExprBuilder::MakeOr({a_gt_0, b_gt_0});
  auto expr = TreeExprBuilder::MakeExpression(node_or, field_result);

  // Build a projector for the expressions.
  std::shared_ptr<Projector> projector;
  auto status = Projector::Make(schema, {expr}, TestConfiguration(), &projector);
  EXPECT_TRUE(status.ok());

  // TRUE_VALID && ?  => TRUE_VALID
  int num_records = 4;
  auto arraya = MakeArrowArrayInt32({2, 2, 2, 2}, {true, true, true, true});
  auto arrayb = MakeArrowArrayInt32({-2, -2, 2, 2}, {true, false, true, false});
  auto exp = MakeArrowArrayBool({true, true, true, true}, {true, true, true, true});
  auto in_batch = arrow::RecordBatch::Make(schema, num_records, {arraya, arrayb});

  arrow::ArrayVector outputs;
  status = projector->Evaluate(*in_batch, pool_, &outputs);
  EXPECT_TRUE(status.ok());
  EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0));

  // TRUE_INVALID && ?
  num_records = 4;
  arraya = MakeArrowArrayInt32({2, 2, 2, 2}, {false, false, false, false});
  arrayb = MakeArrowArrayInt32({-2, -2, 2, 2}, {true, false, true, false});
  exp = MakeArrowArrayBool({false, false, true, false}, {false, false, true, false});
  in_batch = arrow::RecordBatch::Make(schema, num_records, {arraya, arrayb});
  outputs.clear();
  status = projector->Evaluate(*in_batch, pool_, &outputs);
  EXPECT_TRUE(status.ok());
  EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0));

  // FALSE_VALID && ?
  num_records = 4;
  arraya = MakeArrowArrayInt32({-2, -2, -2, -2}, {true, true, true, true});
  arrayb = MakeArrowArrayInt32({-2, -2, 2, 2}, {true, false, true, false});
  exp = MakeArrowArrayBool({false, false, true, false}, {true, false, true, false});
  in_batch = arrow::RecordBatch::Make(schema, num_records, {arraya, arrayb});
  outputs.clear();
  status = projector->Evaluate(*in_batch, pool_, &outputs);
  EXPECT_TRUE(status.ok());
  EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0));

  // FALSE_INVALID && ?
  num_records = 4;
  arraya = MakeArrowArrayInt32({-2, -2, -2, -2}, {false, false, false, false});
  arrayb = MakeArrowArrayInt32({-2, -2, 2, 2}, {true, false, true, false});
  exp = MakeArrowArrayBool({false, false, true, false}, {false, false, true, false});
  in_batch = arrow::RecordBatch::Make(schema, num_records, {arraya, arrayb});
  outputs.clear();
  status = projector->Evaluate(*in_batch, pool_, &outputs);
  EXPECT_TRUE(status.ok());
  EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0));
}

TEST_F(TestBooleanExpr, AndThree) {
  // schema for input fields
  auto fielda = field("a", int32());
  auto fieldb = field("b", int32());
  auto fieldc = field("c", int32());
  auto schema = arrow::schema({fielda, fieldb, fieldc});

  // output fields
  auto field_result = field("res", boolean());

  // build expression.
  // (a > 0) && (b > 0) && (c > 0)
  auto node_a = TreeExprBuilder::MakeField(fielda);
  auto node_b = TreeExprBuilder::MakeField(fieldb);
  auto node_c = TreeExprBuilder::MakeField(fieldc);
  auto literal_0 = TreeExprBuilder::MakeLiteral((int32_t)0);
  auto a_gt_0 =
      TreeExprBuilder::MakeFunction("greater_than", {node_a, literal_0}, boolean());
  auto b_gt_0 =
      TreeExprBuilder::MakeFunction("greater_than", {node_b, literal_0}, boolean());
  auto c_gt_0 =
      TreeExprBuilder::MakeFunction("greater_than", {node_c, literal_0}, boolean());

  auto node_and = TreeExprBuilder::MakeAnd({a_gt_0, b_gt_0, c_gt_0});
  auto expr = TreeExprBuilder::MakeExpression(node_and, field_result);

  // Build a projector for the expressions.
  std::shared_ptr<Projector> projector;
  auto status = Projector::Make(schema, {expr}, TestConfiguration(), &projector);
  EXPECT_TRUE(status.ok());

  int num_records = 8;
  std::vector<bool> validity({true, true, true, true, true, true, true, true});
  auto arraya = MakeArrowArrayInt32({2, 2, 2, 0, 2, 0, 0, 0}, validity);
  auto arrayb = MakeArrowArrayInt32({2, 2, 0, 2, 0, 2, 0, 0}, validity);
  auto arrayc = MakeArrowArrayInt32({2, 0, 2, 2, 0, 0, 2, 0}, validity);
  auto exp = MakeArrowArrayBool({true, false, false, false, false, false, false, false},
                                validity);

  auto in_batch = arrow::RecordBatch::Make(schema, num_records, {arraya, arrayb, arrayc});

  arrow::ArrayVector outputs;
  status = projector->Evaluate(*in_batch, pool_, &outputs);
  EXPECT_TRUE(status.ok());
  EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0));
}

TEST_F(TestBooleanExpr, OrThree) {
  // schema for input fields
  auto fielda = field("a", int32());
  auto fieldb = field("b", int32());
  auto fieldc = field("c", int32());
  auto schema = arrow::schema({fielda, fieldb, fieldc});

  // output fields
  auto field_result = field("res", boolean());

  // build expression.
  // (a > 0) || (b > 0) || (c > 0)
  auto node_a = TreeExprBuilder::MakeField(fielda);
  auto node_b = TreeExprBuilder::MakeField(fieldb);
  auto node_c = TreeExprBuilder::MakeField(fieldc);
  auto literal_0 = TreeExprBuilder::MakeLiteral((int32_t)0);
  auto a_gt_0 =
      TreeExprBuilder::MakeFunction("greater_than", {node_a, literal_0}, boolean());
  auto b_gt_0 =
      TreeExprBuilder::MakeFunction("greater_than", {node_b, literal_0}, boolean());
  auto c_gt_0 =
      TreeExprBuilder::MakeFunction("greater_than", {node_c, literal_0}, boolean());

  auto node_or = TreeExprBuilder::MakeOr({a_gt_0, b_gt_0, c_gt_0});
  auto expr = TreeExprBuilder::MakeExpression(node_or, field_result);

  // Build a projector for the expressions.
  std::shared_ptr<Projector> projector;
  auto status = Projector::Make(schema, {expr}, TestConfiguration(), &projector);
  EXPECT_TRUE(status.ok());

  int num_records = 8;
  std::vector<bool> validity({true, true, true, true, true, true, true, true});
  auto arraya = MakeArrowArrayInt32({2, 2, 2, 0, 2, 0, 0, 0}, validity);
  auto arrayb = MakeArrowArrayInt32({2, 2, 0, 2, 0, 2, 0, 0}, validity);
  auto arrayc = MakeArrowArrayInt32({2, 0, 2, 2, 0, 0, 2, 0}, validity);
  auto exp =
      MakeArrowArrayBool({true, true, true, true, true, true, true, false}, validity);

  auto in_batch = arrow::RecordBatch::Make(schema, num_records, {arraya, arrayb, arrayc});

  arrow::ArrayVector outputs;
  status = projector->Evaluate(*in_batch, pool_, &outputs);
  EXPECT_TRUE(status.ok());
  EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0));
}

TEST_F(TestBooleanExpr, BooleanAndInsideIf) {
  // schema for input fields
  auto fielda = field("a", int32());
  auto fieldb = field("b", int32());
  auto schema = arrow::schema({fielda, fieldb});

  // output fields
  auto field_result = field("res", boolean());

  // build expression.
  // if (a > 2 && b > 2)
  //   a > 3 && b > 3
  // else
  //   a > 1 && b > 1
  auto node_a = TreeExprBuilder::MakeField(fielda);
  auto node_b = TreeExprBuilder::MakeField(fieldb);
  auto literal_1 = TreeExprBuilder::MakeLiteral((int32_t)1);
  auto literal_2 = TreeExprBuilder::MakeLiteral((int32_t)2);
  auto literal_3 = TreeExprBuilder::MakeLiteral((int32_t)3);
  auto a_gt_1 =
      TreeExprBuilder::MakeFunction("greater_than", {node_a, literal_1}, boolean());
  auto a_gt_2 =
      TreeExprBuilder::MakeFunction("greater_than", {node_a, literal_2}, boolean());
  auto a_gt_3 =
      TreeExprBuilder::MakeFunction("greater_than", {node_a, literal_3}, boolean());
  auto b_gt_1 =
      TreeExprBuilder::MakeFunction("greater_than", {node_b, literal_1}, boolean());
  auto b_gt_2 =
      TreeExprBuilder::MakeFunction("greater_than", {node_b, literal_2}, boolean());
  auto b_gt_3 =
      TreeExprBuilder::MakeFunction("greater_than", {node_b, literal_3}, boolean());

  auto and_1 = TreeExprBuilder::MakeAnd({a_gt_1, b_gt_1});
  auto and_2 = TreeExprBuilder::MakeAnd({a_gt_2, b_gt_2});
  auto and_3 = TreeExprBuilder::MakeAnd({a_gt_3, b_gt_3});

  auto node_if = TreeExprBuilder::MakeIf(and_2, and_3, and_1, arrow::boolean());
  auto expr = TreeExprBuilder::MakeExpression(node_if, field_result);

  // Build a projector for the expressions.
  std::shared_ptr<Projector> projector;
  auto status = Projector::Make(schema, {expr}, TestConfiguration(), &projector);
  EXPECT_TRUE(status.ok());

  int num_records = 4;
  std::vector<bool> validity({true, true, true, true});
  auto arraya = MakeArrowArrayInt32({4, 4, 2, 1}, validity);
  auto arrayb = MakeArrowArrayInt32({5, 3, 3, 1}, validity);
  auto exp = MakeArrowArrayBool({true, false, true, false}, validity);

  auto in_batch = arrow::RecordBatch::Make(schema, num_records, {arraya, arrayb});

  arrow::ArrayVector outputs;
  status = projector->Evaluate(*in_batch, pool_, &outputs);
  EXPECT_TRUE(status.ok());
  EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0));
}

TEST_F(TestBooleanExpr, IfInsideBooleanAnd) {
  // schema for input fields
  auto fielda = field("a", int32());
  auto fieldb = field("b", int32());
  auto schema = arrow::schema({fielda, fieldb});

  // output fields
  auto field_result = field("res", boolean());

  // build expression.
  // (if (a > b) a > 3 else b > 3) && (if (a > b) a > 2 else b > 2)

  auto node_a = TreeExprBuilder::MakeField(fielda);
  auto node_b = TreeExprBuilder::MakeField(fieldb);
  auto literal_2 = TreeExprBuilder::MakeLiteral((int32_t)2);
  auto literal_3 = TreeExprBuilder::MakeLiteral((int32_t)3);
  auto a_gt_b =
      TreeExprBuilder::MakeFunction("greater_than", {node_a, node_b}, boolean());
  auto a_gt_2 =
      TreeExprBuilder::MakeFunction("greater_than", {node_a, literal_2}, boolean());
  auto a_gt_3 =
      TreeExprBuilder::MakeFunction("greater_than", {node_a, literal_3}, boolean());
  auto b_gt_2 =
      TreeExprBuilder::MakeFunction("greater_than", {node_b, literal_2}, boolean());
  auto b_gt_3 =
      TreeExprBuilder::MakeFunction("greater_than", {node_b, literal_3}, boolean());

  auto if_3 = TreeExprBuilder::MakeIf(a_gt_b, a_gt_3, b_gt_3, arrow::boolean());
  auto if_2 = TreeExprBuilder::MakeIf(a_gt_b, a_gt_2, b_gt_2, arrow::boolean());
  auto node_and = TreeExprBuilder::MakeAnd({if_3, if_2});
  auto expr = TreeExprBuilder::MakeExpression(node_and, field_result);

  // Build a projector for the expressions.
  std::shared_ptr<Projector> projector;
  auto status = Projector::Make(schema, {expr}, TestConfiguration(), &projector);
  EXPECT_TRUE(status.ok());

  int num_records = 4;
  std::vector<bool> validity({true, true, true, true});
  auto arraya = MakeArrowArrayInt32({4, 3, 3, 2}, validity);
  auto arrayb = MakeArrowArrayInt32({3, 4, 2, 3}, validity);
  auto exp = MakeArrowArrayBool({true, true, false, false}, validity);

  auto in_batch = arrow::RecordBatch::Make(schema, num_records, {arraya, arrayb});

  arrow::ArrayVector outputs;
  status = projector->Evaluate(*in_batch, pool_, &outputs);
  EXPECT_TRUE(status.ok());
  EXPECT_ARROW_ARRAY_EQUALS(exp, outputs.at(0));
}

}  // namespace gandiva