package com.alibaba.druid.sql.dialect.kingbase.parser;

import com.alibaba.druid.DbType;
import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLName;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.statement.SQLDeleteStatement;
import com.alibaba.druid.sql.ast.statement.SQLInsertStatement;
import com.alibaba.druid.sql.ast.statement.SQLSelect;
import com.alibaba.druid.sql.ast.statement.SQLSelectItem;
import com.alibaba.druid.sql.ast.statement.SQLSelectStatement;
import com.alibaba.druid.sql.ast.statement.SQLTableSource;
import com.alibaba.druid.sql.ast.statement.SQLUpdateStatement;
import com.alibaba.druid.sql.ast.statement.SQLWithSubqueryClause;
import com.alibaba.druid.sql.dialect.kingbase.ast.clause.KingbaseWithSubqueryEntry;
import com.alibaba.druid.sql.dialect.kingbase.ast.statement.KingbaseDeleteStatement;
import com.alibaba.druid.sql.dialect.kingbase.ast.statement.KingbaseInsertStatement;
import com.alibaba.druid.sql.dialect.kingbase.ast.statement.KingbaseMultiInsertStatement;
import com.alibaba.druid.sql.dialect.kingbase.ast.statement.KingbaseUpdateStatement;
import com.alibaba.druid.sql.parser.ParserException;
import com.alibaba.druid.sql.parser.SQLExprParser;
import com.alibaba.druid.sql.parser.SQLParserFeature;
import com.alibaba.druid.sql.parser.SQLStatementParser;
import com.alibaba.druid.sql.parser.Token;
import com.alibaba.druid.util.FnvHash;
import java.util.ArrayList;
import java.util.List;

public class KingbaseStatementParser extends SQLStatementParser {
   public KingbaseStatementParser(String sql, SQLParserFeature... features) {
      this(sql, DbType.kingbase, features);
   }

   public KingbaseStatementParser(String sql, DbType dbType, SQLParserFeature... features) {
      super((SQLExprParser)(new KingbaseExprParser(sql, features)));
      this.dbType = DbType.kingbase;
   }

   public SQLStatement parseSelect() {
      KingbaseStatementSelectParser selectParser = this.createSQLSelectParser();
      SQLSelect select = selectParser.select();
      return (SQLStatement)(selectParser.returningFlag ? selectParser.updateStmt : new SQLSelectStatement(select, DbType.kingbase));
   }

   public SQLUpdateStatement parseUpdateStatement() {
      KingbaseUpdateStatement udpateStatement = this.createUpdateStatement();
      if (this.lexer.token() == Token.UPDATE) {
         this.lexer.nextToken();
         SQLTableSource tableSource = this.exprParser.createSelectParser().parseTableSource();
         udpateStatement.setTableSource(tableSource);
      }

      this.parseUpdateSet(udpateStatement);
      if (this.lexer.token() == Token.WHERE) {
         this.lexer.nextToken();
         udpateStatement.setWhere(this.exprParser.expr());
      }

      if (this.lexer.token() == Token.RETURN || this.lexer.token() == Token.RETURNING) {
         this.lexer.nextToken();
         List<SQLSelectItem> returnList = new ArrayList();
         udpateStatement.setReturnList(returnList);
         this.parseReturnList(returnList);
      }

      return udpateStatement;
   }

   protected void parseReturnList(List<SQLSelectItem> returnList) {
      while(true) {
         SQLSelectItem selectItem = this.exprParser.parseSelectItem();
         returnList.add(selectItem);
         if (this.lexer.token() != Token.COMMA) {
            return;
         }

         this.lexer.nextToken();
      }
   }

   public KingbaseStatementSelectParser createSQLSelectParser() {
      return new KingbaseStatementSelectParser(this.exprParser, this.selectListCache);
   }

   public SQLWithSubqueryClause parseWithQuery() {
      SQLWithSubqueryClause withQueryClause = new SQLWithSubqueryClause();
      if (this.lexer.hasComment() && this.lexer.isKeepComments()) {
         withQueryClause.addBeforeComment(this.lexer.readAndResetComments());
      }

      this.accept(Token.WITH);
      if (this.lexer.token() == Token.RECURSIVE || this.lexer.identifierEquals("RECURSIVE")) {
         this.lexer.nextToken();
         withQueryClause.setRecursive(true);
      }

      while(true) {
         KingbaseWithSubqueryEntry entry = new KingbaseWithSubqueryEntry();
         entry.setParent(withQueryClause);
         String alias = this.lexer.stringVal();
         this.lexer.nextToken();
         entry.setAlias(alias);
         if (this.lexer.token() == Token.LPAREN) {
            this.lexer.nextToken();
            this.exprParser.names(entry.getColumns());
            this.accept(Token.RPAREN);
         }

         this.accept(Token.AS);
         this.parseMaterialized(entry);
         this.accept(Token.LPAREN);
         switch (this.lexer.token()) {
            case VALUES:
            case WITH:
            case TABLE:
            case SELECT:
            case LPAREN:
            case FROM:
               entry.setSubQuery(this.createSQLSelectParser().select());
               break;
            case INSERT:
               entry.setReturningStatement(this.parseInsert());
               break;
            case UPDATE:
               entry.setReturningStatement(this.parseUpdateStatement());
               break;
            case DELETE:
               entry.setReturningStatement(this.parseDeleteStatement());
         }

         this.accept(Token.RPAREN);
         withQueryClause.addEntry(entry);
         if (this.lexer.token() != Token.COMMA) {
            return withQueryClause;
         }

         this.lexer.nextToken();
      }
   }

