/*
 * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
 * or more contributor license agreements. Licensed under the Elastic License;
 * you may not use this file except in compliance with the Elastic License.
 */
package org.elasticsearch.xpack.sql.parser;

import org.elasticsearch.xpack.sql.expression.Expression;
import org.elasticsearch.xpack.sql.expression.Literal;
import org.elasticsearch.xpack.sql.expression.NamedExpression;
import org.elasticsearch.xpack.sql.expression.Order;
import org.elasticsearch.xpack.sql.expression.UnresolvedAlias;
import org.elasticsearch.xpack.sql.parser.SqlBaseParser.AliasedQueryContext;
import org.elasticsearch.xpack.sql.parser.SqlBaseParser.AliasedRelationContext;
import org.elasticsearch.xpack.sql.parser.SqlBaseParser.FromClauseContext;
import org.elasticsearch.xpack.sql.parser.SqlBaseParser.JoinCriteriaContext;
import org.elasticsearch.xpack.sql.parser.SqlBaseParser.JoinRelationContext;
import org.elasticsearch.xpack.sql.parser.SqlBaseParser.JoinTypeContext;
import org.elasticsearch.xpack.sql.parser.SqlBaseParser.NamedQueryContext;
import org.elasticsearch.xpack.sql.parser.SqlBaseParser.QueryContext;
import org.elasticsearch.xpack.sql.parser.SqlBaseParser.QueryNoWithContext;
import org.elasticsearch.xpack.sql.parser.SqlBaseParser.QuerySpecificationContext;
import org.elasticsearch.xpack.sql.parser.SqlBaseParser.RelationContext;
import org.elasticsearch.xpack.sql.parser.SqlBaseParser.SubqueryContext;
import org.elasticsearch.xpack.sql.parser.SqlBaseParser.TableNameContext;
import org.elasticsearch.xpack.sql.plan.TableIdentifier;
import org.elasticsearch.xpack.sql.plan.logical.Aggregate;
import org.elasticsearch.xpack.sql.plan.logical.Distinct;
import org.elasticsearch.xpack.sql.plan.logical.Filter;
import org.elasticsearch.xpack.sql.plan.logical.FromlessSelect;
import org.elasticsearch.xpack.sql.plan.logical.Join;
import org.elasticsearch.xpack.sql.plan.logical.Join.JoinType;
import org.elasticsearch.xpack.sql.plan.logical.Limit;
import org.elasticsearch.xpack.sql.plan.logical.LogicalPlan;
import org.elasticsearch.xpack.sql.plan.logical.OrderBy;
import org.elasticsearch.xpack.sql.plan.logical.Project;
import org.elasticsearch.xpack.sql.plan.logical.SubQueryAlias;
import org.elasticsearch.xpack.sql.plan.logical.UnresolvedRelation;
import org.elasticsearch.xpack.sql.plan.logical.With;
import org.elasticsearch.xpack.sql.type.DataTypes;

import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

import static java.util.Collections.emptyList;
import static java.util.stream.Collectors.toList;

abstract class LogicalPlanBuilder extends ExpressionBuilder {

    @Override
    public LogicalPlan visitQuery(QueryContext ctx) {
        LogicalPlan body = plan(ctx.queryNoWith());
        
        List<SubQueryAlias> namedQueries = visitList(ctx.namedQuery(), SubQueryAlias.class);

        // unwrap query (and validate while at it)
        Map<String, SubQueryAlias> cteRelations = new LinkedHashMap<>(namedQueries.size());
        for (SubQueryAlias namedQuery : namedQueries) {
            if (cteRelations.put(namedQuery.alias(), namedQuery) != null) {
                throw new ParsingException(namedQuery.location(), "Duplicate alias %s", namedQuery.alias());
            }
        }

        // return WITH
        return new With(source(ctx), body, cteRelations);
    }

    @Override
    public LogicalPlan visitNamedQuery(NamedQueryContext ctx) {
        return new SubQueryAlias(source(ctx), plan(ctx.queryNoWith()), ctx.name.getText());
    }

    public LogicalPlan visitQueryNoWith(QueryNoWithContext ctx) {
        LogicalPlan plan = plan(ctx.queryTerm());

        if (!ctx.orderBy().isEmpty()) {
            plan = new OrderBy(source(ctx), plan, visitList(ctx.orderBy(), Order.class));
        }
        
        if (ctx.limit != null && ctx.INTEGER_VALUE() != null) {
            plan = new Limit(source(ctx), new Literal(source(ctx), ctx.limit.getText(), DataTypes.INTEGER), plan);
        }
        
        return plan;
    }

