#include <Functions/FunctionsNull.h>
#include <Functions/FunctionsLogical.h>
#include <Functions/FunctionsComparison.h>
#include <Functions/FunctionsConditional.h>
#include <Functions/FunctionFactory.h>
#include <DataTypes/DataTypesNumber.h>
#include <DataTypes/DataTypeNullable.h>
#include <Columns/ColumnNullable.h>

namespace DB
{

void registerFunctionsNull(FunctionFactory & factory)
{
    factory.registerFunction<FunctionIsNull>();
    factory.registerFunction<FunctionIsNotNull>();
    factory.registerFunction<FunctionCoalesce>();
    factory.registerFunction<FunctionIfNull>();
    factory.registerFunction<FunctionNullIf>();
    factory.registerFunction<FunctionAssumeNotNull>();
    factory.registerFunction<FunctionToNullable>();
}

/// Implementation of isNull.

FunctionPtr FunctionIsNull::create(const Context &)
{
    return std::make_shared<FunctionIsNull>();
}

std::string FunctionIsNull::getName() const
{
    return name;
}

DataTypePtr FunctionIsNull::getReturnTypeImpl(const DataTypes &) const
{
    return std::make_shared<DataTypeUInt8>();
}

void FunctionIsNull::executeImpl(Block & block, const ColumnNumbers & arguments, size_t result)
{
    ColumnWithTypeAndName & elem = block.getByPosition(arguments[0]);
    if (elem.column->isNullable())
    {
        /// Merely return the embedded null map.
        ColumnNullable & nullable_col = static_cast<ColumnNullable &>(*elem.column);
        block.getByPosition(result).column = nullable_col.getNullMapColumn();
    }
    else
    {
        /// Since no element is nullable, return a zero-constant column representing
        /// a zero-filled null map.
        block.getByPosition(result).column = DataTypeUInt8().createConstColumn(elem.column->size(), UInt64(0));
    }
}

/// Implementation of isNotNull.

FunctionPtr FunctionIsNotNull::create(const Context &)
{
    return std::make_shared<FunctionIsNotNull>();
}

std::string FunctionIsNotNull::getName() const
{
    return name;
}

DataTypePtr FunctionIsNotNull::getReturnTypeImpl(const DataTypes &) const
{
    return std::make_shared<DataTypeUInt8>();
}

void FunctionIsNotNull::executeImpl(Block & block, const ColumnNumbers & arguments, size_t result)
{
    Block temp_block
    {
        block.getByPosition(arguments[0]),
        {
            nullptr,
            std::make_shared<DataTypeUInt8>(),
            ""
        },
        {
            nullptr,
            std::make_shared<DataTypeUInt8>(),
            ""
        }
    };

    FunctionIsNull{}.executeImpl(temp_block, {0}, 1);
    FunctionNot{}.execute(temp_block, {1}, 2);

    block.getByPosition(result).column = std::move(temp_block.getByPosition(2).column);
}

/// Implementation of coalesce.

static const DataTypePtr getNestedDataType(const DataTypePtr & type)
{
    if (type->isNullable())
        return static_cast<const DataTypeNullable &>(*type).getNestedType();

    return type;
}

FunctionPtr FunctionCoalesce::create(const Context & context)
{
    return std::make_shared<FunctionCoalesce>(context);
}

std::string FunctionCoalesce::getName() const
{
    return name;
}

DataTypePtr FunctionCoalesce::getReturnTypeImpl(const DataTypes & arguments) const
{
    /// Skip all NULL arguments. If any argument is non-Nullable, skip all next arguments.
    DataTypes filtered_args;
    filtered_args.reserve(arguments.size());
    for (const auto & arg : arguments)
    {
        filtered_args.push_back(arg);

        if (!arg->isNullable())
            break;
    }

    DataTypes new_args;
    for (size_t i = 0; i < filtered_args.size(); ++i)
    {
        bool is_last = i + 1 == filtered_args.size();

        if (is_last)
        {
            new_args.push_back(filtered_args[i]);
        }
        else
        {
            new_args.push_back(std::make_shared<DataTypeUInt8>());
            new_args.push_back(getNestedDataType(filtered_args[i]));
        }
    }

    if (new_args.empty())
        throw Exception("Logical error in implementation of function coalesce", ErrorCodes::LOGICAL_ERROR);
    if (new_args.size() == 1)
        return new_args.front();

    auto res = FunctionMultiIf{context}.getReturnTypeImpl(new_args);

    /// if last argument is not nullable, result should be also not nullable
    if (!new_args.back()->isNullable() && res->isNullable())
        res = getNestedDataType(res);

    return res;
}

void FunctionCoalesce::executeImpl(Block & block, const ColumnNumbers & arguments, size_t result)
{
    /// coalesce(arg0, arg1, ..., argN) is essentially
    /// multiIf(isNotNull(arg0), assumeNotNull(arg0), isNotNull(arg1), assumeNotNull(arg1), ..., argN)
    /// with constant NULL arguments removed.

    ColumnNumbers filtered_args;
    filtered_args.reserve(arguments.size());
    for (const auto & arg : arguments)
    {
        const auto & column = block.getByPosition(arg).column;

        if (column->isNull())
            continue;

        filtered_args.push_back(arg);

        if (!column->isNullable())
            break;
    }

    FunctionIsNotNull is_not_null;
    FunctionAssumeNotNull assume_not_null;
    ColumnNumbers multi_if_args;

    Block temp_block = block;

    for (size_t i = 0; i < filtered_args.size(); ++i)
    {
        size_t res_pos = temp_block.columns();
        bool is_last = i + 1 == filtered_args.size();

        if (is_last)
        {
            multi_if_args.push_back(filtered_args[i]);
        }
        else
        {
            temp_block.insert({nullptr, std::make_shared<DataTypeUInt8>(), ""});
            is_not_null.executeImpl(temp_block, {filtered_args[i]}, res_pos);
            temp_block.insert({nullptr, getNestedDataType(block.getByPosition(filtered_args[i]).type), ""});
            assume_not_null.executeImpl(temp_block, {filtered_args[i]}, res_pos + 1);

            multi_if_args.push_back(res_pos);
            multi_if_args.push_back(res_pos + 1);
        }
    }

    /// If all arguments appeared to be NULL.
    if (multi_if_args.empty())
    {
        block.getByPosition(result).column = block.getByPosition(result).type->createConstColumn(block.rows(), Null());
        return;
    }

    if (multi_if_args.size() == 1)
    {
        block.getByPosition(result).column = block.getByPosition(multi_if_args.front()).column;
        return;
    }

    FunctionMultiIf{context}.executeImpl(temp_block, multi_if_args, result);

    auto res = std::move(temp_block.getByPosition(result).column);

    /// if last argument is not nullable, result should be also not nullable
    if (!block.getByPosition(multi_if_args.back()).column->isNullable() && res->isNullable())
        res = static_cast<ColumnNullable &>(*res).getNestedColumn();

    block.getByPosition(result).column = res;
}

/// Implementation of ifNull.

FunctionPtr FunctionIfNull::create(const Context &)
{
    return std::make_shared<FunctionIfNull>();
}

std::string FunctionIfNull::getName() const
{
    return name;
}

DataTypePtr FunctionIfNull::getReturnTypeImpl(const DataTypes & arguments) const
{
    if (arguments[0]->isNull())
        return arguments[1];

    if (!arguments[0]->isNullable())
        return arguments[0];

    return FunctionIf{}.getReturnTypeImpl({std::make_shared<DataTypeUInt8>(), getNestedDataType(arguments[0]), arguments[1]});
}

void FunctionIfNull::executeImpl(Block & block, const ColumnNumbers & arguments, size_t result)
{
    /// Always null.
    if (block.getByPosition(arguments[0]).column->isNull())
    {
        block.getByPosition(result).column = block.getByPosition(arguments[1]).column;
        return;
    }

    /// Could not contain nulls, so nullIf makes no sense.
    if (!block.getByPosition(arguments[0]).column->isNullable())
    {
        block.getByPosition(result).column = block.getByPosition(arguments[0]).column;
        return;
    }

    /// ifNull(col1, col2) == if(isNotNull(col1), assumeNotNull(col1), col2)

    Block temp_block = block;

    size_t is_not_null_pos = temp_block.columns();
    temp_block.insert({nullptr, std::make_shared<DataTypeUInt8>(), ""});
    size_t assume_not_null_pos = temp_block.columns();
    temp_block.insert({nullptr, getNestedDataType(block.getByPosition(arguments[0]).type), ""});

    FunctionIsNotNull{}.executeImpl(temp_block, {arguments[0]}, is_not_null_pos);
    FunctionAssumeNotNull{}.executeImpl(temp_block, {arguments[0]}, assume_not_null_pos);
    FunctionIf{}.executeImpl(temp_block, {is_not_null_pos, assume_not_null_pos, arguments[1]}, result);

    block.getByPosition(result).column = std::move(temp_block.getByPosition(result).column);
}

/// Implementation of nullIf.

FunctionPtr FunctionNullIf::create(const Context & )
{
    return std::make_shared<FunctionNullIf>();
}

std::string FunctionNullIf::getName() const
{
    return name;
}

DataTypePtr FunctionNullIf::getReturnTypeImpl(const DataTypes & arguments) const
{
    return FunctionIf{}.getReturnTypeImpl({std::make_shared<DataTypeUInt8>(), makeNullableDataTypeIfNot(arguments[0]), arguments[0]});
}

void FunctionNullIf::executeImpl(Block & block, const ColumnNumbers & arguments, size_t result)
{
    /// nullIf(col1, col2) == multiIf(col1 == col2, NULL, col1)

    Block temp_block = block;

    size_t res_pos = temp_block.columns();
    temp_block.insert({nullptr, std::make_shared<DataTypeUInt8>(), ""});

    FunctionEquals{}.execute(temp_block, {arguments[0], arguments[1]}, res_pos);

    /// Argument corresponding to the NULL value.
    size_t null_pos = temp_block.columns();

    /// Append a NULL column.
    ColumnWithTypeAndName null_elem;
    null_elem.type = block.getByPosition(result).type;
    null_elem.column = null_elem.type->createConstColumn(temp_block.rows(), Null());
    null_elem.name = "NULL";

    temp_block.insert(null_elem);

    FunctionIf{}.executeImpl(temp_block, {res_pos, null_pos, arguments[0]}, result);

    block.getByPosition(result).column = std::move(temp_block.getByPosition(result).column);
}

/// Implementation of assumeNotNull.

FunctionPtr FunctionAssumeNotNull::create(const Context &)
{
    return std::make_shared<FunctionAssumeNotNull>();
}

std::string FunctionAssumeNotNull::getName() const
{
    return name;
}

DataTypePtr FunctionAssumeNotNull::getReturnTypeImpl(const DataTypes & arguments) const
{
    return getNestedDataType(arguments[0]);
}

void FunctionAssumeNotNull::executeImpl(Block & block, const ColumnNumbers & arguments, size_t result)
{
    const ColumnPtr & col = block.getByPosition(arguments[0]).column;
    ColumnPtr & res_col = block.getByPosition(result).column;

    if (col->isNullable())
    {
        const ColumnNullable & nullable_col = static_cast<const ColumnNullable &>(*col);
        res_col = nullable_col.getNestedColumn();
    }
    else
        res_col = col;
}

/// Implementation of toNullable.

FunctionPtr FunctionToNullable::create(const Context &)
{
    return std::make_shared<FunctionToNullable>();
}

std::string FunctionToNullable::getName() const
{
    return name;
}

DataTypePtr FunctionToNullable::getReturnTypeImpl(const DataTypes & arguments) const
{
    if (arguments[0]->isNull() || arguments[0]->isNullable())
        return arguments[0];
    return std::make_shared<DataTypeNullable>(arguments[0]);
}

void FunctionToNullable::executeImpl(Block & block, const ColumnNumbers & arguments, size_t result)
{
    const ColumnPtr & col = block.getByPosition(arguments[0]).column;

    if (col->isNull() || col->isNullable())
        block.getByPosition(result).column = col;
    else
        block.getByPosition(result).column = std::make_shared<ColumnNullable>(col,
            std::make_shared<ColumnUInt8>(block.rows(), 0));
}

}