// 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.


#ifndef IMPALA_UDF_TEST_HARNESS_H
#define IMPALA_UDF_TEST_HARNESS_H

// THIS FILE IS USED BY THE STANDALONE IMPALA UDF DEVELOPMENT KIT.
// IT MUST BE BUILDABLE WITH C++98 AND WITHOUT ANY INTERNAL IMPALA HEADERS.

#include <iostream>
#include <vector>
#include <boost/function.hpp>
#include <boost/scoped_ptr.hpp>

#include "udf/udf.h"
#include "udf/udf-debug.h"

namespace impala {
  class MemPool;
  class RuntimeState;
}

namespace impala_udf {

/// Utility class to help test UDFs.
class UdfTestHarness {
 public:
  /// Create a test FunctionContext object. 'arg_types' should contain a TypeDesc for each
  /// argument of the UDF not including the FunctionContext*. The caller is responsible
  /// for calling delete on it. This context has additional debugging validation enabled.
  static FunctionContext* CreateTestContext(const FunctionContext::TypeDesc& return_type,
      const std::vector<FunctionContext::TypeDesc>& arg_types,
      impala::RuntimeState* state = NULL, impala::MemPool* pool = NULL);

  /// Use with test contexts to test use of IsArgConstant() and GetConstantArg().
  /// constant_args should contain an AnyVal* for each argument of the UDF not including
  /// the FunctionContext*; constant_args[i] corresponds to the i-th argument.
  /// Non-constant arguments should be set to NULL, and constant arguments should be set
  /// to the constant value.
  //
  /// The AnyVal* values are owned by the caller.
  //
  /// Can only be called on contexts created by CreateTestContext().
  static void SetConstantArgs(
      FunctionContext* context, const std::vector<AnyVal*>& constant_args);

  /// Test contexts should be closed in order to check for UDF memory leaks. Leaks cause
  /// the error to be set on context.
  static void CloseContext(FunctionContext* context);

  /// Template function to execute a UDF and validate the result. They should be
  /// used like:
  /// ValidateUdf(udf_fn, arg1, arg2, ..., expected_result);
  /// Only functions with up to 8 arguments are supported
  //
  /// For variable argument udfs, the variable arguments should be passed as
  /// a std::vector:
  ///   ValidateUdf(udf_fn, arg1, arg2, const vector<arg3>& args, expected_result);
  template<typename RET>
  static bool ValidateUdf(boost::function<RET(FunctionContext*)> fn,
      const RET& expected, UdfPrepare init_fn = NULL, UdfClose close_fn = NULL,
      const std::vector<AnyVal*>& constant_args = std::vector<AnyVal*>()) {
    FunctionContext::TypeDesc return_type; // TODO
    std::vector<FunctionContext::TypeDesc> arg_types;
    boost::scoped_ptr<FunctionContext> context(CreateTestContext(return_type, arg_types));
    SetConstantArgs(context.get(), constant_args);
    if (!RunPrepareFn(init_fn, context.get())) return false;
    RET ret = fn(context.get());
    RunCloseFn(close_fn, context.get());
    return Validate(context.get(), expected, ret);
  }

  template<typename RET, typename A1>
  static bool ValidateUdf(boost::function<RET(FunctionContext*, const A1&)> fn,
      const A1& a1, const RET& expected, UdfPrepare init_fn = NULL,
      UdfClose close_fn = NULL,
      const std::vector<AnyVal*>& constant_args = std::vector<AnyVal*>()) {
    FunctionContext::TypeDesc return_type; // TODO
    std::vector<FunctionContext::TypeDesc> arg_types; // TODO
    boost::scoped_ptr<FunctionContext> context(CreateTestContext(return_type, arg_types));
    SetConstantArgs(context.get(), constant_args);
    if (!RunPrepareFn(init_fn, context.get())) return false;
    RET ret = fn(context.get(), a1);
    RunCloseFn(close_fn, context.get());
    return Validate(context.get(), expected, ret);
  }

