/*
 * 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.
 */
package org.apache.beam.sdk.extensions.sql.impl.rel;

import static org.apache.beam.vendor.calcite.v1_20_0.com.google.common.base.Preconditions.checkArgument;
import static org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.avatica.util.DateTimeUtils.MILLIS_PER_DAY;

import java.io.IOException;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Queue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.stream.Collectors;
import javax.annotation.Nullable;
import org.apache.beam.runners.direct.DirectOptions;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.Pipeline.PipelineVisitor;
import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.PipelineResult.State;
import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils.CharType;
import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils.DateType;
import org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils.TimeType;
import org.apache.beam.sdk.metrics.Counter;
import org.apache.beam.sdk.metrics.MetricNameFilter;
import org.apache.beam.sdk.metrics.MetricQueryResults;
import org.apache.beam.sdk.metrics.MetricResult;
import org.apache.beam.sdk.metrics.Metrics;
import org.apache.beam.sdk.metrics.MetricsFilter;
import org.apache.beam.sdk.options.ApplicationNameOptions;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.runners.TransformHierarchy.Node;
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollection.IsBounded;
import org.apache.beam.sdk.values.PValue;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.adapter.enumerable.EnumerableRel;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.adapter.enumerable.EnumerableRelImplementor;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.adapter.enumerable.PhysType;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.adapter.enumerable.PhysTypeImpl;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.linq4j.Enumerable;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.linq4j.Linq4j;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.linq4j.tree.BlockBuilder;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.linq4j.tree.Expression;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.linq4j.tree.Expressions;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.plan.ConventionTraitDef;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.plan.RelOptCluster;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.plan.RelOptCost;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.plan.RelOptPlanner;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.plan.RelTraitSet;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.RelNode;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.convert.ConverterImpl;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.metadata.RelMetadataQuery;
import org.apache.beam.vendor.calcite.v1_20_0.org.apache.calcite.rel.type.RelDataType;
import org.joda.time.Duration;
import org.joda.time.ReadableInstant;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** BeamRelNode to replace a {@code Enumerable} node. */
public class BeamEnumerableConverter extends ConverterImpl implements EnumerableRel {
  private static final Logger LOG = LoggerFactory.getLogger(BeamEnumerableConverter.class);

  public BeamEnumerableConverter(RelOptCluster cluster, RelTraitSet traits, RelNode input) {
    super(cluster, ConventionTraitDef.INSTANCE, traits, input);
  }

  @Override
  public RelNode copy(RelTraitSet traitSet, List<RelNode> inputs) {
    return new BeamEnumerableConverter(getCluster(), traitSet, sole(inputs));
  }

  @Override
  public RelOptCost computeSelfCost(RelOptPlanner planner, RelMetadataQuery mq) {
    // This should always be a last resort.
    return planner.getCostFactory().makeHugeCost();
  }

  @Override
  public Result implement(EnumerableRelImplementor implementor, Prefer prefer) {
    final BlockBuilder list = new BlockBuilder();
    final RelDataType rowType = getRowType();
    final PhysType physType =
        PhysTypeImpl.of(implementor.getTypeFactory(), rowType, prefer.preferArray());
    final Expression node = implementor.stash((BeamRelNode) getInput(), BeamRelNode.class);
    list.add(Expressions.call(BeamEnumerableConverter.class, "toEnumerable", node));
    return implementor.result(physType, list.toBlock());
  }

  public static Enumerable<Object> toEnumerable(BeamRelNode node) {
    final ClassLoader originalClassLoader = Thread.currentThread().getContextClassLoader();
    try {
      Thread.currentThread().setContextClassLoader(BeamEnumerableConverter.class.getClassLoader());
      final PipelineOptions options = createPipelineOptions(node.getPipelineOptions());
      return toEnumerable(options, node);
    } finally {
      Thread.currentThread().setContextClassLoader(originalClassLoader);
    }
  }

  public static List<Row> toRowList(BeamRelNode node) {
    final ClassLoader originalClassLoader = Thread.currentThread().getContextClassLoader();
    try {
      Thread.currentThread().setContextClassLoader(BeamEnumerableConverter.class.getClassLoader());
      final PipelineOptions options = createPipelineOptions(node.getPipelineOptions());
      return toRowList(options, node);
    } finally {
      Thread.currentThread().setContextClassLoader(originalClassLoader);
    }
  }

  public static PipelineOptions createPipelineOptions(Map<String, String> map) {
    final String[] args = new String[map.size()];
    int i = 0;
    for (Map.Entry<String, String> entry : map.entrySet()) {
      args[i++] = "--" + entry.getKey() + "=" + entry.getValue();
    }
    PipelineOptions options = PipelineOptionsFactory.fromArgs(args).withValidation().create();
    options.as(ApplicationNameOptions.class).setAppName("BeamSql");
    return options;
  }