    @Override
    public LogicalPlan visitQuerySpecification(QuerySpecificationContext ctx) {
        LogicalPlan query = (ctx.fromClause() != null)? plan(ctx.fromClause()) : new FromlessSelect(source(ctx));
                
        // add WHERE
        if (ctx.where != null) {
            query = new Filter(source(ctx), query, expression(ctx.where));
        }

        List<NamedExpression> selectTarget = emptyList();

        // SELECT a, b, c ...
        if (!ctx.selectItem().isEmpty()) {
            selectTarget = expressions(ctx.selectItem()).stream()
                    .map(e -> (e instanceof NamedExpression) ? (NamedExpression) e : new UnresolvedAlias(e))
                    .collect(toList());
        }

        // GROUP BY
        if (ctx.groupBy() != null) {
            List<Expression> groupBy = expressions(ctx.groupBy().groupingElement());
            query = new Aggregate(source(ctx.groupBy()), query, groupBy, selectTarget);
        }
        else if (!selectTarget.isEmpty()) {
            query = new Project(source(ctx.selectItem(0)), query, selectTarget);
        }

        // HAVING
        if (ctx.having != null) {
            query = new Filter(source(ctx.having), query, expression(ctx.having));
        }

        if (ctx.setQuantifier() != null && ctx.setQuantifier().DISTINCT() != null) {
            query = new Distinct(source(ctx.setQuantifier()), query);
        }
        return query;
    }

    @Override
    public LogicalPlan visitFromClause(FromClauseContext ctx) {
        // if there are multiple FROM clauses, convert each pair in a inner join
        List<LogicalPlan> plans = plans(ctx.relation());
        return plans.stream()
                .reduce((left, right) -> new Join(source(ctx), left, right, Join.JoinType.IMPLICIT, null))
                .get();
    }

    @Override
    public LogicalPlan visitRelation(RelationContext ctx) {
        // check if there are multiple join clauses. ANTLR produces a right nested tree with the left join clause
        // at the top. However the fields previously references might be used in the following clauses.
        // As such, swap/reverse the tree.
        
        LogicalPlan result = plan(ctx.relationPrimary());
        for (JoinRelationContext j : ctx.joinRelation()) {
            result = doJoin(result, j);
        }

        return result;
    }

    private Join doJoin(LogicalPlan left, JoinRelationContext ctx) {
        JoinTypeContext joinType = ctx.joinType();

        Join.JoinType type = JoinType.INNER;
        if (joinType != null) {
            if (joinType.FULL() != null) {
                type = JoinType.FULL;
            }
            if (joinType.LEFT() != null) {
                type = JoinType.LEFT;
            }
            if (joinType.RIGHT() != null) {
                type = JoinType.RIGHT;
            }
        }

        Expression condition = null;
        JoinCriteriaContext criteria = ctx.joinCriteria();
        if (criteria != null) {
            if (criteria.USING() != null) {
                throw new UnsupportedOperationException();
            }
            if (criteria.booleanExpression() != null) {
                condition = expression(criteria.booleanExpression());
            }
        }

        return new Join(source(ctx), left, plan(ctx.right), type, condition);
    }

    @Override
    public Object visitAliasedRelation(AliasedRelationContext ctx) {
        return new SubQueryAlias(source(ctx), plan(ctx.relation()), visitQualifiedName(ctx.qualifiedName()));
    }

    @Override
    public Object visitAliasedQuery(AliasedQueryContext ctx) {
        return new SubQueryAlias(source(ctx), plan(ctx.queryNoWith()), visitQualifiedName(ctx.qualifiedName()));
    }

    @Override
    public Object visitSubquery(SubqueryContext ctx) {
        return plan(ctx.queryNoWith());
    }

    @Override
    public LogicalPlan visitTableName(TableNameContext ctx) {
        String alias = visitQualifiedName(ctx.qualifiedName());
        TableIdentifier tableIdentifier = visitTableIdentifier(ctx.tableIdentifier());
        return new UnresolvedRelation(source(ctx), tableIdentifier, alias);
    }
}
@@ -3966,16 +3966,33 @@
 e(ctx), 
+Integer.parseInt(
 ctx.limi
@@ -4002,16 +4002,17 @@
 etText()
+)
 , DataTy

Improve parsing of Limit