/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.shardingsphere.sql.parser.postgresql.visitor.statement.type;

import org.antlr.v4.runtime.misc.Interval;
import org.apache.shardingsphere.sql.parser.api.ASTNode;
import org.apache.shardingsphere.sql.parser.api.visitor.statement.type.DMLStatementVisitor;
import org.apache.shardingsphere.sql.parser.autogen.PostgreSQLStatementParser.CallArgumentContext;
import org.apache.shardingsphere.sql.parser.autogen.PostgreSQLStatementParser.CallContext;
import org.apache.shardingsphere.sql.parser.autogen.PostgreSQLStatementParser.CopyContext;
import org.apache.shardingsphere.sql.parser.autogen.PostgreSQLStatementParser.CopyWithTableBinaryContext;
import org.apache.shardingsphere.sql.parser.autogen.PostgreSQLStatementParser.CopyWithTableOrQueryBinaryCsvContext;
import org.apache.shardingsphere.sql.parser.autogen.PostgreSQLStatementParser.CopyWithTableOrQueryContext;
import org.apache.shardingsphere.sql.parser.autogen.PostgreSQLStatementParser.DoStatementContext;
import org.apache.shardingsphere.sql.parser.autogen.PostgreSQLStatementParser.PreparableStmtContext;
import org.apache.shardingsphere.sql.parser.autogen.PostgreSQLStatementParser.ReturningClauseContext;
import org.apache.shardingsphere.sql.parser.postgresql.visitor.statement.PostgreSQLStatementVisitor;
import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.ReturningSegment;
import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.column.ColumnSegment;
import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.expr.ExpressionSegment;
import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.expr.complex.CommonExpressionSegment;
import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.item.ProjectionsSegment;
import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.prepare.PrepareStatementQuerySegment;
import org.apache.shardingsphere.sql.parser.statement.core.segment.generic.table.SimpleTableSegment;
import org.apache.shardingsphere.sql.parser.statement.core.statement.dml.DeleteStatement;
import org.apache.shardingsphere.sql.parser.statement.core.statement.dml.InsertStatement;
import org.apache.shardingsphere.sql.parser.statement.core.statement.dml.SelectStatement;
import org.apache.shardingsphere.sql.parser.statement.core.statement.dml.UpdateStatement;
import org.apache.shardingsphere.sql.parser.statement.core.value.collection.CollectionValue;
import org.apache.shardingsphere.sql.parser.statement.core.value.identifier.IdentifierValue;
import org.apache.shardingsphere.sql.parser.statement.postgresql.dml.PostgreSQLCallStatement;
import org.apache.shardingsphere.sql.parser.statement.postgresql.dml.PostgreSQLCopyStatement;
import org.apache.shardingsphere.sql.parser.statement.postgresql.dml.PostgreSQLDoStatement;

import java.util.Collection;
import java.util.LinkedList;

/**
 * DML statement visitor for PostgreSQL.
 */
public final class PostgreSQLDMLStatementVisitor extends PostgreSQLStatementVisitor implements DMLStatementVisitor {
    
    @Override
    public ASTNode visitCall(final CallContext ctx) {
        PostgreSQLCallStatement result = new PostgreSQLCallStatement();
        result.setProcedureName(((IdentifierValue) visit(ctx.identifier())).getValue());
        if (null != ctx.callArguments()) {
            Collection<ExpressionSegment> params = new LinkedList<>();
            for (CallArgumentContext each : ctx.callArguments().callArgument()) {
                params.add((ExpressionSegment) visit(each));
            }
            result.getParameters().addAll(params);
        }
        return result;
    }
    
    @Override
    public ASTNode visitCallArgument(final CallArgumentContext ctx) {
        if (null == ctx.positionalNotation()) {
            String text = ctx.namedNotation().start.getInputStream().getText(new Interval(ctx.namedNotation().start.getStartIndex(), ctx.namedNotation().stop.getStopIndex()));
            return new CommonExpressionSegment(ctx.namedNotation().getStart().getStartIndex(), ctx.namedNotation().getStop().getStopIndex(), text);
        }
        return visit(ctx.positionalNotation().aExpr());
    }
    
    @Override
    public ASTNode visitDoStatement(final DoStatementContext ctx) {
        return new PostgreSQLDoStatement();
    }
    
    @Override
    public ASTNode visitCopy(final CopyContext ctx) {
        if (null != ctx.copyWithTableOrQuery()) {
            return visit(ctx.copyWithTableOrQuery());
        }
        if (null != ctx.copyWithTableOrQueryBinaryCsv()) {
            return visit(ctx.copyWithTableOrQueryBinaryCsv());
        }
        return visit(ctx.copyWithTableBinary());
    }
    
    @Override
    public ASTNode visitCopyWithTableOrQuery(final CopyWithTableOrQueryContext ctx) {
        PostgreSQLCopyStatement result = new PostgreSQLCopyStatement();
        if (null != ctx.qualifiedName()) {
            result.setTableSegment((SimpleTableSegment) visit(ctx.qualifiedName()));
            if (null != ctx.columnNames()) {
                result.getColumns().addAll(((CollectionValue<ColumnSegment>) visit(ctx.columnNames())).getValue());
            }
        }
        if (null != ctx.preparableStmt()) {
            result.setPrepareStatementQuery(extractPrepareStatementQuerySegmentFromPreparableStmt(ctx.preparableStmt()));
        }
        return result;
    }
    
    private PrepareStatementQuerySegment extractPrepareStatementQuerySegmentFromPreparableStmt(final PreparableStmtContext ctx) {
        PrepareStatementQuerySegment result = new PrepareStatementQuerySegment(ctx.start.getStartIndex(), ctx.stop.getStopIndex());
        if (null != ctx.select()) {
            result.setSelect((SelectStatement) visit(ctx.select()));
        } else if (null != ctx.insert()) {
            result.setInsert((InsertStatement) visit(ctx.insert()));
        } else if (null != ctx.update()) {
            result.setUpdate((UpdateStatement) visit(ctx.update()));
        } else {
            result.setDelete((DeleteStatement) visit(ctx.delete()));
        }
        return result;
    }
    
    @Override
    public ASTNode visitCopyWithTableOrQueryBinaryCsv(final CopyWithTableOrQueryBinaryCsvContext ctx) {
        PostgreSQLCopyStatement result = new PostgreSQLCopyStatement();
        if (null != ctx.qualifiedName()) {
            result.setTableSegment((SimpleTableSegment) visit(ctx.qualifiedName()));
            if (null != ctx.columnNames()) {
                result.getColumns().addAll(((CollectionValue<ColumnSegment>) visit(ctx.columnNames())).getValue());
            }
        }
        if (null != ctx.preparableStmt()) {
            result.setPrepareStatementQuery(extractPrepareStatementQuerySegmentFromPreparableStmt(ctx.preparableStmt()));
        }
        return result;
    }
    
    @Override
    public ASTNode visitCopyWithTableBinary(final CopyWithTableBinaryContext ctx) {
        PostgreSQLCopyStatement result = new PostgreSQLCopyStatement();
        if (null != ctx.qualifiedName()) {
            result.setTableSegment((SimpleTableSegment) visit(ctx.qualifiedName()));
        }
        return result;
    }
    
    @Override
    public ASTNode visitReturningClause(final ReturningClauseContext ctx) {
        return new ReturningSegment(ctx.start.getStartIndex(), ctx.stop.getStopIndex(), (ProjectionsSegment) visit(ctx.targetList()));
    }
}
