package com.test.parsersandbox;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import org.apache.calcite.runtime.CalciteContextException;
import org.apache.calcite.sql.SqlCall;
import org.apache.calcite.sql.SqlIdentifier;
import org.apache.calcite.sql.SqlNode;
import org.apache.calcite.sql.SqlSelect;
import org.apache.calcite.sql.SqlWith;
import org.apache.calcite.sql.SqlWithItem;
import org.apache.calcite.sql.validate.SqlQualified;
import org.apache.calcite.sql.validate.SqlScopedShuttle;
import org.apache.calcite.sql.validate.SqlValidatorImpl;
import org.apache.calcite.sql.validate.SqlValidatorNamespace;
import org.apache.calcite.sql.validate.SqlValidatorScope;
import org.checkerframework.checker.nullness.qual.Nullable;

public class AccumulatingFieldsShuttle extends SqlScopedShuttle {
  
  private List<List<String>> accumulatedFields = new ArrayList<>();

  public AccumulatingFieldsShuttle(SqlValidatorScope initialScope) {
    super(initialScope);
  }

  public static AccumulatingFieldsShuttle of(SqlValidatorImpl validator, SqlNode sql) {
    if (sql instanceof SqlSelect) {
      SqlSelect select = (SqlSelect) sql;
      final SqlValidatorScope selectScope = validator.getSelectScope(select);
      return new AccumulatingFieldsShuttle(validator, selectScope);
    } else if (sql instanceof SqlWith) {
      SqlWith withSql = (SqlWith) sql;
      SqlWithItem withItem = (SqlWithItem) withSql.withList.get(0);
      final SqlValidatorScope withScope = validator.getWithScope(withItem);
      return new AccumulatingFieldsShuttle(validator, withScope);
    }
    throw new UnsupportedOperationException(sql.getKind().toString());
  }

  @Override
  public @Nullable SqlNode visit(SqlIdentifier id) {
    try {
      SqlQualified qualified = getScope().fullyQualify(id);
      SqlValidatorNamespace namespace = qualified.namespace;
      if (namespace != null && namespace.getTable() != null) {
        accumulatedFields.add(
            Stream.concat(
                    namespace.getTable().getQualifiedName().stream(), qualified.suffix().stream())
                .collect(Collectors.toList()));
      }

    } catch (CalciteContextException exc) {
      // Ignore identifiers not found in the catalog for now
    }
    return super.visit(id);
  }

  public List<List<String>> getAccumulatedFields() {
    return accumulatedFields;
  }

  @Override
  protected @Nullable SqlNode visitScoped(SqlCall call) {
    return call.getOperator().acceptCall(this, call);
  }
}