  template<typename RET, typename A1>
  static bool ValidateUdf(boost::function<RET(FunctionContext*, int, const A1*)> fn,
      const std::vector<A1>& a1, const RET& expected, UdfPrepare init_fn = NULL,
      UdfClose close_fn = NULL,
      const std::vector<AnyVal*>& constant_args = std::vector<AnyVal*>()) {
    FunctionContext::TypeDesc return_type; // TODO
    std::vector<FunctionContext::TypeDesc> arg_types; // TODO
    boost::scoped_ptr<FunctionContext> context(CreateTestContext(return_type, arg_types));
    SetConstantArgs(context.get(), constant_args);
    if (!RunPrepareFn(init_fn, context.get())) return false;
    RET ret = fn(context.get(), a1.size(), a1.data());
    RunCloseFn(close_fn, context.get());
    return Validate(context.get(), expected, ret);
  }

  template<typename RET, typename A1, typename A2>
  static bool ValidateUdf(
      boost::function<RET(FunctionContext*, const A1&, const A2&)> fn,
      const A1& a1, const A2& a2, const RET& expected, UdfPrepare init_fn = NULL,
      UdfClose close_fn = NULL,
      const std::vector<AnyVal*>& constant_args = std::vector<AnyVal*>()) {
    FunctionContext::TypeDesc return_type; // TODO
    std::vector<FunctionContext::TypeDesc> arg_types; // TODO
    boost::scoped_ptr<FunctionContext> context(CreateTestContext(return_type, arg_types));
    SetConstantArgs(context.get(), constant_args);
    if (!RunPrepareFn(init_fn, context.get())) return false;
    RET ret = fn(context.get(), a1, a2);
    RunCloseFn(close_fn, context.get());
    return Validate(context.get(), expected, ret);
  }

  template<typename RET, typename A1, typename A2>
  static bool ValidateUdf(
      boost::function<RET(FunctionContext*, const A1&, int, const A2*)> fn,
      const A1& a1, const std::vector<A2>& a2, const RET& expected,
      UdfPrepare init_fn = NULL, UdfClose close_fn = NULL,
      const std::vector<AnyVal*>& constant_args = std::vector<AnyVal*>()) {
    FunctionContext::TypeDesc return_type; // TODO
    std::vector<FunctionContext::TypeDesc> arg_types; // TODO
    boost::scoped_ptr<FunctionContext> context(CreateTestContext(return_type, arg_types));
    SetConstantArgs(context.get(), constant_args);
    if (!RunPrepareFn(init_fn, context.get())) return false;
    RET ret = fn(context.get(), a1, a2.size(), a2.data());
    RunCloseFn(close_fn, context.get());
    return Validate(context.get(), expected, ret);
  }

  template<typename RET, typename A1, typename A2, typename A3>
  static bool ValidateUdf(
      boost::function<RET(FunctionContext*, const A1&, const A2&, const A3&)> fn,
      const A1& a1, const A2& a2, const A3& a3, const RET& expected,
      UdfPrepare init_fn = NULL, UdfClose close_fn = NULL,
      const std::vector<AnyVal*>& constant_args = std::vector<AnyVal*>()) {
    FunctionContext::TypeDesc return_type; // TODO
    std::vector<FunctionContext::TypeDesc> arg_types; // TODO
    boost::scoped_ptr<FunctionContext> context(CreateTestContext(return_type, arg_types));
    SetConstantArgs(context.get(), constant_args);
    if (!RunPrepareFn(init_fn, context.get())) return false;
    RET ret = fn(context.get(), a1, a2, a3);
    RunCloseFn(close_fn, context.get());
    return Validate(context.get(), expected, ret);
  }

  template<typename RET, typename A1, typename A2, typename A3>
  static bool ValidateUdf(
      boost::function<RET(FunctionContext*, const A1&, const A2&, int, const A3*)> fn,
      const A1& a1, const A2& a2, const std::vector<A3>& a3, const RET& expected,
      UdfPrepare init_fn = NULL, UdfClose close_fn = NULL,
      const std::vector<AnyVal*>& constant_args = std::vector<AnyVal*>()) {
    FunctionContext::TypeDesc return_type; // TODO
    std::vector<FunctionContext::TypeDesc> arg_types; // TODO
    boost::scoped_ptr<FunctionContext> context(CreateTestContext(return_type, arg_types));
    SetConstantArgs(context.get(), constant_args);
    if (!RunPrepareFn(init_fn, context.get())) return false;
    RET ret = fn(context.get(), a1, a2, a3.size(), a3.data());
    RunCloseFn(close_fn, context.get());
    return Validate(context.get(), expected, ret);
  }