   private void parseMaterialized(KingbaseWithSubqueryEntry entry) {
      if (this.lexer.token() == Token.NOT) {
         this.lexer.nextToken();
         if (this.lexer.identifierEquals(FnvHash.Constants.MATERIALIZED)) {
            this.lexer.nextToken();
            entry.setType(KingbaseWithSubqueryEntry.Type.NOT_MATERIALIZED);
         }
      } else if (this.lexer.identifierEquals(FnvHash.Constants.MATERIALIZED)) {
         this.lexer.nextToken();
         entry.setType(KingbaseWithSubqueryEntry.Type.MATERIALIZED);
      }

   }

   protected KingbaseUpdateStatement createUpdateStatement() {
      return new KingbaseUpdateStatement();
   }

   public SQLStatement parseInsert() {
      if (this.lexer.token() == Token.INSERT) {
         this.accept(Token.INSERT);
      }

      if (this.lexer.token() == Token.INTO) {
         KingbaseInsertStatement stmt = new KingbaseInsertStatement();
         this.parseInsert0(stmt);
         return stmt;
      } else {
         KingbaseMultiInsertStatement stmt = this.parseMultiInsert();
         return stmt;
      }
   }

   protected void parseInsert0(KingbaseInsertStatement insertStatement) {
      this.parseInsert0(insertStatement, true);
   }

   protected void parseInsert0(KingbaseInsertStatement insertStatement, boolean acceptSubQuery) {
      if (this.lexer.token() == Token.INTO) {
         this.lexer.nextToken();
         if (this.lexer.token() == Token.LPAREN) {
            SQLSelect select = this.createSQLSelectParser().select();
            insertStatement.setSelect(select);
         } else {
            this.parseTable(insertStatement);
            if (this.lexer.token() == Token.LITERAL_ALIAS) {
               insertStatement.setAlias(this.tableAlias());
            }
         }

         this.parseInsert0_hinits(insertStatement);
         if (this.lexer.token() == Token.IDENTIFIER || this.lexer.token() == Token.AS) {
            if (this.lexer.token() == Token.AS) {
               this.lexer.nextToken();
            }

            insertStatement.setAlias(this.lexer.stringVal());
            this.lexer.nextToken();
         }
      }

      if (this.lexer.token() == Token.OVERRIDING) {
         String str = this.lexer.stringVal() + " ";
         this.lexer.nextToken();
         str = str + this.lexer.stringVal() + " ";
         this.lexer.nextToken();
         str = str + this.lexer.stringVal() + " ";
         this.lexer.nextToken();
         insertStatement.setOverridingStr(str);
      }

      if (this.lexer.token() == Token.SET) {
         this.lexer.nextToken();
         insertStatement.setSetWay(true);
         this.parseSetColumn(insertStatement);
      }

      if (this.lexer.token() == Token.LPAREN) {
         this.lexer.nextToken();
         this.parseInsertColumns(insertStatement);
         this.accept(Token.RPAREN);
      }

      if (this.lexer.token() == Token.VALUES) {
         this.lexer.nextToken();

         while(true) {
            if (this.lexer.token() == Token.LPAREN) {
               this.lexer.nextToken();
               SQLInsertStatement.ValuesClause values = new SQLInsertStatement.ValuesClause();
               this.exprParser.exprList(values.getValues(), values);
               insertStatement.addValueCause(values);
               this.accept(Token.RPAREN);
            } else {
               SQLInsertStatement.ValuesClause values = new SQLInsertStatement.ValuesClause();
               SQLExpr value = this.exprParser.expr();
               values.addValue(value);
               insertStatement.addValueCause(values);
            }

            if (this.lexer.token() != Token.COMMA) {
               break;
            }

            this.lexer.nextToken();
         }
      } else if (!acceptSubQuery || this.lexer.token() != Token.SELECT && this.lexer.token() != Token.LPAREN) {
         if (this.lexer.identifierEquals(FnvHash.Constants.VALUE)) {
            throw new ParserException("'values' expected, but 'value'. " + this.lexer.info());
         }

         if (this.lexer.token() == Token.DEFAULT) {
            insertStatement.setDefaultValue(true);
            this.lexer.nextToken();
            this.accept(Token.VALUES);
         }
      } else {
         SQLSelect select = this.createSQLSelectParser().select();
         insertStatement.setQuery(select);
      }

      if (this.lexer.token() == Token.ON) {
         this.lexer.nextToken();
         this.accept(Token.CONFLICT);
         insertStatement.setOnConflict(true);
      }

      if (this.lexer.token() == Token.LPAREN) {
         List<SQLExpr> doUpdateList = insertStatement.getDoUpdateList();

         do {
            this.lexer.nextToken();
            SQLExpr primary = this.exprParser.primary();
            doUpdateList.add(primary);
         } while(this.lexer.token() == Token.COMMA && this.lexer.token() != Token.RPAREN);

         this.accept(Token.RPAREN);
      }

      if (this.lexer.token() == Token.DO) {
         this.lexer.nextToken();
         if (this.lexer.token() == Token.NOTHING) {
            this.lexer.nextToken();
            insertStatement.setDoNothing(true);
         } else if (this.lexer.token() == Token.UPDATE) {
            this.lexer.nextToken();
            this.lexer.nextToken();
            KingbaseInsertStatement doUpdateStatement = new KingbaseInsertStatement();
            doUpdateStatement.setTableSource(insertStatement.getTableSource());
            doUpdateStatement.setParent(insertStatement);
            insertStatement.setDoUpdateStatement(doUpdateStatement);
            this.parseSetColumn(doUpdateStatement);
            insertStatement.setDoUpdate(true);
            if (this.lexer.token() == Token.WHERE) {
               this.lexer.nextToken();
               doUpdateStatement.setDoUpdateWhere(this.exprParser.expr());
            }
         }
      }

      if (this.lexer.token() == Token.RETURN || this.lexer.token() == Token.RETURNING) {
         this.lexer.nextToken();
         List<SQLSelectItem> returnList = new ArrayList();
         insertStatement.setReturnList(returnList);
         this.parseReturnList(returnList);
      }

   }

