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

import com.alibaba.druid.DbType;
import com.alibaba.druid.sql.ast.SQLCommentHint;
import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLName;
import com.alibaba.druid.sql.ast.SQLObject;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.SQLStatementImpl;
import com.alibaba.druid.sql.ast.expr.SQLNullExpr;
import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr;
import com.alibaba.druid.sql.ast.expr.SQLVariantRefExpr;
import com.alibaba.druid.sql.ast.statement.SQLBlockStatement;
import com.alibaba.druid.sql.ast.statement.SQLCallStatement;
import com.alibaba.druid.sql.ast.statement.SQLExplainStatement;
import com.alibaba.druid.sql.ast.statement.SQLExportTableStatement;
import com.alibaba.druid.sql.ast.statement.SQLExprStatement;
import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
import com.alibaba.druid.sql.ast.statement.SQLImportTableStatement;
import com.alibaba.druid.sql.ast.statement.SQLSelect;
import com.alibaba.druid.sql.ast.statement.SQLSelectStatement;
import com.alibaba.druid.sql.ast.statement.SQLValuesTableSource;
import com.alibaba.druid.sql.dialect.dm.ast.stmt.DmContinueStatement;
import com.alibaba.druid.sql.dialect.dm.ast.stmt.DmGotoStatement;
import com.alibaba.druid.sql.dialect.dm.ast.stmt.DmLabelStatement;
import com.alibaba.druid.sql.dialect.dm.ast.stmt.DmPrintStatement;
import com.alibaba.druid.sql.dialect.hive.ast.stmt.HiveLoadDataStatement;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlExplainPlanCacheStatement;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlHintStatement;
import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlSelectQueryBlock;
import com.alibaba.druid.sql.parser.Lexer;
import com.alibaba.druid.sql.parser.ParserException;
import com.alibaba.druid.sql.parser.SQLCreateTableParser;
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.List;

public class HiveStatementParser extends SQLStatementParser {
    public HiveStatementParser(String sql) {
        super((SQLExprParser) (new HiveExprParser(sql)));
        this.dbType = DbType.hive;
    }

    public HiveStatementParser(String sql, SQLParserFeature... features) {
        super((SQLExprParser) (new HiveExprParser(sql, features)));
        this.dbType = DbType.hive;
    }

    public HiveStatementParser(Lexer lexer) {
        super((SQLExprParser) (new HiveExprParser(lexer)));
        this.dbType = DbType.hive;
    }

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

    public SQLCreateTableParser getSQLCreateTableParser() {
        return new HiveCreateTableParser(this.exprParser);
    }

    public boolean parseStatementListDialect(List<SQLStatement> statementList) {
        return false;
    }

    protected HiveLoadDataStatement parseLoad() {
        this.accept(Token.LOAD);
        this.acceptIdentifier("DATA");
        HiveLoadDataStatement stmt = new HiveLoadDataStatement();
        if (this.lexer.token() == Token.LOCAL) {
            this.lexer.nextToken();
            stmt.setLocal(true);
        }

        this.acceptIdentifier("INPATH");
        SQLExpr inpath = this.exprParser.expr();
        stmt.setInpath(inpath);
        if (this.lexer.token() == Token.OVERWRITE) {
            this.lexer.nextToken();
            stmt.setOverwrite(true);
        }

        this.accept(Token.INTO);
        this.accept(Token.TABLE);
        SQLExpr table = this.exprParser.expr();
        stmt.setInto(table);
        if (this.lexer.token() == Token.PARTITION) {
            this.lexer.nextToken();
            this.accept(Token.LPAREN);
            this.exprParser.exprList(stmt.getPartition(), stmt);
            this.accept(Token.RPAREN);
        }

        return stmt;
    }

