/*
 * Copyright (c) 2015. Qubole 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.
 */
package com.qubole.quark.fatjdbc.executor;

import org.apache.calcite.adapter.java.JavaTypeFactory;
import org.apache.calcite.avatica.AvaticaParameter;
import org.apache.calcite.avatica.AvaticaStatement;
import org.apache.calcite.avatica.ColumnMetaData;
import org.apache.calcite.avatica.Meta;
import org.apache.calcite.jdbc.JavaTypeFactoryImpl;
import org.apache.calcite.linq4j.Ord;
import org.apache.calcite.plan.RelOptUtil;
import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.rel.type.RelDataTypeFactory;
import org.apache.calcite.rel.type.RelDataTypeFactoryImpl;
import org.apache.calcite.rel.type.RelDataTypeField;
import org.apache.calcite.rel.type.RelDataTypeFieldImpl;
import org.apache.calcite.rel.type.RelRecordType;
import org.apache.calcite.sql.type.SqlTypeName;
import org.apache.calcite.util.Util;

import com.google.common.cache.Cache;
import com.google.common.collect.ImmutableList;

import com.qubole.quark.catalog.db.pojo.DataSource;
import com.qubole.quark.catalog.db.pojo.View;
import com.qubole.quark.ee.QuarkExecutor;
import com.qubole.quark.ee.QuarkExecutorFactory;
import com.qubole.quark.fatjdbc.QuarkConnectionImpl;
import com.qubole.quark.fatjdbc.QuarkJdbcStatement;
import com.qubole.quark.fatjdbc.QuarkMetaResultSet;
import com.qubole.quark.planner.parser.ParserResult;

import java.lang.reflect.Type;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Types;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;

/**
 * Created by amoghm on 3/4/16.
 */
public class PlanExecutor {
  private final Meta.StatementHandle h;
  private final QuarkConnectionImpl connection;
  private final long maxRowCount;
  private final Cache<String, Connection> connectionCache;

  public PlanExecutor(Meta.StatementHandle h,
               QuarkConnectionImpl connection,
               Cache<String, Connection> connectionCache,
               long maxRowCount) {
    this.h = h;
    this.connection = connection;
    this.connectionCache = connectionCache;
    this.maxRowCount = maxRowCount;
  }

  public QuarkMetaResultSet execute(ParserResult parserResult) throws Exception {
    QuarkJdbcStatement stmt = connection.server.getStatement(h);

    QuarkExecutor executor = QuarkExecutorFactory.getQuarkExecutor(parserResult.getKind(),
        connection.parserFactory, connection.getProperties(), connectionCache);
    Object result = executor.execute(parserResult);

    if (result instanceof Integer) {
      // Alter, Create, Drop DDL commands will either return id or 0
      return QuarkMetaResultSet.count(h.connectionId, h.id, ((Integer) result).intValue());
    } else if (result instanceof ArrayList) {
      // Show DDL returns an arraylist
      Class pojoType = getPojoType(parserResult.getParsedSql());

      return getMetaResultSetFromIterator(
          convertToIterator((ArrayList) result, pojoType),
          connection, parserResult, "", connection.server.getStatement(h), h,
          AvaticaStatement.DEFAULT_FETCH_SIZE, pojoType);
    } else if (result instanceof ResultSet) {
      // Querying JdbcDB
      return QuarkMetaResultSet.create(h.connectionId, h.id, (ResultSet) result,
          maxRowCount);
    } else if (result instanceof Iterator) {
      // Querying QuboleDB
      return getMetaResultSetFromIterator((Iterator<Object>) result,
          connection, parserResult, parserResult.getParsedSql(), stmt, h, maxRowCount, null);
    }

    throw new RuntimeException("Cannot handle execution for: " + parserResult.getParsedSql());
  }

  private Class getPojoType(String sql) throws SQLException {
    if (sql.startsWith("SHOW DATASOURCE")) {
      return DataSource.class;
    } else if (sql.startsWith("SHOW VIEW")) {
      return View.class;
    }
    throw new SQLException("Unable to determine POJO for: " + sql);
  }

  private Iterator<Object> convertToIterator(List list, Class pojoType) throws SQLException {
    List<Object> resultSet = new ArrayList<>();

    for (int i = 0; i < list.size(); i++) {
      String[] row = getValues(list.get(i), pojoType);
      resultSet.add(row);
    }
    return  resultSet.iterator();
  }

  private String[] getValues(Object object, Class pojoType) throws SQLException {
    if (DataSource.class.equals(pojoType)) {
      return ((DataSource) object).values();
    } else if (View.class.equals(pojoType)) {
      return ((View) object).values();
    }
    throw new SQLException("Unknown object type for: " + pojoType);
  }