  static List<Row> toRowList(PipelineOptions options, BeamRelNode node) {
    if (node instanceof BeamIOSinkRel) {
      throw new UnsupportedOperationException("Does not support BeamIOSinkRel in toRowList.");
    } else if (isLimitQuery(node)) {
      throw new UnsupportedOperationException("Does not support queries with LIMIT in toRowList.");
    }
    return collectRows(options, node).stream().collect(Collectors.toList());
  }

  static Enumerable<Object> toEnumerable(PipelineOptions options, BeamRelNode node) {
    if (node instanceof BeamIOSinkRel) {
      return count(options, node);
    } else if (isLimitQuery(node)) {
      return limitCollect(options, node);
    }
    return Linq4j.asEnumerable(rowToAvaticaAndUnboxValues((collectRows(options, node))));
  }

  private static PipelineResult limitRun(
      PipelineOptions options,
      BeamRelNode node,
      DoFn<Row, Void> doFn,
      Queue<Row> values,
      int limitCount) {
    options.as(DirectOptions.class).setBlockOnRun(false);
    Pipeline pipeline = Pipeline.create(options);
    PCollection<Row> resultCollection = BeamSqlRelUtils.toPCollection(pipeline, node);
    resultCollection.apply(ParDo.of(doFn));

    PipelineResult result = pipeline.run();

    State state;
    while (true) {
      // Check pipeline state in every second
      state = result.waitUntilFinish(Duration.standardSeconds(1));
      if (state != null && state.isTerminal()) {
        break;
      }

      try {
        if (values.size() >= limitCount) {
          result.cancel();
          break;
        }
      } catch (IOException e) {
        LOG.warn(e.toString());
        break;
      }
    }

    return result;
  }

  private static void runCollector(PipelineOptions options, BeamRelNode node) {
    Pipeline pipeline = Pipeline.create(options);
    PCollection<Row> resultCollection = BeamSqlRelUtils.toPCollection(pipeline, node);
    resultCollection.apply(ParDo.of(new Collector()));
    PipelineResult result = pipeline.run();
    result.waitUntilFinish();
  }

  private static Queue<Row> collectRows(PipelineOptions options, BeamRelNode node) {
    long id = options.getOptionsId();
    Queue<Row> values = new ConcurrentLinkedQueue<>();

    checkArgument(
        options
            .getRunner()
            .getCanonicalName()
            .equals("org.apache.beam.runners.direct.DirectRunner"),
        "collectRowList is only available in direct runner.");

    Collector.globalValues.put(id, values);

    runCollector(options, node);

    Collector.globalValues.remove(id);
    return values;
  }

  private static Enumerable<Object> limitCollect(PipelineOptions options, BeamRelNode node) {
    long id = options.getOptionsId();
    ConcurrentLinkedQueue<Row> values = new ConcurrentLinkedQueue<>();

    checkArgument(
        options
            .getRunner()
            .getCanonicalName()
            .equals("org.apache.beam.runners.direct.DirectRunner"),
        "SELECT without INSERT is only supported in DirectRunner in SQL Shell.");

    int limitCount = getLimitCount(node);

    Collector.globalValues.put(id, values);
    limitRun(options, node, new Collector(), values, limitCount);
    Collector.globalValues.remove(id);

    // remove extra retrieved values
    while (values.size() > limitCount) {
      values.remove();
    }

    return Linq4j.asEnumerable(rowToAvaticaAndUnboxValues(values));
  }

  private static class Collector extends DoFn<Row, Void> {

    // This will only work on the direct runner.
    private static final Map<Long, Queue<Row>> globalValues = new ConcurrentHashMap<>();

    @Nullable private volatile Queue<Row> values;

    @StartBundle
    public void startBundle(StartBundleContext context) {
      long id = context.getPipelineOptions().getOptionsId();
      values = globalValues.get(id);
    }

    @ProcessElement
    public void processElement(ProcessContext context) {
      values.add(context.element());
    }
  }

  private static List<Object> rowToAvaticaAndUnboxValues(Queue<Row> values) {
    return values.stream()
        .map(
            row -> {
              Object[] objects = rowToAvatica(row);
              if (objects.length == 1) {
                // if objects.length == 1, that means input Row contains only 1 column/element,
                // then an Object instead of Object[] should be returned because of
                // CalciteResultSet's behaviour that tries to convert one column row to an Object.
                return objects[0];
              } else {
                return objects;
              }
            })
        .collect(Collectors.toList());
  }