    protected SQLStatement parseExport() {
        this.lexer.nextToken();
        this.accept(Token.TABLE);
        SQLExportTableStatement stmt = new SQLExportTableStatement();
        stmt.setTable(new SQLExprTableSource(this.exprParser.name()));
        if (this.lexer.token() == Token.PARTITION) {
            this.lexer.nextToken();
            this.accept(Token.LPAREN);
            this.parseAssignItems(stmt.getPartition(), stmt, false);
            this.accept(Token.RPAREN);
        }

        if (this.lexer.token() == Token.TO) {
            this.lexer.nextToken();
            SQLExpr to = this.exprParser.primary();
            stmt.setTo(to);
        }

        return stmt;
    }

    protected SQLStatement parseImport() {
        this.lexer.nextToken();
        SQLImportTableStatement stmt = new SQLImportTableStatement();
        if (this.lexer.identifierEquals(FnvHash.Constants.EXTERNAL)) {
            this.lexer.nextToken();
            stmt.setExtenal(true);
        }

        if (this.lexer.token() == Token.TABLE) {
            this.lexer.nextToken();
            stmt.setTable(new SQLExprTableSource(this.exprParser.name()));
        }

        if (this.lexer.token() == Token.PARTITION) {
            this.lexer.nextToken();
            this.accept(Token.LPAREN);
            this.parseAssignItems(stmt.getPartition(), stmt, false);
            this.accept(Token.RPAREN);
        }

        if (this.lexer.token() == Token.FROM) {
            this.lexer.nextToken();
            SQLExpr to = this.exprParser.primary();
            stmt.setFrom(to);
        }

        return stmt;
    }

    public HiveExprParser getExprParser() {
        return (HiveExprParser) this.exprParser;
    }

