package util.sqlparse.visitor.spark;

import bean.DataBase;
import com.alibaba.druid.DbType;
import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.statement.SQLCreateFunctionStatement;
import com.alibaba.druid.sql.ast.statement.SQLDropFunctionStatement;
import com.alibaba.druid.sql.ast.statement.SQLDropTableStatement;
import com.alibaba.druid.sql.ast.statement.SQLDropViewStatement;
import com.alibaba.druid.sql.ast.statement.SQLSelectStatement;
import com.alibaba.druid.sql.ast.statement.SQLShowTablesStatement;
import com.alibaba.druid.sql.ast.statement.SQLTruncateStatement;
import com.alibaba.druid.sql.dialect.spark.ast.statement.SparkAlterDatabaseStatement;
import com.alibaba.druid.sql.dialect.spark.ast.statement.SparkAlterTableStatement;
import com.alibaba.druid.sql.dialect.spark.ast.statement.SparkAlterViewStatement;
import com.alibaba.druid.sql.dialect.spark.ast.statement.SparkCreateDatabaseStatement;
import com.alibaba.druid.sql.dialect.spark.ast.statement.SparkCreateTableStatement;
import com.alibaba.druid.sql.dialect.spark.ast.statement.SparkCreateViewStatement;
import com.alibaba.druid.sql.dialect.spark.ast.statement.SparkInsertStatement;
import com.alibaba.druid.sql.dialect.spark.ast.statement.SparkLoadStatement;
import com.alibaba.druid.sql.dialect.spark.ast.statement.SparkRepairTableStatement;
import java.util.List;
import java.util.Map;
import util.sqlparse.visitor.common.bean.SQLResult;
import util.sqlparse.visitor.common.bean.StatementType;
import util.sqlparse.visitor.common.bean.ValueInfo;
import util.sqlparse.visitor.spark.parser.SQLStatementParser;
import util.sqlparse.visitor.spark.visitor.FormatVisitor;
import util.sqlparse.visitor.spark.visitor.ValueReplacer;
import util.sqlparse.visitor.spark.visitor.WhereVisitor;

public class SQLParser {
   public SQLResult parse(String sql, DataBase dataBase, String schema) {
      if (sql != null && sql.length() != 0) {
         return SQLStatementParser.parse(sql, dataBase, schema);
      } else {
         throw new IllegalArgumentException("sql");
      }
   }

   public SQLResult parse(SQLStatement statement, DataBase dataBase, String schema) {
      if (statement == null) {
         throw new IllegalArgumentException("statement");
      } else {
         return SQLStatementParser.parse(statement, dataBase, schema);
      }
   }

   public static StatementType getStatementType(SQLStatement statement) {
      if (statement instanceof SQLSelectStatement) {
         return StatementType.select;
      } else if (statement instanceof SparkInsertStatement) {
         return ((SparkInsertStatement)statement).isOverwrite() ? StatementType.update : StatementType.insert;
      } else if (statement instanceof SparkLoadStatement) {
         return StatementType.loadData;
      } else if (statement instanceof SparkAlterDatabaseStatement) {
         return StatementType.other;
      } else if (statement instanceof SparkAlterTableStatement) {
         return StatementType.alterTable;
      } else if (statement instanceof SparkAlterViewStatement) {
         return StatementType.alterView;
      } else if (statement instanceof SQLTruncateStatement) {
         return StatementType.truncate;
      } else if (statement instanceof SparkCreateDatabaseStatement) {
         return StatementType.createSchema;
      } else if (statement instanceof SparkCreateTableStatement) {
         return StatementType.createTable;
      } else if (statement instanceof SparkCreateViewStatement) {
         return StatementType.createView;
      } else if (statement instanceof SQLCreateFunctionStatement) {
         return StatementType.createFunction;
      } else if (statement instanceof SparkRepairTableStatement) {
         return StatementType.repairTable;
      } else if (statement instanceof SparkCreateViewStatement) {
         return StatementType.createView;
      } else if (statement instanceof SQLDropFunctionStatement) {
         return StatementType.dropFunction;
      } else if (statement instanceof SQLDropTableStatement) {
         return StatementType.dropTable;
      } else if (statement instanceof SQLDropViewStatement) {
         return StatementType.dropView;
      } else {
         return statement instanceof SQLShowTablesStatement ? StatementType.showTables : StatementType.other;
      }
   }

   public void replace(List<ValueInfo> values, Map<String, String> replaceValue) {
      ValueReplacer replacer = new ValueReplacer(replaceValue);
      replacer.replaceValues(values);
   }

   public void replace(List<ValueInfo> values, List<Map> whereList) {
      ValueReplacer replacer = new ValueReplacer(whereList);
      replacer.replaceValuesWithList(values);
   }

   public List<String> where(String sql) {
      List<SQLStatement> statements = SQLUtils.parseStatements(sql, DbType.spark);
      WhereVisitor visitor = new WhereVisitor();
      visitor.perform(statements);
      return visitor.getWheres();
   }

   public String format(String sql) {
      List<SQLStatement> statements = SQLUtils.parseStatements(sql, DbType.spark);
      if (statements != null && statements.size() != 0) {
         FormatVisitor visitor = new FormatVisitor();
         visitor.perform(statements);
         return ((SQLStatement)statements.get(0)).toString();
      } else {
         return "";
      }
   }
}
