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

import com.alibaba.druid.DbType;
import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLObject;
import com.alibaba.druid.sql.ast.expr.SQLSizeExpr;
import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
import com.alibaba.druid.sql.ast.statement.SQLSelectQuery;
import com.alibaba.druid.sql.ast.statement.SQLTableSampling;
import com.alibaba.druid.sql.ast.statement.SQLTableSource;
import com.alibaba.druid.sql.dialect.hive.ast.stmt.HiveSelectQueryBlock;
import com.alibaba.druid.sql.parser.Lexer;
import com.alibaba.druid.sql.parser.SQLExprParser;
import com.alibaba.druid.sql.parser.SQLSelectListCache;
import com.alibaba.druid.sql.parser.SQLSelectParser;
import com.alibaba.druid.sql.parser.Token;
import com.alibaba.druid.util.FnvHash;

public class HiveSelectParser extends SQLSelectParser {
   public HiveSelectParser(SQLExprParser exprParser) {
      super(exprParser);
   }

   public HiveSelectParser(SQLExprParser exprParser, SQLSelectListCache selectListCache) {
      super(exprParser, selectListCache);
   }

   public HiveSelectParser(String sql) {
      this((SQLExprParser)(new HiveExprParser(sql)));
   }

   protected SQLExprParser createExprParser() {
      return new HiveExprParser(this.lexer);
   }

   public SQLTableSource parseTableSourceRest(SQLTableSource tableSource) {
      if (this.lexer.identifierEquals(FnvHash.Constants.TABLESAMPLE) && tableSource instanceof SQLExprTableSource) {
         Lexer.SavePoint mark = this.lexer.mark();
         this.lexer.nextToken();
         if (this.lexer.token() == Token.LPAREN) {
            this.lexer.nextToken();
            SQLTableSampling sampling = new SQLTableSampling();
            if (this.lexer.identifierEquals(FnvHash.Constants.BUCKET)) {
               this.lexer.nextToken();
               SQLExpr bucket = this.exprParser.primary();
               sampling.setBucket(bucket);
               if (this.lexer.token() == Token.OUT) {
                  this.lexer.nextToken();
                  this.accept(Token.OF);
                  SQLExpr outOf = this.exprParser.primary();
                  sampling.setOutOf(outOf);
               }

               if (this.lexer.token() == Token.ON) {
                  this.lexer.nextToken();
                  SQLExpr on = this.exprParser.expr();
                  sampling.setOn(on);
               }
            }

            if (this.lexer.token() == Token.LITERAL_INT || this.lexer.token() == Token.LITERAL_FLOAT) {
               SQLExpr val = this.exprParser.primary();
               if (this.lexer.identifierEquals(FnvHash.Constants.ROWS)) {
                  this.lexer.nextToken();
                  sampling.setRows(val);
               } else {
                  this.acceptIdentifier("PERCENT");
                  sampling.setPercent(val);
               }
            }

            if (this.lexer.token() == Token.IDENTIFIER) {
               String strVal = this.lexer.stringVal();
               char first = strVal.charAt(0);
               char last = strVal.charAt(strVal.length() - 1);
               if (last >= 'a' && last <= 'z') {
                  last = (char)(last - 32);
               }

               boolean match = false;
               if (first == '.' || first >= '0' && first <= '9') {
                  switch (last) {
                     case 'B':
                     case 'G':
                     case 'K':
                     case 'M':
                     case 'P':
                     case 'T':
                        match = true;
                     case 'C':
                     case 'D':
                     case 'E':
                     case 'F':
                     case 'H':
                     case 'I':
                     case 'J':
                     case 'L':
                     case 'N':
                     case 'O':
                     case 'Q':
                     case 'R':
                     case 'S':
                  }
               }

               SQLSizeExpr size = new SQLSizeExpr(strVal.substring(0, strVal.length() - 2), last);
               sampling.setByteLength(size);
               this.lexer.nextToken();
            }

            SQLExprTableSource table = (SQLExprTableSource)tableSource;
            table.setSampling(sampling);
            this.accept(Token.RPAREN);
         } else {
            this.lexer.reset(mark);
         }
      }

      return super.parseTableSourceRest(tableSource);
   }

   public SQLSelectQuery query(SQLObject parent, boolean acceptUnion) {
      if (this.lexer.token() == Token.LPAREN) {
         this.lexer.nextToken();
         SQLSelectQuery select = this.query();
         this.accept(Token.RPAREN);
         return this.queryRest(select, acceptUnion);
      } else if (this.lexer.token() == Token.VALUES) {
         return this.valuesQuery(acceptUnion);
      } else {
         HiveSelectQueryBlock queryBlock = new HiveSelectQueryBlock();
         queryBlock.setParent(parent);
         if (this.lexer.hasComment() && this.lexer.isKeepComments()) {
            queryBlock.addBeforeComment(this.lexer.readAndResetComments());
         }

         this.accept(Token.SELECT);
         if (this.lexer.token() == Token.HINT) {
            this.exprParser.parseHints(queryBlock.getHints());
         }

         if (DbType.informix == this.dbType) {
            if (this.lexer.identifierEquals(FnvHash.Constants.SKIP)) {
               this.lexer.nextToken();
               SQLExpr offset = this.exprParser.primary();
               queryBlock.setOffset(offset);
            }

            if (this.lexer.identifierEquals(FnvHash.Constants.FIRST)) {
               this.lexer.nextToken();
               SQLExpr first = this.exprParser.primary();
               queryBlock.setFirst(first);
            }
         }

         if (this.lexer.token() == Token.DISTINCT) {
            queryBlock.setDistionOption(2);
            this.lexer.nextToken();
         } else if (this.lexer.token() == Token.UNIQUE) {
            queryBlock.setDistionOption(3);
            this.lexer.nextToken();
         } else if (this.lexer.token() == Token.ALL) {
            queryBlock.setDistionOption(1);
            this.lexer.nextToken();
         }

         this.parseSelectList(queryBlock);
         if (this.lexer.token() == Token.INTO) {
            this.lexer.nextToken();
            SQLExpr expr = this.expr();
            if (this.lexer.token() != Token.COMMA) {
               queryBlock.setInto(expr);
            }
         }

         this.parseFrom(queryBlock);
         this.parseWhere(queryBlock);
         this.parseGroupBy(queryBlock);
         if (this.lexer.identifierEquals(FnvHash.Constants.WINDOW)) {
            this.parseWindow(queryBlock);
         }

         this.parseSortBy(queryBlock);
         this.parseFetchClause(queryBlock);
         if (this.lexer.token() == Token.FOR) {
            this.lexer.nextToken();
            this.accept(Token.UPDATE);
            queryBlock.setForUpdate(true);
            if (!this.lexer.identifierEquals(FnvHash.Constants.NO_WAIT) && !this.lexer.identifierEquals(FnvHash.Constants.NOWAIT)) {
               if (this.lexer.identifierEquals(FnvHash.Constants.WAIT)) {
                  this.lexer.nextToken();
                  SQLExpr waitTime = this.exprParser.primary();
                  queryBlock.setWaitTime(waitTime);
               }
            } else {
               this.lexer.nextToken();
               queryBlock.setNoWait(true);
            }
         }

         return this.queryRest(queryBlock, acceptUnion);
      }
   }
}