  private RelDataType getRowType(Class clazz) throws SQLException {
    if (DataSource.class.equals(clazz)) {
      return getDataSourceRowType();
    } else if (View.class.equals(clazz)) {
      return getViewRowType();
    }

    throw new SQLException("RowType not defined for class: " + clazz);
  }

  private RelDataType getDataSourceRowType() throws SQLException {

    List<RelDataTypeField> relDataTypeFields =
        ImmutableList.<RelDataTypeField>of(
            new RelDataTypeFieldImpl("id", 1, getIntegerJavaType()),
            new RelDataTypeFieldImpl("type", 2, getStringJavaType()),
            new RelDataTypeFieldImpl("url", 3, getStringJavaType()),
            new RelDataTypeFieldImpl("name", 4, getStringJavaType()),
            new RelDataTypeFieldImpl("ds_set_id", 5, getIntegerJavaType()),
            new RelDataTypeFieldImpl("datasource_type", 6, getStringJavaType()),
            new RelDataTypeFieldImpl("auth_token", 7, getStringJavaType()),
            new RelDataTypeFieldImpl("dbtap_id", 8, getIntegerJavaType()),
            new RelDataTypeFieldImpl("username", 9, getStringJavaType()),
            new RelDataTypeFieldImpl("password", 10, getStringJavaType()));

    return new RelRecordType(relDataTypeFields);
  }

  private RelDataType getViewRowType() {

    List<RelDataTypeField> relDataTypeFields =
        ImmutableList.<RelDataTypeField>of(
            new RelDataTypeFieldImpl("id", 1, getIntegerJavaType()),
            new RelDataTypeFieldImpl("name", 2, getStringJavaType()),
            new RelDataTypeFieldImpl("description", 3, getStringJavaType()),
            new RelDataTypeFieldImpl("cost", 4, getIntegerJavaType()),
            new RelDataTypeFieldImpl("query", 5, getStringJavaType()),
            new RelDataTypeFieldImpl("destination_id", 6, getIntegerJavaType()),
            new RelDataTypeFieldImpl("schema_name", 7, getStringJavaType()),
            new RelDataTypeFieldImpl("table_name", 8, getStringJavaType()),
            new RelDataTypeFieldImpl("ds_set_id", 9, getIntegerJavaType()));

    return new RelRecordType(relDataTypeFields);
  }

  private RelDataTypeFactoryImpl.JavaType getIntegerJavaType() {
    RelDataTypeFactoryImpl relDataTypeFactoryImpl = new JavaTypeFactoryImpl();
    return relDataTypeFactoryImpl.new JavaType(Integer.class);
  }

  private RelDataTypeFactoryImpl.JavaType getStringJavaType() {
    RelDataTypeFactoryImpl relDataTypeFactoryImpl = new JavaTypeFactoryImpl();
    return relDataTypeFactoryImpl.new JavaType(String.class,
        !(String.class.isPrimitive()), Util.getDefaultCharset(), null);
  }

  private QuarkMetaResultSet getMetaResultSetFromIterator(Iterator<Object> iterator,
                                                  QuarkConnectionImpl connection,
                                                  ParserResult result,
                                                  String sql,
                                                  QuarkJdbcStatement stmt,
                                                  Meta.StatementHandle h,
                                                  long maxRowCount,
                                                  Class clazz) throws SQLException {
    QuarkMetaResultSet metaResultSet;
    final JavaTypeFactory typeFactory =
        connection.getSqlQueryParser().getTypeFactory();
    final RelDataType x;
    switch (result.getKind()) {
      case INSERT:
      case EXPLAIN:
        x = RelOptUtil.createDmlRowType(result.getKind(), typeFactory);
        break;
      case OTHER_DDL:
        x = getRowType(clazz);
        break;
      default:
        x = result.getRelNode().getRowType();
    }
    RelDataType jdbcType = makeStruct(typeFactory, x);
    final List<ColumnMetaData> columns =
        getColumnMetaDataList(typeFactory, x, jdbcType);
    Meta.Signature signature = new Meta.Signature(columns,
        sql,
        new ArrayList<AvaticaParameter>(),
        new HashMap<String, Object>(),
        Meta.CursorFactory.ARRAY,
        Meta.StatementType.SELECT);
    stmt.setSignature(signature);
    stmt.setResultSet(iterator);
    if (signature.statementType.canUpdate()) {
      metaResultSet = QuarkMetaResultSet.count(h.connectionId, h.id,
          ((Number) iterator.next()).intValue());
    } else {
      metaResultSet = QuarkMetaResultSet.create(h.connectionId, h.id,
          iterator, maxRowCount, signature);
    }
    return metaResultSet;
  }