  template<typename RET, typename A1, typename A2, typename A3, typename A4>
  static bool ValidateUdf(
      boost::function<RET(FunctionContext*, const A1&, const A2&, const A3&,
          const A4&)> fn,
      const A1& a1, const A2& a2, const A3& a3, const A4& a4, const RET& expected,
      UdfPrepare init_fn = NULL, UdfClose close_fn = NULL,
      const std::vector<AnyVal*>& constant_args = std::vector<AnyVal*>()) {
    FunctionContext::TypeDesc return_type; // TODO
    std::vector<FunctionContext::TypeDesc> arg_types; // TODO
    boost::scoped_ptr<FunctionContext> context(CreateTestContext(return_type, arg_types));
    SetConstantArgs(context.get(), constant_args);
    if (!RunPrepareFn(init_fn, context.get())) return false;
    RET ret = fn(context.get(), a1, a2, a3, a4);
    RunCloseFn(close_fn, context.get());
    return Validate(context.get(), expected, ret);
  }

  template<typename RET, typename A1, typename A2, typename A3, typename A4>
  static bool ValidateUdf(
      boost::function<RET(FunctionContext*, const A1&, const A2&, const A3&,
          int, const A4*)> fn,
      const A1& a1, const A2& a2, const A3& a3, const std::vector<A4>& a4,
      const RET& expected, UdfPrepare init_fn = NULL, UdfClose close_fn = NULL,
      const std::vector<AnyVal*>& constant_args = std::vector<AnyVal*>()) {
    FunctionContext::TypeDesc return_type; // TODO
    std::vector<FunctionContext::TypeDesc> arg_types; // TODO
    boost::scoped_ptr<FunctionContext> context(CreateTestContext(return_type, arg_types));
    SetConstantArgs(context.get(), constant_args);
    if (!RunPrepareFn(init_fn, context.get())) return false;
    RET ret = fn(context.get(), a1, a2, a3, a4.size(), a4.data());
    RunCloseFn(close_fn, context.get());
    return Validate(context.get(), expected, ret);
  }

  template<typename RET, typename A1, typename A2, typename A3, typename A4,
      typename A5>
  static bool ValidateUdf(
      boost::function<RET(FunctionContext*, const A1&, const A2&, const A3&,
          const A4&, const A5&)> fn,
      const A1& a1, const A2& a2, const A3& a3, const A4& a4,const A5& a5,
      const RET& expected, UdfPrepare init_fn = NULL, UdfClose close_fn = NULL,
      const std::vector<AnyVal*>& constant_args = std::vector<AnyVal*>()) {
    FunctionContext::TypeDesc return_type; // TODO
    std::vector<FunctionContext::TypeDesc> arg_types; // TODO
    boost::scoped_ptr<FunctionContext> context(CreateTestContext(return_type, arg_types));
    SetConstantArgs(context.get(), constant_args);
    if (!RunPrepareFn(init_fn, context.get())) return false;
    RET ret = fn(context.get(), a1, a2, a3, a4, a5);
    RunCloseFn(close_fn, context.get());
    return Validate(context.get(), expected, ret);
  }

  template<typename RET, typename A1, typename A2, typename A3, typename A4,
      typename A5, typename A6>
  static bool ValidateUdf(
      boost::function<RET(FunctionContext*, const A1&, const A2&, const A3&,
          const A4&, const A5&, const A6&)> fn,
      const A1& a1, const A2& a2, const A3& a3, const A4& a4, const A5& a5,
      const A6& a6, const RET& expected, UdfPrepare init_fn = NULL,
      UdfClose close_fn = NULL,
      const std::vector<AnyVal*>& constant_args = std::vector<AnyVal*>()) {
    FunctionContext::TypeDesc return_type; // TODO
    std::vector<FunctionContext::TypeDesc> arg_types; // TODO
    boost::scoped_ptr<FunctionContext> context(CreateTestContext(return_type, arg_types));
    SetConstantArgs(context.get(), constant_args);
    if (!RunPrepareFn(init_fn, context.get())) return false;
    RET ret = fn(context.get(), a1, a2, a3, a4, a5, a6);
    RunCloseFn(close_fn, context.get());
    return Validate(context.get(), expected, ret);
  }