  private static Object[] rowToAvatica(Row row) {
    Schema schema = row.getSchema();
    Object[] convertedColumns = new Object[schema.getFields().size()];
    int i = 0;
    for (Schema.Field field : schema.getFields()) {
      convertedColumns[i] = fieldToAvatica(field.getType(), row.getValue(i));
      ++i;
    }
    return convertedColumns;
  }

  private static Object fieldToAvatica(Schema.FieldType type, Object beamValue) {
    if (beamValue == null) {
      return null;
    }

    switch (type.getTypeName()) {
      case LOGICAL_TYPE:
        String logicalId = type.getLogicalType().getIdentifier();
        if (logicalId.equals(TimeType.IDENTIFIER)) {
          return (int) ((ReadableInstant) beamValue).getMillis();
        } else if (logicalId.equals(DateType.IDENTIFIER)) {
          return (int) (((ReadableInstant) beamValue).getMillis() / MILLIS_PER_DAY);
        } else if (logicalId.equals(CharType.IDENTIFIER)) {
          return beamValue;
        } else {
          throw new IllegalArgumentException("Unknown DateTime type " + logicalId);
        }
      case DATETIME:
        return ((ReadableInstant) beamValue).getMillis();
      case BYTE:
      case INT16:
      case INT32:
      case INT64:
      case DECIMAL:
      case FLOAT:
      case DOUBLE:
      case STRING:
      case BOOLEAN:
      case BYTES:
        return beamValue;
      case ARRAY:
        return ((List<?>) beamValue)
            .stream()
                .map(elem -> fieldToAvatica(type.getCollectionElementType(), elem))
                .collect(Collectors.toList());
      case MAP:
        return ((Map<?, ?>) beamValue)
            .entrySet().stream()
                .collect(
                    Collectors.toMap(
                        entry -> entry.getKey(),
                        entry ->
                            fieldToAvatica(type.getCollectionElementType(), entry.getValue())));
      case ROW:
        // TODO: needs to be a Struct
        return beamValue;
      default:
        throw new IllegalStateException(
            String.format("Unreachable case for Beam typename %s", type.getTypeName()));
    }
  }

  private static Enumerable<Object> count(PipelineOptions options, BeamRelNode node) {
    Pipeline pipeline = Pipeline.create(options);
    BeamSqlRelUtils.toPCollection(pipeline, node).apply(ParDo.of(new RowCounter()));
    PipelineResult result = pipeline.run();

    long count = 0;
    if (!containsUnboundedPCollection(pipeline)) {
      result.waitUntilFinish();
      MetricQueryResults metrics =
          result
              .metrics()
              .queryMetrics(
                  MetricsFilter.builder()
                      .addNameFilter(MetricNameFilter.named(BeamEnumerableConverter.class, "rows"))
                      .build());
      Iterator<MetricResult<Long>> iterator = metrics.getCounters().iterator();
      if (iterator.hasNext()) {
        count = iterator.next().getAttempted();
      }
    }
    return Linq4j.singletonEnumerable(count);
  }

  private static class RowCounter extends DoFn<Row, Void> {
    final Counter rows = Metrics.counter(BeamEnumerableConverter.class, "rows");

    @ProcessElement
    public void processElement(ProcessContext context) {
      rows.inc();
    }
  }

  private static boolean isLimitQuery(BeamRelNode node) {
    return (node instanceof BeamSortRel && ((BeamSortRel) node).isLimitOnly())
        || (node instanceof BeamCalcRel && ((BeamCalcRel) node).isInputSortRelAndLimitOnly());
  }

  private static int getLimitCount(BeamRelNode node) {
    if (node instanceof BeamSortRel) {
      return ((BeamSortRel) node).getCount();
    } else if (node instanceof BeamCalcRel) {
      return ((BeamCalcRel) node).getLimitCountOfSortRel();
    }

    throw new RuntimeException(
        "Cannot get limit count from RelNode tree with root " + node.getRelTypeName());
  }

  private static boolean containsUnboundedPCollection(Pipeline p) {
    class BoundednessVisitor extends PipelineVisitor.Defaults {
      IsBounded boundedness = IsBounded.BOUNDED;

      @Override
      public void visitValue(PValue value, Node producer) {
        if (value instanceof PCollection) {
          boundedness = boundedness.and(((PCollection) value).isBounded());
        }
      }
    }

    BoundednessVisitor visitor = new BoundednessVisitor();
    p.traverseTopologically(visitor);
    return visitor.boundedness == IsBounded.UNBOUNDED;
  }
}