  private List<ColumnMetaData> getColumnMetaDataList(
      JavaTypeFactory typeFactory, RelDataType x, RelDataType jdbcType) {
    final List<ColumnMetaData> columns = new ArrayList<>();
    for (Ord<RelDataTypeField> pair : Ord.zip(jdbcType.getFieldList())) {
      final RelDataTypeField field = pair.e;
      final RelDataType type = field.getType();
      final RelDataType fieldType =
          x.isStruct() ? x.getFieldList().get(pair.i).getType() : type;
      columns.add(
          metaData(typeFactory, columns.size(), field.getName(), type,
              fieldType, null));
    }
    return columns;
  }

  private ColumnMetaData metaData(JavaTypeFactory typeFactory, int ordinal,
                                  String fieldName, RelDataType type, RelDataType fieldType,
                                  List<String> origins) {
    final ColumnMetaData.AvaticaType avaticaType =
        avaticaType(typeFactory, type, fieldType);
    return new ColumnMetaData(
        ordinal,
        false,
        true,
        false,
        false,
        type.isNullable()
            ? DatabaseMetaData.columnNullable
            : DatabaseMetaData.columnNoNulls,
        true,
        type.getPrecision(),
        fieldName,
        origin(origins, 0),
        origin(origins, 2),
        getPrecision(type),
        getScale(type),
        origin(origins, 1),
        null,
        avaticaType,
        true,
        false,
        false,
        avaticaType.columnClassName());
  }

  private ColumnMetaData.AvaticaType avaticaType(JavaTypeFactory typeFactory,
                                                 RelDataType type, RelDataType fieldType) {
    final String typeName = getTypeName(type);
    if (type.getComponentType() != null) {
      final ColumnMetaData.AvaticaType componentType =
          avaticaType(typeFactory, type.getComponentType(), null);
      final Type clazz = typeFactory.getJavaClass(type.getComponentType());
      final ColumnMetaData.Rep rep = ColumnMetaData.Rep.of(clazz);
      assert rep != null;
      return ColumnMetaData.array(componentType, typeName, rep);
    } else {
      final int typeOrdinal = getTypeOrdinal(type);
      switch (typeOrdinal) {
        case Types.STRUCT:
          final List<ColumnMetaData> columns = new ArrayList<>();
          for (RelDataTypeField field : type.getFieldList()) {
            columns.add(
                metaData(typeFactory, field.getIndex(), field.getName(),
                    field.getType(), null, null));
          }
          return ColumnMetaData.struct(columns);
        default:
          final Type clazz =
              typeFactory.getJavaClass(Util.first(fieldType, type));
          final ColumnMetaData.Rep rep = ColumnMetaData.Rep.of(clazz);
          assert rep != null;
          return ColumnMetaData.scalar(typeOrdinal, typeName, rep);
      }
    }
  }
  private static String getTypeName(RelDataType type) {
    SqlTypeName sqlTypeName = type.getSqlTypeName();
    if (type instanceof RelDataTypeFactoryImpl.JavaType) {
      // We'd rather print "INTEGER" than "JavaType(int)".
      return sqlTypeName.getName();
    }
    switch (sqlTypeName) {
      case INTERVAL_YEAR_MONTH:
        // e.g. "INTERVAL_MONTH" or "INTERVAL_YEAR_MONTH"
        return "INTERVAL_"
            + type.getIntervalQualifier().toString().replace(' ', '_');
      default:
        return type.toString(); // e.g. "VARCHAR(10)", "INTEGER ARRAY"
    }
  }

  private static String origin(List<String> origins, int offsetFromEnd) {
    return origins == null || offsetFromEnd >= origins.size()
        ? null
        : origins.get(origins.size() - 1 - offsetFromEnd);
  }

  private int getTypeOrdinal(RelDataType type) {
    return type.getSqlTypeName().getJdbcOrdinal();
  }

  private static int getScale(RelDataType type) {
    return type.getScale() == RelDataType.SCALE_NOT_SPECIFIED
        ? 0
        : type.getScale();
  }

  private static int getPrecision(RelDataType type) {
    return type.getPrecision() == RelDataType.PRECISION_NOT_SPECIFIED
        ? 0
        : type.getPrecision();
  }

  private static RelDataType makeStruct(
      RelDataTypeFactory typeFactory,
      RelDataType type) {
    if (type.isStruct()) {
      return type;
    }
    return typeFactory.builder().add("$0", type).build();
  }
}