#include "global.h"

#include "column.h"
#include "function.h"
#include "input.h"
#include "named_node_resolution.h"
#include "named_node_visibility.h"
#include "parse_tree.h"
#include "parser.h"
#include "use.h"

#include <yql/essentials/sql/v1/ide/completion/syntax/ansi.h>

#include <library/cpp/iterator/functools.h>

#include <util/string/join.h>

namespace NSQLComplete {

namespace {

template <class C>
void Move(TColumnContext& lhs, TColumnContext& rhs, C TColumnContext::*member) {
    C& lhsM = lhs.*member;
    C& rhsM = rhs.*member;

    lhsM.reserve(lhsM.size() + rhsM.size());
    std::move(rhsM.begin(), rhsM.end(), std::back_inserter(lhsM));
    SortUnique(lhsM);
}

} // namespace

bool operator<(const TColumnId& lhs, const TColumnId& rhs) {
    return std::tie(lhs.TableAlias, lhs.Name) < std::tie(rhs.TableAlias, rhs.Name);
}

TColumnContext TColumnContext::ExtractAliased(TMaybe<TStringBuf> alias) {
    if (alias.Empty()) {
        return *this;
    }

    auto aliasedTables = std::ranges::partition(Tables, [&](const auto& table) {
        return table.Alias != alias;
    });

    auto aliasedColumns = std::ranges::partition(Columns, [&](const auto& column) {
        return column.TableAlias != alias;
    });

    TVector<TAliased<TTableId>> tables(aliasedTables.begin(), aliasedTables.end());
    TVector<TColumnId> columns(aliasedColumns.begin(), aliasedColumns.end());

    Tables.erase(aliasedTables.begin(), aliasedTables.end());
    Columns.erase(aliasedColumns.begin(), aliasedColumns.end());

    THashMap<TString, THashSet<TString>> without;
    if (auto it = WithoutByTableAlias.find(*alias); it != WithoutByTableAlias.end()) {
        without[*alias] = std::move(it->second);
        WithoutByTableAlias.erase(it);
    }

    return {
        .Tables = std::move(tables),
        .Columns = std::move(columns),
        .WithoutByTableAlias = std::move(without),
    };
}

bool TColumnContext::IsAsterisk() const {
    return Columns.size() == 1 &&
           Columns[0].TableAlias.empty() &&
           Columns[0].Name == "*";
}

TColumnContext TColumnContext::Renamed(TStringBuf alias) && {
    for (TAliased<TTableId>& table : Tables) {
        table.Alias = alias;
    }

    for (TColumnId& column : Columns) {
        column.TableAlias = alias;
    }

    THashSet<TString>& without = WithoutByTableAlias[alias];
    for (auto& [tableAlias, excluded] : WithoutByTableAlias) {
        if (tableAlias == alias) {
            continue;
        }

        without.insert(excluded.begin(), excluded.end());
    }

    if (without.empty()) {
        WithoutByTableAlias = {};
    } else {
        WithoutByTableAlias = {{TString(alias), std::move(without)}};
    }

    return *this;
}

TColumnContext operator|(TColumnContext lhs, TColumnContext rhs) {
    Move(lhs, rhs, &TColumnContext::Tables);

    Move(lhs, rhs, &TColumnContext::Columns);

    for (auto& [tableAlias, excluded] : rhs.WithoutByTableAlias) {
        auto& without = lhs.WithoutByTableAlias[tableAlias];
        without.insert(excluded.begin(), excluded.end());
    }

    return lhs;
}

TColumnContext TColumnContext::Asterisk() {
    return {.Columns = {{.Name = "*"}}};
}

class TSpecializedGlobalAnalysis final: public IGlobalAnalysis {
public:
    explicit TSpecializedGlobalAnalysis(IParser::TPtr parser)
        : Parser_(std::move(parser))
    {
    }

    TGlobalContext Analyze(TCompletionInput input, TEnvironment env) override {
        TParsedInput parsed = Parser_->Parse(std::move(input));

        INamedNodes::TPtr nodes = ResolveNamedNodes(parsed, env);

        TGlobalContext ctx;
        ctx.Use = FindUseStatement(parsed, *nodes);
        ctx.Names = VisibleNamedNodes(parsed);
        ctx.EnclosingFunction = EnclosingFunction(parsed, *nodes);
        ctx.Column = InferColumnContext(parsed, *nodes);

        if (ctx.Use && ctx.Column) {
            EnrichTableClusters(*ctx.Column, *ctx.Use);
        }

        return ctx;
    }

private:
    void EnrichTableClusters(TColumnContext& column, const TClusterContext& use) {
        for (auto& table : column.Tables) {
            if (table.Cluster.empty()) {
                table.Cluster = use.Name;
            }
        }
    }

    IParser::TPtr Parser_;
};

class TGlobalAnalysis: public IGlobalAnalysis {
public:
    TGlobalAnalysis()
        : DefaultAnalysis_(MakeParser(/* isAnsiLexer = */ false))
        , AnsiAnalysis_(MakeParser(/* isAnsiLexer = */ true))
    {
    }

    TGlobalContext Analyze(TCompletionInput input, TEnvironment env) override {
        const bool isAnsiLexer = IsAnsiQuery(TString(input.Text));
        return GetSpecialized(isAnsiLexer).Analyze(std::move(input), std::move(env));
    }

private:
    IGlobalAnalysis& GetSpecialized(bool isAnsiLexer) {
        if (isAnsiLexer) {
            return AnsiAnalysis_;
        }
        return DefaultAnalysis_;
    }

    TSpecializedGlobalAnalysis DefaultAnalysis_;
    TSpecializedGlobalAnalysis AnsiAnalysis_;
};

IGlobalAnalysis::TPtr MakeGlobalAnalysis() {
    return MakeHolder<TGlobalAnalysis>();
}

} // namespace NSQLComplete

template <>
void Out<NSQLComplete::TClusterContext>(IOutputStream& out, const NSQLComplete::TClusterContext& value) {
    if (!value.Provider.empty()) {
        out << value.Provider << ":";
    }
    out << value.Name;
}

template <>
void Out<NSQLComplete::TFunctionContext>(IOutputStream& out, const NSQLComplete::TFunctionContext& value) {
    out << "TFunctionContext { ";
    out << "Name: " << value.Name;
    out << ", ArgN: " << value.ArgumentNumber;
    out << ", Arg0: " << value.Arg0.GetOrElse("None");
    out << ", Cluster: " << value.Cluster;
    out << " }";
}

template <>
void Out<NSQLComplete::TColumnContext>(IOutputStream& out, const NSQLComplete::TColumnContext& value) {
    out << "TColumnContext { ";
    out << "Tables: " << JoinSeq(", ", value.Tables);
    out << ", Columns: " << JoinSeq(", ", value.Columns);

    if (!value.WithoutByTableAlias.empty()) {
        out << ", WithoutByTableAlias: ";
        for (const auto& [tableAlias, columns] : value.WithoutByTableAlias) {
            out << tableAlias << ".[" << JoinSeq(", ", columns) << "], ";
        }
    }

    out << " }";
}