// Copyright 2014 Cloudera Inc.
//
// Licensed 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 "exec/data-source-scan-node.h"

#include <boost/foreach.hpp>
#include <vector>
#include <gutil/strings/substitute.h>

#include "exec/data-source-row-converter.h"
#include "exec/parquet-common.h"
#include "exec/read-write-util.h"
#include "exprs/expr.h"
#include "runtime/mem-pool.h"
#include "runtime/runtime-state.h"
#include "runtime/row-batch.h"
#include "runtime/string-value.h"
#include "runtime/tuple-row.h"
#include "util/jni-util.h"
#include "util/periodic-counter-updater.h"
#include "util/runtime-profile.h"

#include "common/names.h"

using namespace strings;
using namespace impala;
using namespace impala::extdatasource;

DEFINE_int32(data_source_batch_size, 1024, "Batch size for calls to GetNext() on "
    "external data sources.");

namespace impala {

DataSourceScanNode::DataSourceScanNode(ObjectPool* pool, const TPlanNode& tnode,
    const DescriptorTbl& descs)
    : ScanNode(pool, tnode, descs),
      data_src_node_(tnode.data_source_node),
      tuple_idx_(0) {
}

DataSourceScanNode::~DataSourceScanNode() {
}

Status DataSourceScanNode::Prepare(RuntimeState* state) {
  RETURN_IF_ERROR(ScanNode::Prepare(state));
  tuple_desc_ = state->desc_tbl().GetTupleDescriptor(data_src_node_.tuple_id);
  DCHECK(tuple_desc_ != NULL);

  data_source_executor_.reset(new ExternalDataSourceExecutor());
  RETURN_IF_ERROR(data_source_executor_->Init(data_src_node_.data_source.hdfs_location,
      data_src_node_.data_source.class_name, data_src_node_.data_source.api_version,
      data_src_node_.init_string));

  // Initialize materialized_slots_
  BOOST_FOREACH(SlotDescriptor* slot, tuple_desc_->slots()) {
    if (!slot->is_materialized()) continue;
    materialized_slots_.push_back(slot);
  }

  row_converter_.reset(
      new DataSourceRowConverter(tuple_desc_, NULL, materialized_slots_));
  return Status::OK();
}

Status DataSourceScanNode::Open(RuntimeState* state) {
  SCOPED_TIMER(runtime_profile_->total_time_counter());
  RETURN_IF_ERROR(ExecNode::Open(state));
  RETURN_IF_CANCELLED(state);

  // Prepare the schema for TOpenParams.row_schema
  vector<extdatasource::TColumnDesc> cols;
  BOOST_FOREACH(const SlotDescriptor* slot, materialized_slots_) {
    extdatasource::TColumnDesc col;
    int col_idx = slot->col_pos();
    col.__set_name(tuple_desc_->table_desc()->col_descs()[col_idx].name());
    col.__set_type(slot->type().ToThrift());
    cols.push_back(col);
  }
  extdatasource::TTableSchema row_schema;
  row_schema.__set_cols(cols);

  TOpenParams params;
  params.__set_query_id(state->query_id());
  params.__set_table_name(tuple_desc_->table_desc()->name());
  params.__set_init_string(data_src_node_.init_string);
  params.__set_authenticated_user_name(state->effective_user());
  params.__set_row_schema(row_schema);
  params.__set_batch_size(FLAGS_data_source_batch_size);
  params.__set_predicates(data_src_node_.accepted_predicates);
  TOpenResult result;
  RETURN_IF_ERROR(data_source_executor_->Open(params, &result));
  RETURN_IF_ERROR(Status(result.status));
  scan_handle_ = result.scan_handle;
  return GetNextInputBatch();
}

Status DataSourceScanNode::GetNextInputBatch() {
  input_batch_.reset(new TGetNextResult());
  TGetNextParams params;
  params.__set_scan_handle(scan_handle_);
  RETURN_IF_ERROR(data_source_executor_->GetNext(params, input_batch_.get()));
  RETURN_IF_ERROR(Status(input_batch_->status));

  if (input_batch_->__isset.rows) {
    RETURN_IF_ERROR(row_converter_->ResetRowBatch(&input_batch_->rows, true));
  }

  if (!InputBatchHasNext() && !input_batch_->eos) {
    // The data source should have set eos, but if it didn't we should just log a
    // warning and continue as if it had.
    VLOG_QUERY << "Data source " << data_src_node_.data_source.name << " returned no "
      << "rows but did not set 'eos'. No more rows will be fetched from the "
      << "data source.";
    input_batch_->eos = true;
  }
  return Status::OK();
}

Status DataSourceScanNode::GetNext(RuntimeState* state, RowBatch* row_batch, bool* eos) {
  RETURN_IF_ERROR(ExecDebugAction(TExecNodePhase::GETNEXT, state));
  RETURN_IF_CANCELLED(state);
  RETURN_IF_ERROR(QueryMaintenance(state));
  SCOPED_TIMER(runtime_profile_->total_time_counter());
  if (ReachedLimit()) {
    *eos = true;
    return Status::OK();
  }
  *eos = false;

  // create new tuple buffer for row_batch
  MemPool* tuple_pool = row_batch->tuple_data_pool();
  int tuple_buffer_size = row_batch->MaxTupleBufferSize();
  void* tuple_buffer = tuple_pool->Allocate(tuple_buffer_size);
  tuple_ = reinterpret_cast<Tuple*>(tuple_buffer);
  ExprContext** ctxs = &conjunct_ctxs_[0];
  int num_ctxs = conjunct_ctxs_.size();

  while (true) {
    {
      SCOPED_TIMER(materialize_tuple_timer());
      // copy rows until we hit the limit/capacity or until we exhaust input_batch_
      while (!ReachedLimit() && !row_batch->AtCapacity(tuple_pool) &&
          InputBatchHasNext()) {
        RETURN_IF_ERROR(row_converter_->MaterializeNextRow(tuple_, tuple_pool));
        int row_idx = row_batch->AddRow();
        TupleRow* tuple_row = row_batch->GetRow(row_idx);
        tuple_row->SetTuple(tuple_idx_, tuple_);

        if (ExecNode::EvalConjuncts(ctxs, num_ctxs, tuple_row)) {
          row_batch->CommitLastRow();
          char* new_tuple = reinterpret_cast<char*>(tuple_);
          new_tuple += tuple_desc_->byte_size();
          tuple_ = reinterpret_cast<Tuple*>(new_tuple);
          ++num_rows_returned_;
        }
      }
      COUNTER_SET(rows_returned_counter_, num_rows_returned_);

      if (ReachedLimit() || row_batch->AtCapacity() || input_batch_->eos) {
        *eos = ReachedLimit() || input_batch_->eos;
        return Status::OK();
      }
    }

    // Need more rows
    DCHECK(!InputBatchHasNext());
    RETURN_IF_ERROR(GetNextInputBatch());
  }
}

Status DataSourceScanNode::Reset(RuntimeState* state) {
  DCHECK(false) << "NYI";
  return Status("NYI");
}

void DataSourceScanNode::Close(RuntimeState* state) {
  if (is_closed()) return;
  SCOPED_TIMER(runtime_profile_->total_time_counter());
  PeriodicCounterUpdater::StopRateCounter(total_throughput_counter());
  PeriodicCounterUpdater::StopTimeSeriesCounter(bytes_read_timeseries_counter_);
  input_batch_.reset();
  TCloseParams params;
  params.__set_scan_handle(scan_handle_);
  TCloseResult result;
  Status status = data_source_executor_->Close(params, &result);
  if (!status.ok()) state->LogError(status.msg());
  ExecNode::Close(state);
}

void DataSourceScanNode::DebugString(int indentation_level, stringstream* out) const {
  string indent(indentation_level * 2, ' ');
  *out << indent << "DataSourceScanNode(tupleid=" << data_src_node_.tuple_id << ")";
}

}  // namespace impala