   public void parseSetColumn(KingbaseInsertStatement insertStatement) {
      insertStatement.setSet(true);
      if (this.lexer.token() == Token.LPAREN) {
         this.lexer.nextToken();
         this.parseInsertColumns(insertStatement);
         this.accept(Token.RPAREN);
      } else {
         SQLExpr expr = this.exprParser.primary();
         expr.setParent(insertStatement);
         insertStatement.getColumns().add(expr);
      }

      if (this.lexer.token() == Token.EQ) {
         this.lexer.nextToken();
      }

      List<SQLInsertStatement.ValuesClause> valuesList = insertStatement.getValuesList();
      SQLInsertStatement.ValuesClause values = null;
      if (valuesList.size() <= 0) {
         values = new SQLInsertStatement.ValuesClause();
         values.setParent(insertStatement);
         insertStatement.addValueCause(values);
      } else {
         values = (SQLInsertStatement.ValuesClause)valuesList.get(0);
      }

      if (this.lexer.token() == Token.LPAREN) {
         this.lexer.nextToken();
         this.exprParser.exprList(values.getValues(), values);
         this.accept(Token.RPAREN);
      } else {
         SQLExpr expr = this.exprParser.primary();
         expr.setParent(values);
         values.getValues().add(expr);
      }

      if (this.lexer.token() == Token.COMMA) {
         this.lexer.nextToken();
         this.parseSetColumn(insertStatement);
      }

   }

   public SQLDeleteStatement parseDeleteStatement() {
      KingbaseDeleteStatement deleteStatement = new KingbaseDeleteStatement();
      if (this.lexer.token() == Token.DELETE) {
         this.lexer.nextToken();
         if (this.lexer.token() == Token.FROM) {
            this.lexer.nextToken();
         }

         if (this.lexer.token() == Token.ONLY) {
            this.lexer.nextToken();
            deleteStatement.setOnly(true);
         }

         if (this.lexer.token() == Token.COMMENT) {
            this.lexer.nextToken();
         }

         SQLName tableName = this.exprParser.name();
         deleteStatement.setTableName(tableName);
         if (this.lexer.token() == Token.STAR) {
            deleteStatement.setInherit(true);
            this.lexer.nextToken();
         }

         if (this.lexer.token() == Token.FROM) {
            this.lexer.nextToken();
            SQLTableSource tableSource = this.createSQLSelectParser().parseTableSource();
            deleteStatement.setFrom(tableSource);
         }

         if (this.lexer.token() == Token.AS) {
            this.lexer.nextToken();
            String alias = this.exprParser.name().getSimpleName();
            deleteStatement.setAlias(alias);
         }

         if (this.lexer.token() == Token.IDENTIFIER) {
            String alias = this.exprParser.name().getSimpleName();
            deleteStatement.setAlias(alias);
         }

         if (this.lexer.token() == Token.USING) {
            this.lexer.nextToken();
            KingbaseStatementSelectParser selectParser = this.createSQLSelectParser();
            SQLTableSource tableSource = selectParser.parseTableSource();
            deleteStatement.setUsing(tableSource);
         }
      }

      if (this.lexer.token() == Token.WHERE) {
         this.lexer.nextToken();
         SQLExpr where = this.exprParser.expr();
         deleteStatement.setWhere(where);
      }

      if (this.lexer.token() == Token.RETURN || this.lexer.token() == Token.RETURNING) {
         this.lexer.nextToken();
         List<SQLSelectItem> returnList = new ArrayList();
         deleteStatement.setReturnList(returnList);
         this.parseReturnList(returnList);
      }

      return deleteStatement;
   }

