package util.sqlparse;

import bean.Column;
import bean.DataBase;
import com.alibaba.druid.sql.SQLUtils;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.util.StringUtils;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import util.JdbcUtil;
import util.sqlparse.visitor.common.Context;
import util.sqlparse.visitor.common.bean.FieldInfo;
import util.sqlparse.visitor.common.bean.SQLResult;
import util.sqlparse.visitor.common.bean.StatementType;
import util.sqlparse.visitor.common.bean.TableInfo;
import util.sqlparse.visitor.common.bean.ValueInfo;
import util.sqlparse.visitor.common.memo.FieldMemo;
import util.sqlparse.visitor.common.memo.TableMemo;
import util.sqlparse.visitor.common.scope.Scope;
import util.sqlparse.visitor.greenplum.visitor.FieldVisitor;
import util.sqlparse.visitor.greenplum.visitor.ScopeVisitor;
import util.sqlparse.visitor.spark.SQLParser;
import util.sqlparse.visitor.spark.SparkNameWrapper;
import util.sqlparse.visitor.spark.visitor.Desensitization;
import util.sqlparse.visitor.spark.visitor.SparkScopeDialector;

public class SparkSQLParse {
   public static final String DB_TYPE = "spark";
   public HashMap<String, List<Column>> tableList = new HashMap();
   public List<Column> columnList = new ArrayList();
   public Map<String, Object> whereMap = new HashMap();
   public List<Map<String, Object>> whereValueList = new ArrayList();
   public Map<String, String> tableMap = new HashMap();
   public Map<String, String> schemaMap = new HashMap();
   public Map<String, String> params;
   public List<Map<String, Object>> whereList = new ArrayList();

   public SparkSQLParse(Map<String, String> params) {
      this.params = params;
   }

   public SparkSQLParse() {
   }

   public String rowVisitControl(Map<String, Object> param) throws SQLException {
      Desensitization desensitization = new Desensitization();
      Map map = desensitization.parseRowVisit(param);
      return (String)map.get("newSql");
   }

   public String parseSql() {
      Desensitization desensitization = new Desensitization();
      Map map = desensitization.parseData(this.params);
      return (String)map.get("newSql");
   }

   public String parseSql_column() {
      Desensitization desensitization = new Desensitization();
      Map map = desensitization.parseColumn(this.params);
      return (String)map.get("newSql");
   }

   public String parseSql_table() {
      try {
         Desensitization desensitization = new Desensitization();
         desensitization.parseTable(this.params);
      } catch (Exception throwables) {
         throwables.printStackTrace();
      }

      return (String)this.params.get("newSql");
   }

   public String parseSql_data() {
      Desensitization desensitization = new Desensitization();
      Map map = desensitization.parseData(this.params);
      return (String)map.get("newSql");
   }

   public void parseSqlStructureAndReplace(Map<String, Object> paramMap, Map<String, String> repalceMap) throws SQLException {
      String sql = (String)this.params.get("sql");
      String dbType = (String)this.params.get("dbType");
      String sqlType = JdbcUtil.getSqlType(sql, dbType);
      if ("select".equals(sqlType) || "union".equals(sqlType) || "update".equals(sqlType) || "insert".equals(sqlType) || "delete".equals(sqlType)) {
         String permission = (String)this.params.get("permission");
         String url = (String)this.params.get("url");
         String schema = (String)this.params.get("schema");
         String ip_port = url.split("/")[2];
         SQLParser sqlParser = new SQLParser();
         DataBase dataBase = new DataBase();

         for(DataBase db : JdbcUtil.dataBaseList) {
            if (ip_port.equalsIgnoreCase(db.getIp_port())) {
               dataBase = db;
               break;
            }
         }

         SQLResult parse = sqlParser.parse(sql, dataBase, schema);
         List<FieldInfo> columnlist = new ArrayList();
         if ("masking".equalsIgnoreCase(permission)) {
            columnlist = parse.outputs;
         } else if ("full".equalsIgnoreCase(permission)) {
            columnlist = parse.fields;
         } else if (!"insert".equalsIgnoreCase(sqlType) && !"update".equalsIgnoreCase(sqlType)) {
            columnlist = parse.outputs;
         }

         for(FieldInfo fieldInfo : columnlist) {
            List<FieldInfo> fieldInfoList = new ArrayList();
            if (fieldInfo.getRelations().size() > 1) {
               fieldInfoList = fieldInfo.getRelations();
            } else {
               fieldInfoList.add(fieldInfo);
            }

            for(FieldInfo info : fieldInfoList) {
               String name = info.getName();
               Map<String, TableInfo> tables = info.getTables();
               String tableName = "";

               for(String s : tables.keySet()) {
                  tableName = s.split("\\.")[1];
               }

               Column column = new Column();
               column.setColumnName(tableName + "." + name);
               column.setColumnAlisa(info.getAlias());
               column.setSubQuery(false);
               this.columnList.add(column);
            }
         }

         SparkNameWrapper wrapper = new SparkNameWrapper(dataBase.simpleCache);

         for(TableInfo table : parse.tables) {
            TableMemo tb = table.getTable();
            String schName = wrapper.unwrapSchema(tb.schema);
            String tbName = wrapper.unwrapTable(schName, tb.name);
            this.tableMap.put(tbName, tbName);
            this.schemaMap.put(tbName, schName);
         }

         for(ValueInfo value : parse.values) {
            FieldMemo field = value.field;
            this.whereMap.put(field.name, value.value.toString());
            Map<String, Object> map = new HashMap();
            map.put("column", field.getAtomName());
            map.put("columnValue", value.value.toString());
            map.put("type", value.operator);
            this.whereList.add(map);
         }

         sqlParser.replace(parse.values, repalceMap);
         paramMap.put("whereList", this.whereList);
         paramMap.put("statement", parse.statement);
         paramMap.put("type", parse.sqlType.name().toLowerCase());
         paramMap.put("sql", parse.statement.toString());
         wrapper.unwrapFullColumns(this.columnList);
      }

      paramMap.put("tableMap", this.tableMap);
      paramMap.put("columnList", this.columnList);
      paramMap.put("schemaMap", this.schemaMap);
      paramMap.put("whereMap", this.whereMap);
      paramMap.put("whereList", this.whereList);
   }

