package com.facebook.presto.sql.tree;

import com.facebook.presto.sql.compiler.Schema;
import com.google.common.base.Joiner;
import com.google.common.base.Strings;

import java.io.PrintStream;
import java.util.IdentityHashMap;

public class TreePrinter
{
    private static final String INDENT = "   ";

    private final IdentityHashMap<QualifiedNameReference, QualifiedName> resolvedNameReferences;
    private final IdentityHashMap<Node, Schema> outputSchemas;
    private final PrintStream out;

    public TreePrinter(IdentityHashMap<QualifiedNameReference, QualifiedName> resolvedNameReferences,
            IdentityHashMap<Node, Schema> outputSchemas,
            PrintStream out)
    {
        this.resolvedNameReferences = new IdentityHashMap<>(resolvedNameReferences);
        this.outputSchemas = new IdentityHashMap<>(outputSchemas);
        this.out = out;
    }

    public void print(Node root)
    {
        AstVisitor<Void, Integer> printer = new DefaultTraversalVisitor<Void, Integer>()
        {
            @Override
            protected Void visitNode(Node node, Integer indentLevel)
            {
                throw new UnsupportedOperationException("not yet implemented: " + node);
            }

            @Override
            protected Void visitQuery(Query node, Integer indentLevel)
            {
                print(indentLevel, "Query " + outputSchemas.get(node));

                indentLevel++;
                process(node.getSelect(), indentLevel);

                print(indentLevel, "From");
                for (Relation relation : node.getFrom()) {
                    process(relation, indentLevel + 1);
                }

                if (node.getWhere() != null) {
                    print(indentLevel, "Where");
                    process(node.getWhere(), indentLevel + 1);
                }

                if (!node.getGroupBy().isEmpty()) {
                    print(indentLevel, "GroupBy");
                    for (Expression expression : node.getGroupBy()) {
                        process(expression, indentLevel + 1);
                    }
                }

                if (node.getHaving() != null) {
                    print(indentLevel, "Having");
                    process(node.getHaving(), indentLevel + 1);
                }

                if (!node.getOrderBy().isEmpty()) {
                    print(indentLevel, "OrderBy");
                    for (SortItem sortItem : node.getOrderBy()) {
                        process(sortItem, indentLevel + 1);
                    }
                }

                if (node.getLimit() != null) {
                    print(indentLevel, "Limit: " + node.getLimit());
                }

                return null;
            }

            @Override
            protected Void visitSelect(Select node, Integer indentLevel)
            {
                String distinct = "";
                if (node.isDistinct()) {
                    distinct = "[DISTINCT]";
                }
                print(indentLevel, "Select" + distinct);

                super.visitSelect(node, indentLevel + 1); // visit children

                return null;
            }

            @Override
            protected Void visitAllColumns(AllColumns node, Integer indent)
            {
                if (node.getPrefix() == null) {
                    print(indent, "*");
                }
                else {
                    print(indent, node.getPrefix() + ".*");
                }

                return null;
            }

            @Override
            protected Void visitComparisonExpression(ComparisonExpression node, Integer indentLevel)
            {
                print(indentLevel, node.getType().toString());

                super.visitComparisonExpression(node, indentLevel + 1);

                return null;
            }

            @Override
            protected Void visitArithmeticExpression(ArithmeticExpression node, Integer indentLevel)
            {
                print(indentLevel, node.getType().toString());

                super.visitArithmeticExpression(node, indentLevel + 1);

                return null;
            }

            @Override
            protected Void visitLogicalBinaryExpression(LogicalBinaryExpression node, Integer indentLevel)
            {
                print(indentLevel, node.getType().toString());

                super.visitLogicalBinaryExpression(node, indentLevel + 1);

                return null;
            }

            @Override
            protected Void visitStringLiteral(StringLiteral node, Integer indentLevel)
            {
                print(indentLevel, "String[" + node.getValue() + "]");
                return null;
            }

            @Override
            protected Void visitLongLiteral(LongLiteral node, Integer indentLevel)
            {
                print(indentLevel, "Long[" + node.getValue() + "]");
                return null;
            }

            @Override
            protected Void visitLikePredicate(LikePredicate node, Integer indentLevel)
            {
                print(indentLevel, "LIKE");

                super.visitLikePredicate(node, indentLevel + 1);

                return null;
            }

            @Override
            protected Void visitQualifiedNameReference(QualifiedNameReference node, Integer indentLevel)
            {
                QualifiedName resolved = resolvedNameReferences.get(node);
                String resolvedName = "";
                if (resolved != null) {
                    resolvedName = "=>" + resolved.toString();
                }
                print(indentLevel, "QualifiedName[" + node.getName() + resolvedName + "]");
                return null;
            }

            @Override
            protected Void visitFunctionCall(FunctionCall node, Integer indentLevel)
            {
                String name = Joiner.on('.').join(node.getName().getParts());
                print(indentLevel, "FunctionCall[" + name + "]");

                super.visitFunctionCall(node, indentLevel + 1);

                return null;
            }

            @Override
            protected Void visitTable(Table node, Integer indentLevel)
            {
                String name = Joiner.on('.').join(node.getName().getParts());
                print(indentLevel, "Table[" + name + "] " + outputSchemas.get(node));

                return null;
            }

            @Override
            protected Void visitAliasedRelation(AliasedRelation node, Integer indentLevel)
            {
                print(indentLevel, "Alias[" + node.getAlias() + "] " + outputSchemas.get(node));

                super.visitAliasedRelation(node, indentLevel + 1);

                return null;
            }

            @Override
            protected Void visitAliasedExpression(AliasedExpression node, Integer indentLevel)
            {
                print(indentLevel, "Alias[" + node.getAlias() + "]");

                super.visitAliasedExpression(node, indentLevel + 1);

                return null;
            }

            @Override
            protected Void visitSubquery(Subquery node, Integer indentLevel)
            {
                print(indentLevel, "SubQuery " + outputSchemas.get(node));

                super.visitSubquery(node, indentLevel + 1);

                return null;
            }

            @Override
            protected Void visitInPredicate(InPredicate node, Integer indentLevel)
            {
                print(indentLevel, "IN");

                super.visitInPredicate(node, indentLevel + 1);

                return null;
            }

            @Override
            protected Void visitSubqueryExpression(SubqueryExpression node, Integer indentLevel)
            {
                print(indentLevel, "SubQuery " + outputSchemas.get(node));

                super.visitSubqueryExpression(node, indentLevel + 1);

                return null;
            }
        };

        printer.process(root, 0);
    }

    private void print(Integer indentLevel, String value)
    {
        out.println(Strings.repeat(INDENT, indentLevel) + value);
    }
}

@@ -3329,16 +3329,20 @@
 ix()
- == null
+.isPresent()
 ) %7B%0A
@@ -3375,17 +3375,37 @@
 indent, 
-%22
+node.getPrefix() + %22.
 *%22);%0A   
@@ -3476,37 +3476,17 @@
 indent, 
-node.getPrefix() + %22.
+%22
 *%22);%0A   

Fix issue where AllColumns with optional prefix was not being printed properly