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

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.expr.SQLIdentifierExpr;
import com.alibaba.druid.sql.ast.statement.SQLAssignItem;
import com.alibaba.druid.sql.ast.statement.SQLExprTableSource;
import com.alibaba.druid.sql.ast.statement.SQLWithSubqueryClause;
import com.alibaba.druid.sql.dialect.hive.ast.HiveGrantPrinciple;
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.HiveAlterDatabaseStatement;
import com.alibaba.druid.sql.dialect.hive.ast.stmt.HiveCreateFunctionStatement;
import com.alibaba.druid.sql.dialect.hive.ast.stmt.HiveInsertStatement;
import com.alibaba.druid.sql.dialect.hive.ast.stmt.HiveLoadDataStatement;
import com.alibaba.druid.sql.dialect.hive.ast.stmt.HiveMsckRepairStatement;
import com.alibaba.druid.sql.dialect.hive.ast.stmt.HiveSelectQueryBlock;
import com.alibaba.druid.sql.dialect.hive.ast.stmt.HiveShowStatement;
import com.alibaba.druid.sql.repository.SchemaRepository;
import com.alibaba.druid.sql.visitor.SchemaStatVisitor;
import com.alibaba.druid.stat.TableStat;

public class HiveSchemaStatVisitor extends SchemaStatVisitor implements HiveASTVisitor {
   public HiveSchemaStatVisitor() {
      super(DbType.hive);
   }

   public HiveSchemaStatVisitor(DbType dbType) {
      super(dbType);
   }

   public HiveSchemaStatVisitor(SchemaRepository repository) {
      super(repository);
   }

   public boolean visit(HiveInsert x) {
      this.setMode(x, TableStat.Mode.Insert);
      SQLExprTableSource tableSource = x.getTableSource();
      SQLExpr tableName = tableSource.getExpr();
      if (tableName instanceof SQLName) {
         TableStat stat = this.getTableStat((SQLName)tableName);
         stat.incrementInsertCount();
      }

      for(SQLAssignItem partition : x.getPartitions()) {
         partition.accept(this);
      }

      this.accept(x.getQuery());
      return false;
   }

   public void endVisit(HiveInsert x) {
   }

   public boolean visit(HiveMultiInsertStatement x) {
      if (this.repository != null && x.getParent() == null) {
         this.repository.resolve((SQLStatement)x);
      }

      return true;
   }

   public void endVisit(HiveMultiInsertStatement x) {
   }

   public boolean visit(HiveInsertStatement x) {
      if (this.repository != null && x.getParent() == null) {
         this.repository.resolve((SQLStatement)x);
      }

      SQLWithSubqueryClause with = x.getWith();
      if (with != null) {
         with.accept(this);
      }

      this.setMode(x, TableStat.Mode.Insert);
      SQLExprTableSource tableSource = x.getTableSource();
      SQLExpr tableName = tableSource.getExpr();
      if (tableName instanceof SQLName) {
         TableStat stat = this.getTableStat((SQLName)tableName);
         stat.incrementInsertCount();

         for(SQLExpr column : x.getColumns()) {
            if (column instanceof SQLIdentifierExpr) {
               this.addColumn((SQLName)tableName, ((SQLIdentifierExpr)column).normalizedName());
            }
         }
      }

      for(SQLAssignItem partition : x.getPartitions()) {
         partition.accept(this);
      }

      this.accept(x.getQuery());
      return false;
   }

   public void endVisit(HiveInsertStatement x) {
   }

   public boolean visit(HiveCreateFunctionStatement x) {
      return false;
   }

   public void endVisit(HiveCreateFunctionStatement x) {
   }

   public boolean visit(HiveLoadDataStatement x) {
      TableStat tableStat = this.getTableStat(x.getInto());
      if (tableStat != null) {
         tableStat.incrementInsertCount();
      }

      return false;
   }

   public void endVisit(HiveLoadDataStatement x) {
   }

   public boolean visit(HiveMsckRepairStatement x) {
      return false;
   }

   public void endVisit(HiveMsckRepairStatement x) {
   }

   public boolean visit(HiveGrantPrinciple x) {
      return false;
   }

   public void endVisit(HiveGrantPrinciple x) {
   }

   public boolean visit(HiveShowStatement x) {
      return true;
   }

   public void endVisit(HiveShowStatement x) {
   }

   public boolean visit(HiveAlterDatabaseStatement x) {
      return true;
   }

   public void endVisit(HiveAlterDatabaseStatement x) {
   }

   public boolean visit(HiveSelectQueryBlock x) {
      return false;
   }

   public void endVisit(HiveSelectQueryBlock x) {
   }
}