   public void parseSqlStructure(Map<String, Object> paramMap) throws SQLException {
      String sql = (String)this.params.get("sql");
      String dbType = (String)this.params.get("dbType");
      boolean isEncryption = this.params.containsKey("encryption");
      List<SQLStatement> statements = SQLUtils.parseStatements(sql, "spark");
      if (statements != null && statements.size() != 0) {
         StatementType statementType = SQLParser.getStatementType((SQLStatement)statements.get(0));
         String sqlType = JdbcUtil.getSqlType(sql, dbType);
         if (statementType != StatementType.other) {
            String permission = (String)this.params.get("permission");
            String url = (String)this.params.get("url");
            String schema = (String)this.params.get("schema");
            String ip_port = url.split("/")[2];
            SQLParser sqlParser = new SQLParser();
            DataBase dataBase = new DataBase();

            for(DataBase db : JdbcUtil.dataBaseList) {
               if (ip_port.equalsIgnoreCase(db.getIp_port())) {
                  dataBase = db;
                  break;
               }
            }

            SQLResult parse = sqlParser.parse((SQLStatement)statements.get(0), dataBase, schema);
            List<FieldInfo> columnlist;
            if ("masking".equalsIgnoreCase(permission)) {
               columnlist = parse.outputs;
            } else if ("full".equalsIgnoreCase(permission)) {
               columnlist = parse.fields;
            } else if (isEncryption) {
               if (!"insert".equalsIgnoreCase(sqlType) && !"update".equalsIgnoreCase(sqlType)) {
                  columnlist = parse.outputs;
               } else {
                  columnlist = parse.fields;
               }
            } else {
               columnlist = parse.outputs;
            }

            SparkNameWrapper wrapper = new SparkNameWrapper(dataBase.simpleCache);
            Set<String> tabSet = new HashSet();

            for(FieldInfo fieldInfo : columnlist) {
               List<FieldInfo> fieldInfoList = new ArrayList();
               if (fieldInfo.getRelations().size() > 0) {
                  fieldInfoList = fieldInfo.getRelations();
               } else {
                  fieldInfoList.add(fieldInfo);
               }

               for(FieldInfo info : fieldInfoList) {
                  String name = info.getName();
                  if (!StringUtils.isEmpty(name)) {
                     Map<String, TableInfo> tables = info.getTables();
                     String tableName = "";
                     String schemaName = "";

                     for(String s : tables.keySet()) {
                        String[] segment = s.split("\\.");
                        if (!StringUtils.isEmpty(segment[0]) && !StringUtils.isEmpty(segment[1])) {
                           schemaName = segment[0];
                           tableName = segment[1];
                           break;
                        }
                     }

                     if (tableName.length() == 0 && info.getTables().size() > 0) {
                        for(Map.Entry<String, TableInfo> entry : info.getTables().entrySet()) {
                           TableInfo tbInfo = (TableInfo)entry.getValue();
                           if (tbInfo != null && tbInfo.getTable() != null && !StringUtils.isEmpty(tbInfo.getTable().schema) && !StringUtils.isEmpty(tbInfo.getTable().name)) {
                              tableName = tbInfo.getTable().name;
                              schemaName = tbInfo.getTable().schema;
                              break;
                           }
                        }
                     }

                     schemaName = wrapper.unwrapSchema(schemaName);
                     tableName = wrapper.unwrapTable(schemaName, tableName);
                     tabSet.add(tableName);
                     Column column = new Column();
                     column.setColumnName(tableName + "." + name);
                     column.setTable(tableName);
                     column.setSchema(schemaName);
                     column.setColumnAlisa(info.getAlias());
                     column.setSubQuery(false);
                     this.columnList.add(column);
                  }
               }
            }

            for(TableInfo table : parse.tables) {
               TableMemo tb = table.getTable();
               String schName = wrapper.unwrapSchema(tb.schema);
               String tbName = wrapper.unwrapTable(schName, tb.name);
               if (tabSet.contains(tbName)) {
                  this.tableMap.put(tbName, tbName);
                  this.schemaMap.put(tbName, schName);
               }
            }

            String type = this.getStatementType(parse.statementType);
            paramMap.put("type", type);

            for(ValueInfo value : parse.values) {
               FieldMemo field = value.field;
               this.whereMap.put(field.name, value.value.toString());
               Map<String, Object> map = new HashMap();
               map.put("column", value.getValueNames());
               map.put("columnValue", value.value.toString());
               map.put("type", value.operator);
               map.put("value", value);
               this.whereList.add(map);
            }

            paramMap.put("whereList", this.whereList);
            paramMap.put("statement", parse.statement);
            paramMap.put("type", statementType.toString().toLowerCase());
            paramMap.put("sql", parse.statement.toString());
            wrapper.unwrapFullColumns(this.columnList);
            paramMap.put("tableMap", this.tableMap);
            paramMap.put("columnList", this.columnList);
            paramMap.put("schemaMap", this.schemaMap);
            paramMap.put("whereMap", this.whereMap);
         }
      }
   }

