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

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.SQLAssignItem;
import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
import com.alibaba.druid.sql.ast.statement.SQLInsertStatement;
import com.alibaba.druid.sql.ast.statement.SQLSelect;
import com.alibaba.druid.sql.ast.statement.SQLSubqueryTableSource;
import com.alibaba.druid.sql.ast.statement.SQLTableSource;
import com.alibaba.druid.sql.dialect.hive.ast.HiveInsert;
import com.alibaba.druid.sql.dialect.hive.ast.HiveMultiInsertStatement;
import com.alibaba.druid.sql.dialect.hive.ast.stmt.HiveInsertStatement;
import com.alibaba.druid.sql.parser.Lexer;
import com.alibaba.druid.sql.parser.SQLExprParser;
import com.alibaba.druid.sql.parser.SQLSelectParser;
import com.alibaba.druid.sql.parser.SQLStatementParser;
import com.alibaba.druid.sql.parser.Token;

public class HiveInsertParser extends SQLStatementParser {
   public HiveInsertParser(SQLExprParser exprParser) {
      super(exprParser);
   }

   public HiveInsertParser(Lexer lexer) {
      super((SQLExprParser)(new HiveExprParser(lexer)));
   }

   public SQLStatement parseInsert() {
      if (this.lexer.token() != Token.FROM) {
         return this.parseHiveInsertStmt();
      } else {
         this.lexer.nextToken();
         HiveMultiInsertStatement stmt = new HiveMultiInsertStatement();
         if (this.lexer.token() == Token.IDENTIFIER) {
            SQLName tableName = this.exprParser.name();
            SQLExprTableSource from = new SQLExprTableSource(tableName);
            SQLTableSource tableSource = this.createSQLSelectParser().parseTableSourceRest(from);
            stmt.setFrom(tableSource);
            if (this.lexer.token() == Token.IDENTIFIER) {
               from.setAlias(this.lexer.stringVal());
               this.lexer.nextToken();
            }
         } else {
            this.accept(Token.LPAREN);
            SQLSelectParser selectParser = this.createSQLSelectParser();
            SQLSelect select = selectParser.select();
            this.accept(Token.RPAREN);
            String alias = this.lexer.stringVal();
            this.accept(Token.IDENTIFIER);
            SQLTableSource from = new SQLSubqueryTableSource(select, alias);
            switch (this.lexer.token()) {
               case LEFT:
               case RIGHT:
               case FULL:
               case JOIN:
                  from = selectParser.parseTableSourceRest(from);
               default:
                  stmt.setFrom(from);
            }
         }

         do {
            HiveInsert insert = this.parseHiveInsert();
            stmt.addItem(insert);
         } while(this.lexer.token() == Token.INSERT);

         return stmt;
      }
   }

   public HiveInsertStatement parseHiveInsertStmt() {
      HiveInsertStatement insert = new HiveInsertStatement();
      insert.setDbType(this.dbType);
      if (this.lexer.isKeepComments() && this.lexer.hasComment()) {
         insert.addBeforeComment(this.lexer.readAndResetComments());
      }

      SQLSelectParser selectParser = this.createSQLSelectParser();
      this.accept(Token.INSERT);
      if (this.lexer.token() == Token.INTO) {
         this.lexer.nextToken();
      } else {
         this.accept(Token.OVERWRITE);
         insert.setOverwrite(true);
      }

      if (this.lexer.token() == Token.LOCAL) {
         insert.local = true;
         this.lexer.nextToken();
      }

      if (this.lexer.token() == Token.DIRECTORY) {
         insert.setDirSouece(true);
         this.lexer.nextToken();
      }

      if (this.lexer.token() == Token.DIRECTORY) {
         insert.setDirSouece(true);
         this.lexer.nextToken();
      }

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

      insert.setTableSource(this.exprParser.name());
      boolean columnsParsed = false;
      if (this.lexer.token() == Token.LPAREN) {
         Lexer.SavePoint mark = this.lexer.mark();
         this.lexer.nextToken();
         if (this.lexer.token() == Token.SELECT) {
            this.lexer.reset(mark);
         } else {
            this.parseInsertColumns(insert);
            columnsParsed = true;
            this.accept(Token.RPAREN);
         }
      }

      if (this.lexer.token() == Token.PARTITION) {
         this.lexer.nextToken();
         this.accept(Token.LPAREN);

         while(true) {
            SQLAssignItem ptExpr = new SQLAssignItem();
            ptExpr.setTarget(this.exprParser.name());
            if (this.lexer.token() == Token.EQ || this.lexer.token() == Token.EQEQ) {
               this.lexer.nextTokenValue();
               SQLExpr ptValue = this.exprParser.expr();
               ptExpr.setValue(ptValue);
            }

            insert.addPartition(ptExpr);
            if (this.lexer.token() != Token.COMMA) {
               this.accept(Token.RPAREN);
               break;
            }

            this.lexer.nextToken();
         }
      }

      if (!columnsParsed && this.lexer.token() == Token.LPAREN) {
         Lexer.SavePoint savePoint = this.lexer.mark();
         this.lexer.nextToken();
         if (this.lexer.token() != Token.SELECT) {
            this.parseInsertColumns(insert);
            this.accept(Token.RPAREN);
         } else {
            this.lexer.reset(savePoint);
         }
      }

      if (this.lexer.token() == Token.IF) {
         this.lexer.nextToken();
         this.accept(Token.NOT);
         this.accept(Token.EXISTS);
         insert.setIfNotExists(true);
      }

      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);
               insert.addValueCause(values);
               this.accept(Token.RPAREN);
            }

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

            this.lexer.nextToken();
         }
      } else {
         SQLSelect query = selectParser.select();
         insert.setQuery(query);
      }

      return insert;
   }
}