   public SQLStatement parseWith() {
      SQLWithSubqueryClause with = this.parseWithQuery();
      if (this.lexer.token() == Token.INSERT) {
         KingbaseInsertStatement stmt = (KingbaseInsertStatement)this.parseInsert();
         stmt.setWith(with);
         return stmt;
      } else if (this.lexer.token() == Token.SELECT) {
         SQLSelectStatement stmt = (SQLSelectStatement)this.parseSelect();
         stmt.getSelect().setWithSubQuery(with);
         return stmt;
      } else if (this.lexer.token() == Token.DELETE) {
         KingbaseDeleteStatement stmt = (KingbaseDeleteStatement)this.parseDeleteStatement();
         stmt.setWith(with);
         return stmt;
      } else if (this.lexer.token() == Token.UPDATE) {
         KingbaseUpdateStatement stmt = (KingbaseUpdateStatement)this.parseUpdateStatement();
         stmt.setWith(with);
         return stmt;
      } else {
         throw new ParserException("TODO. " + this.lexer.info());
      }
   }

   public KingbaseMultiInsertStatement parseMultiInsert() {
      KingbaseMultiInsertStatement stmt = new KingbaseMultiInsertStatement();
      if (this.lexer.token() == Token.ALL) {
         this.lexer.nextToken();
         stmt.setOption(KingbaseMultiInsertStatement.Option.ALL);
      } else if (this.lexer.token() == Token.FIRST || this.lexer.identifierEquals("FIRST")) {
         this.lexer.nextToken();
         stmt.setOption(KingbaseMultiInsertStatement.Option.FIRST);
      }

      while(this.lexer.token() == Token.INTO) {
         KingbaseMultiInsertStatement.InsertIntoClause clause = new KingbaseMultiInsertStatement.InsertIntoClause();
         boolean acceptSubQuery = stmt.getEntries().size() == 0;
         this.parseInsert0(clause, acceptSubQuery);
         stmt.addEntry(clause);
      }

      SQLSelect query = null;
      if (this.lexer.token() == Token.WHEN) {
         KingbaseMultiInsertStatement.ConditionalInsertClause clause = new KingbaseMultiInsertStatement.ConditionalInsertClause();

         while(this.lexer.token() == Token.WHEN) {
            this.lexer.nextToken();
            KingbaseMultiInsertStatement.ConditionalInsertClauseItem item = new KingbaseMultiInsertStatement.ConditionalInsertClauseItem();
            item.setWhen(this.exprParser.expr());
            this.accept(Token.THEN);
            KingbaseMultiInsertStatement.InsertIntoClause insertInto = new KingbaseMultiInsertStatement.InsertIntoClause();
            this.parseInsert0(insertInto);
            item.setThen(insertInto);
            clause.addItem(item);
         }

         if (clause.getItems().size() > 0) {
            KingbaseMultiInsertStatement.ConditionalInsertClauseItem item = (KingbaseMultiInsertStatement.ConditionalInsertClauseItem)clause.getItems().get(clause.getItems().size() - 1);
            KingbaseMultiInsertStatement.InsertIntoClause insert = item.getThen();
            if (insert.getQuery() != null) {
               query = insert.getQuery();
               insert.setQueryNull();
            }
         }

         if (this.lexer.token() == Token.ELSE) {
            this.lexer.nextToken();
            KingbaseMultiInsertStatement.InsertIntoClause insertInto = new KingbaseMultiInsertStatement.InsertIntoClause();
            this.parseInsert0(insertInto, false);
            clause.setElseItem(insertInto);
            if (insertInto.getQuery() != null) {
               query = insertInto.getQuery();
               insertInto.setQueryNull();
            }
         }

         stmt.addEntry(clause);
      }

      if (query == null) {
         SQLSelect subQuery = this.createSQLSelectParser().select();
         stmt.setSubQuery(subQuery);
         subQuery.setParent(stmt);
      } else {
         stmt.setSubQuery(query);
      }

      return stmt;
   }
}