    public void parseStatementList(List<SQLStatement> statementList, int max, SQLObject parent) {
        if ("select @@session.tx_read_only".equals(this.lexer.text) && this.lexer.token() == Token.SELECT) {
            SQLSelect select = new SQLSelect();
            MySqlSelectQueryBlock queryBlock = new MySqlSelectQueryBlock();
            queryBlock.addSelectItem(new SQLPropertyExpr(new SQLVariantRefExpr("@@session"), "tx_read_only"));
            select.setQuery(queryBlock);
            SQLSelectStatement stmt = new SQLSelectStatement(select);
            statementList.add(stmt);
            this.lexer.reset(29, '\u001a', Token.EOF);
        } else {
            boolean semi = false;

            for (int i = 0; max == -1 || statementList.size() < max; ++i) {
                while (this.lexer.token() == Token.MULTI_LINE_COMMENT || this.lexer.token() == Token.LINE_COMMENT) {
                    this.lexer.nextToken();
                }

                switch (this.lexer.token()) {
                    case EOF:
                    case END:
                    case UNTIL:
                    case ELSE:
                    case WHEN: {
                        if (this.lexer.isKeepComments() && this.lexer.hasComment() && statementList.size() > 0) {
                            SQLStatement stmt = (SQLStatement) statementList.get(statementList.size() - 1);
                            stmt.addAfterComment(this.lexer.readAndResetComments());
                        }

                        return;
                    }
                    case SEMI: {
                        int line0 = this.lexer.getLine();
                        this.lexer.nextToken();
                        int line1 = this.lexer.getLine();
                        if (statementList.size() > 0) {
                            SQLStatement lastStmt = (SQLStatement) statementList.get(statementList.size() - 1);
                            lastStmt.setAfterSemi(true);
                            if (this.lexer.isKeepComments()) {
                                SQLStatement stmt = (SQLStatement) statementList.get(statementList.size() - 1);
                                if (line1 - line0 <= 1) {
                                    stmt.addAfterComment(this.lexer.readAndResetComments());
                                }
                            }
                        }

                        semi = true;
                        break;
                    }
                    case WITH: {
                        SQLStatement stmt = this.parseWith();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case SELECT: {
                        MySqlHintStatement hintStatement = null;
                        SQLStatement stmt = this.parseSelect();
                        stmt.setParent(parent);
                        if (hintStatement != null && stmt instanceof SQLStatementImpl) {
                            SQLStatementImpl stmtImpl = (SQLStatementImpl) stmt;
                            List<SQLCommentHint> hints = stmtImpl.getHeadHintsDirect();
                            if (hints == null) {
                                stmtImpl.setHeadHints(hintStatement.getHints());
                            } else {
                                hints.addAll(hintStatement.getHints());
                            }

                            statementList.set(statementList.size() - 1, stmt);
                        } else {
                            statementList.add(stmt);
                        }

                        semi = false;
                        break;
                    }
                    case LOAD: {
                        SQLStatement stmt = this.parseLoad();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case UPDATE: {
                        Lexer.SavePoint savePoint = this.lexer.mark();
                        this.lexer.nextToken();
                        this.lexer.reset(savePoint);
                        SQLStatement stmt = this.parseUpdateStatement();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case CREATE: {
                        HiveCreateParser parser = new HiveCreateParser(this.exprParser);
                        SQLStatement stmt = parser.parseCreate();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case INSERT: {
                        HiveInsertParser parser = new HiveInsertParser(this.exprParser);
                        SQLStatement stmt = parser.parseInsert();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case DELETE: {
                        SQLStatement stmt = this.parseDeleteStatement();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case EXPLAIN: {
                        this.lexer.computeRowAndColumn();
                        int sourceLine = this.lexer.getPosLine();
                        int sourceColumn = this.lexer.getPosColumn();
                        Lexer.SavePoint savePoint = this.lexer.mark();
                        this.lexer.nextToken();
                        if (this.lexer.identifierEquals("PLANCACHE")) {
                            this.lexer.nextToken();
                            MySqlExplainPlanCacheStatement stmt = new MySqlExplainPlanCacheStatement();
                            stmt.setSourceLine(sourceLine);
                            stmt.setSourceLine(sourceColumn);
                            statementList.add(stmt);
                        } else {
                            this.lexer.reset(savePoint);
                            SQLExplainStatement stmt = this.parseExplain();
                            stmt.setSourceLine(sourceLine);
                            stmt.setSourceLine(sourceColumn);
                            stmt.setParent(parent);
                            statementList.add(stmt);
                        }
                        break;
                    }
                    case SET: {
                        SQLStatement stmt = this.parseSet();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case ALTER: {
                        HiveAlterStatementParser parser = new HiveAlterStatementParser(this.exprParser);
                        SQLStatement stmt = parser.parseAlter();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case TRUNCATE: {
                        SQLStatement stmt = this.parseTruncate();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case USE: {
                        SQLStatement stmt = this.parseUse();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case GRANT: {
                        HiveGrantParser parser = new HiveGrantParser(this.exprParser);
                        SQLStatement stmt = parser.parseGrant();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case REVOKE: {
                        HiveGrantParser parser = new HiveGrantParser(this.exprParser);
                        SQLStatement stmt = parser.parseRevoke();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case SHOW: {
                        HiveShowParser parser = new HiveShowParser(this.exprParser);
                        SQLStatement stmt = parser.parseShow();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case ABORT: {
                        HiveShowParser parser = new HiveShowParser(this.exprParser);
                        SQLStatement stmt = parser.parseAbort();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case MERGE: {
                        SQLStatement stmt = this.parseMerge();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case REPEAT: {
                        SQLStatement stmt = this.parseRepeat();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case WHILE: {
                        SQLStatement stmt = this.parseWhile();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case IF: {
                        SQLStatement stmt = this.parseIf();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case CASE: {
                        SQLStatement stmt = this.parseCase();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case OPEN: {
                        SQLStatement stmt = this.parseOpen();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case NULL: {
                        this.lexer.nextToken();
                        SQLExprStatement stmt = new SQLExprStatement(new SQLNullExpr());
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case FETCH: {
                        SQLStatement stmt = this.parseFetch();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case DROP: {
                        SQLStatement stmt = this.parseDrop();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case COMMENT: {
                        SQLStatement stmt = this.parseComment();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case KILL: {
                        SQLStatement stmt = this.parseKill();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case CLOSE: {
                        SQLStatement stmt = this.parseClose();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case RETURN: {
                        SQLStatement stmt = this.parseReturn();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case UPSERT: {
                        SQLStatement stmt = this.parseUpsert();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    case LEAVE: {
                        SQLStatement stmt = this.parseLeave();
                        stmt.setParent(parent);
                        statementList.add(stmt);
                        break;
                    }
                    default:
                        if (this.lexer.token() == Token.CASE) {
                            SQLStatement stmt = this.parseCase();
                            stmt.setParent(parent);
                            statementList.add(stmt);
                        } else if (this.lexer.token() != Token.LBRACE && !this.lexer.identifierEquals("CALL")) {
                            if (this.lexer.identifierEquals("UPSERT")) {
                                SQLStatement stmt = this.parseUpsert();
                                statementList.add(stmt);
                            } else {
                                if (this.lexer.identifierEquals("LIST")) {
                                    Lexer.SavePoint mark = this.lexer.mark();
                                    SQLStatement stmt = this.parseList();
                                    if (stmt != null) {
                                        statementList.add(stmt);
                                        continue;
                                    }

                                    this.lexer.reset(mark);
                                }

                                if (this.lexer.identifierEquals("RENAME")) {
                                    SQLStatement stmt = this.parseRename();
                                    statementList.add(stmt);
                                } else if (this.lexer.token() == Token.CONTINUE) {
                                    this.lexer.nextToken();
                                    DmContinueStatement stmt = new DmContinueStatement();
                                    if (this.lexer.token() == Token.IDENTIFIER) {
                                        String label = this.lexer.stringVal();
                                        this.lexer.nextToken();
                                        stmt.setLabel(label);
                                    }

                                    if (this.lexer.token() == Token.WHEN) {
                                        this.lexer.nextToken();
                                        stmt.setWhen(this.exprParser.expr());
                                    }

                                    stmt.setParent(parent);
                                    statementList.add(stmt);
                                } else if (this.lexer.token() == Token.GOTO) {
                                    this.lexer.nextToken();
                                    SQLName label = this.exprParser.name();
                                    DmGotoStatement stmt = new DmGotoStatement(label);
                                    stmt.setParent(parent);
                                    statementList.add(stmt);
                                } else if (this.lexer.token() == Token.LTLT) {
                                    this.lexer.nextToken();
                                    SQLName label = this.exprParser.name();
                                    DmLabelStatement stmt = new DmLabelStatement(label);
                                    this.accept(Token.GTGT);
                                    stmt.setParent(parent);
                                    statementList.add(stmt);
                                } else if (this.lexer.identifierEquals("RELEASE")) {
                                    SQLStatement stmt = this.parseReleaseSavePoint();
                                    statementList.add(stmt);
                                } else if (this.lexer.token() != Token.BEGIN && this.lexer.token() != Token.DECLARE) {
                                    if (this.lexer.token() == Token.SAVEPOINT) {
                                        SQLStatement stmt = this.parseSavePoint();
                                        statementList.add(stmt);
                                    } else if (this.lexer.identifierEquals("REFRESH")) {
                                        SQLStatement stmt = this.parseRefresh();
                                        statementList.add(stmt);
                                    } else if (this.lexer.identifierEquals(FnvHash.Constants.COPY)) {
                                        SQLStatement stmt = this.parseCopy();
                                        statementList.add(stmt);
                                    } else if (this.lexer.token() != Token.DESC && !this.lexer.identifierEquals(FnvHash.Constants.DESCRIBE)) {
                                        if (this.lexer.identifierEquals("ROLLBACK")) {
                                            SQLStatement stmt = this.parseRollback();
                                            statementList.add(stmt);
                                            if (parent instanceof SQLBlockStatement && DbType.mysql == this.dbType) {
                                                return;
                                            }
                                        } else if (this.lexer.identifierEquals("DUMP")) {
                                            SQLStatement stmt = this.parseDump();
                                            statementList.add(stmt);
                                        } else if (this.lexer.token() == Token.COMMIT) {
                                            SQLStatement stmt = this.parseCommit();
                                            statementList.add(stmt);
                                            if (parent instanceof SQLBlockStatement && DbType.mysql == this.dbType) {
                                                return;
                                            }
                                        } else if (this.lexer.identifierEquals(FnvHash.Constants.RETURN)) {
                                            SQLStatement stmt = this.parseReturn();
                                            statementList.add(stmt);
                                        } else if (this.lexer.identifierEquals(FnvHash.Constants.PURGE)) {
                                            SQLStatement stmt = this.parsePurge();
                                            statementList.add(stmt);
                                        } else if (this.lexer.identifierEquals(FnvHash.Constants.FLASHBACK)) {
                                            SQLStatement stmt = this.parseFlashback();
                                            statementList.add(stmt);
                                        } else if (this.lexer.identifierEquals(FnvHash.Constants.WHO)) {
                                            SQLStatement stmt = this.parseWhoami();
                                            statementList.add(stmt);
                                        } else if (this.lexer.token() == Token.FOR) {
                                            SQLStatement stmt = this.parseFor();
                                            statementList.add(stmt);
                                            stmt.setParent(parent);
                                        } else if (this.lexer.token() == Token.LPAREN) {
                                            char markChar = this.lexer.current();
                                            int markBp = this.lexer.bp();

                                            do {
                                                this.lexer.nextToken();
                                            } while (this.lexer.token() == Token.LPAREN);

                                            if (this.lexer.token() != Token.SELECT) {
                                                throw new ParserException("TODO " + this.lexer.info());
                                            }

                                            this.lexer.reset(markBp, markChar, Token.LPAREN);
                                            SQLStatement stmt = this.parseSelect();
                                            statementList.add(stmt);
                                        } else if (this.lexer.token() == Token.VALUES) {
                                            SQLValuesTableSource values = this.createSQLSelectParser().parseValues();
                                            SQLSelectStatement stmt = new SQLSelectStatement();
                                            stmt.setSelect(new SQLSelect(values));
                                            statementList.add(stmt);
                                            stmt.setParent(parent);
                                        } else if (this.lexer.identifierEquals("PRINT")) {
                                            this.lexer.nextToken();
                                            DmPrintStatement stmt = new DmPrintStatement();
                                            stmt.setString(this.exprParser.expr());
                                            stmt.setParent(parent);
                                            statementList.add(stmt);
                                        } else {
                                            if (this.lexer.token() == Token.IDENTIFIER || this.lexer.token() == Token.VARIANT) {
                                                Lexer.SavePoint savePoint = this.lexer.mark();
                                                if (this.lexer.token() == Token.VARIANT) {
                                                    this.exprParser.name();
                                                } else {
                                                    this.lexer.nextToken();
                                                    if (this.lexer.token() == Token.DOT) {
                                                        this.lexer.nextToken();
                                                        this.lexer.nextToken();
                                                    }
                                                }

                                                if (this.lexer.token() == Token.COLONEQ || this.lexer.token() == Token.EQ) {
                                                    this.lexer.reset(savePoint);
                                                    SQLStatement stmt = this.parseSet();
                                                    stmt.setParent(parent);
                                                    statementList.add(stmt);
                                                    continue;
                                                }

                                                if (this.lexer.token() == Token.LPAREN) {
                                                    this.lexer.reset(savePoint);
                                                    SQLCallStatement stmt = this.parseCall();
                                                    statementList.add(stmt);
                                                    continue;
                                                }

                                                this.lexer.reset(savePoint);
                                            }

                                            int size = statementList.size();
                                            if (this.parseStatementListDialect(statementList)) {
                                                if (parent != null) {
                                                    for (int j = size; j < statementList.size(); ++j) {
                                                        SQLStatement dialectStmt = (SQLStatement) statementList.get(j);
                                                        dialectStmt.setParent(parent);
                                                    }
                                                }
                                            } else {
                                                this.printError(this.lexer.token());
                                            }
                                        }
                                    } else {
                                        SQLStatement stmt = this.parseDescribe();
                                        statementList.add(stmt);
                                    }
                                } else {
                                    SQLStatement stmt = this.parseBlock();
                                    stmt.setParent(parent);
                                    statementList.add(stmt);
                                }
                            }
                        } else {
                            SQLCallStatement stmt = this.parseCall();
                            statementList.add(stmt);
                        }
                }
            }

        }
    }
}