  template<typename RET, typename A1, typename A2, typename A3, typename A4,
      typename A5, typename A6, typename A7>
  static bool ValidateUdf(
      boost::function<RET(FunctionContext*, const A1&, const A2&, const A3&,
          const A4&, const A5&, const A6&, const A7&)> fn,
      const A1& a1, const A2& a2, const A3& a3, const A4& a4, const A5& a5,
      const A6& a6, const A7& a7, const RET& expected, UdfPrepare init_fn = NULL,
      UdfClose close_fn = NULL,
      const std::vector<AnyVal*>& constant_args = std::vector<AnyVal*>()) {
    FunctionContext::TypeDesc return_type; // TODO
    std::vector<FunctionContext::TypeDesc> arg_types; // TODO
    boost::scoped_ptr<FunctionContext> context(CreateTestContext(return_type, arg_types));
    SetConstantArgs(context.get(), constant_args);
    if (!RunPrepareFn(init_fn, context.get())) return false;
    RET ret = fn(context.get(), a1, a2, a3, a4, a5, a6, a7);
    RunCloseFn(close_fn, context.get());
    return Validate(context.get(), expected, ret);
  }

  template<typename RET, typename A1, typename A2, typename A3, typename A4,
      typename A5, typename A6, typename A7, typename A8>
  static bool ValidateUdf(
      boost::function<RET(FunctionContext*, const A1&, const A2&, const A3&,
          const A4&, const A5&, const A6&, const A7&)> fn,
      const A1& a1, const A2& a2, const A3& a3, const A4& a4, const A5& a5,
      const A6& a6, const A7& a7, const A8& a8, const RET& expected,
      UdfPrepare init_fn = NULL, UdfClose close_fn = NULL,
      const std::vector<AnyVal*>& constant_args = std::vector<AnyVal*>()) {
    FunctionContext::TypeDesc return_type; // TODO
    std::vector<FunctionContext::TypeDesc> arg_types; // TODO
    boost::scoped_ptr<FunctionContext> context(CreateTestContext(return_type, arg_types));
    SetConstantArgs(context.get(), constant_args);
    if (!RunPrepareFn(init_fn, context.get())) return false;
    RET ret = fn(context.get(), a1, a2, a3, a4, a5, a6, a7, a8);
    RunCloseFn(close_fn, context.get());
    return Validate(context.get(), expected, ret);
  }

 private:
  static bool ValidateError(FunctionContext* context) {
    if (context->has_error()) {
      std::cerr << "Udf Failed: " << context->error_msg() << std::endl;
      return false;
    }
    return true;
  }

  template<typename RET>
  static bool Validate(FunctionContext* context, const RET& expected, const RET& actual) {
    bool valid = true;
    if (!context->has_error() && actual != expected) {
      std::cerr << "UDF did not return the correct result:" << std::endl
                << "  Expected: " << DebugString(expected) << std::endl
                << "  Actual: " << DebugString(actual) << std::endl;
      valid = false;
    }
    CloseContext(context);
    if (!ValidateError(context)) valid = false;
    return valid;
  }

  static bool RunPrepareFn(UdfPrepare prepare_fn, FunctionContext* context) {
    if (prepare_fn != NULL) {
      // TODO: FRAGMENT_LOCAL
      prepare_fn(context, FunctionContext::THREAD_LOCAL);
      if (!ValidateError(context)) return false;
    }
    return true;
  }

  static void RunCloseFn(UdfClose close_fn, FunctionContext* context) {
    if (close_fn != NULL) {
      // TODO: FRAGMENT_LOCAL
      close_fn(context, FunctionContext::THREAD_LOCAL);
    }
  }
};

}

#endif