#pragma once

namespace DB
{

/// Fills the array_join_result_to_source: on which columns-arrays to replicate, and how to call them after that.
class ArrayJoinedColumnsVisitor
{
public:
    ArrayJoinedColumnsVisitor(NameToNameMap & array_join_name_to_alias_,
                              NameToNameMap & array_join_alias_to_name_,
                              NameToNameMap & array_join_result_to_source_)
    :   array_join_name_to_alias(array_join_name_to_alias_),
        array_join_alias_to_name(array_join_alias_to_name_),
        array_join_result_to_source(array_join_result_to_source_)
    {}

    void visit(ASTPtr & ast) const
    {
        if (!tryVisit<ASTTablesInSelectQuery>(ast) &&
            !tryVisit<ASTIdentifier>(ast))
            visitChildren(ast);
    }

private:
    NameToNameMap & array_join_name_to_alias;
    NameToNameMap & array_join_alias_to_name;
    NameToNameMap & array_join_result_to_source;

    void visit(const ASTTablesInSelectQuery *, ASTPtr &) const
    {}

    void visit(const ASTIdentifier * node, ASTPtr &) const
    {
        if (node->general())
        {
            auto splitted = Nested::splitName(node->name);  /// ParsedParams, Key1

            if (array_join_alias_to_name.count(node->name))
            {
                /// ARRAY JOIN was written with an array column. Example: SELECT K1 FROM ... ARRAY JOIN ParsedParams.Key1 AS K1
                array_join_result_to_source[node->name] = array_join_alias_to_name[node->name];    /// K1 -> ParsedParams.Key1
            }
            else if (array_join_alias_to_name.count(splitted.first) && !splitted.second.empty())
            {
                /// ARRAY JOIN was written with a nested table. Example: SELECT PP.KEY1 FROM ... ARRAY JOIN ParsedParams AS PP
                array_join_result_to_source[node->name]    /// PP.Key1 -> ParsedParams.Key1
                    = Nested::concatenateName(array_join_alias_to_name[splitted.first], splitted.second);
            }
            else if (array_join_name_to_alias.count(node->name))
            {
                /** Example: SELECT ParsedParams.Key1 FROM ... ARRAY JOIN ParsedParams.Key1 AS PP.Key1.
                * That is, the query uses the original array, replicated by itself.
                */
                array_join_result_to_source[    /// PP.Key1 -> ParsedParams.Key1
                    array_join_name_to_alias[node->name]] = node->name;
            }
            else if (array_join_name_to_alias.count(splitted.first) && !splitted.second.empty())
            {
                /** Example: SELECT ParsedParams.Key1 FROM ... ARRAY JOIN ParsedParams AS PP.
                */
                array_join_result_to_source[    /// PP.Key1 -> ParsedParams.Key1
                Nested::concatenateName(array_join_name_to_alias[splitted.first], splitted.second)] = node->name;
            }
        }
    }

    void visit(const ASTSubquery *, ASTPtr &) const
    {}

    void visit(const ASTSelectQuery *, ASTPtr &) const
    {}

    void visitChildren(ASTPtr & ast) const
    {
        for (auto & child : ast->children)
            if (!tryVisit<ASTSubquery>(child) &&
                !tryVisit<ASTSelectQuery>(child))
                visit(child);
    }

    template <typename T>
    bool tryVisit(ASTPtr & ast) const
    {
        if (const T * t = typeid_cast<const T *>(ast.get()))
        {
            visit(t, ast);
            return true;
        }
        return false;
    }
};

}