   private String getStatementType(StatementType stmtType) {
      String type = "OTHER";
      switch (stmtType) {
         case select:
            type = "SELECT";
            break;
         case insert:
            type = "INSERT";
            break;
         case update:
            type = "UPDATE";
            break;
         case delete:
            type = "DELETE";
            break;
         case replace:
            type = "REPLACE";
            break;
         case grant:
            type = "GRANT";
            break;
         case revoke:
            type = "REVOKE";
            break;
         case alterView:
            type = "ALTERVIEW";
            break;
         case createView:
            type = "CREATEVIEW";
            break;
         case dropView:
            type = "DROPVIEW";
            break;
         case truncate:
            type = "TRUNCATE";
            break;
         case createIndex:
            type = "CREATEINDEX";
            break;
         case dropIndex:
            type = "DROPINDEX";
            break;
         case dropEvent:
         case dropFunction:
         case dropTable:
         case dropProc:
         case dropSequence:
            type = "DROP";
            break;
         case alterFunction:
         case alterEvent:
         case alterTable:
         case alterProc:
         case alterSequence:
         case alterIndex:
            type = "ALTER";
            break;
         case createEvent:
         case createTable:
         case createFunction:
         case createProc:
         case createSequence:
            type = "CREATE";
      }

      return type;
   }

   public List<String> parseWheres(String sql) {
      SQLParser parser = new SQLParser();
      return parser.where(sql);
   }

   public String formatSQL(String sql) {
      SQLParser parser = new SQLParser();
      return parser.format(sql);
   }

   public Scope getScope(String ip_port, String schema, List<SQLStatement> statements) {
      Context context = new Context();
      DataBase dataBase = new DataBase();

      for(DataBase db : JdbcUtil.dataBaseList) {
         if (ip_port.equalsIgnoreCase(db.getIp_port())) {
            dataBase = db;
            break;
         }
      }

      context.initialize(dataBase, schema, "spark");
      Scope scope = new Scope();
      scope.dialector = new SparkScopeDialector(scope);
      scope.context = context;
      scope.isCheckIsolated = true;
      ScopeVisitor scopeVisitor = new ScopeVisitor(scope);
      scopeVisitor.perform(statements);
      FieldVisitor fieldVisitor = new FieldVisitor(scope);
      fieldVisitor.perform(statements);
      return scope;
   }
}
