// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements.  See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership.  The ASF licenses this file

// to you under the Apache License, Version 2.0 (the

// "License"); you may not use this file except in compliance

// with the License.  You may obtain a copy of the License at

//

//   http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing,

// software distributed under the License is distributed on an

// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

// KIND, either express or implied.  See the License for the

// specific language governing permissions and limitations

// under the License.

#include <fmt/format.h>

#include <glog/logging.h>

#include <stddef.h>

#include <algorithm>

#include <boost/iterator/iterator_facade.hpp>

#include <memory>

#include <string>

#include <utility>

#include "common/status.h"

#include "core/assert_cast.h"

#include "core/block/block.h"

#include "core/block/column_numbers.h"

#include "core/block/column_with_type_and_name.h"

#include "core/column/column.h"

#include "core/column/column_array.h"

#include "core/column/column_nullable.h"

#include "core/column/column_vector.h"

#include "core/data_type/data_type.h"

#include "core/data_type/data_type_array.h"

#include "core/data_type/data_type_nullable.h"

#include "core/data_type/data_type_number.h"

#include "core/types.h"

#include "exprs/aggregate/aggregate_function.h"

#include "exprs/function/function.h"

#include "exprs/function/function_helpers.h"

#include "exprs/function/simple_function_factory.h"

namespace doris {

class FunctionContext;

} // namespace doris

namespace doris {

class FunctionArrayEnumerate : public IFunction {

public:

    static constexpr auto name = "array_enumerate";

    static FunctionPtr create() { return std::make_shared<FunctionArrayEnumerate>(); }

    String get_name() const override { return name; }

    size_t get_number_of_arguments() const override { return 1; }

    DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {

        const DataTypeArray* array_type =

                check_and_get_data_type<DataTypeArray>(remove_nullable(arguments[0]).get());

        if (!array_type) {

            throw doris::Exception(

                    ErrorCode::INVALID_ARGUMENT,

                    "First argument for function {} .must be an array but it type is {}",

                    get_name(), arguments[0]->get_name());

        }

        auto nested_type = assert_cast<const DataTypeArray&>(*array_type).get_nested_type();

        bool is_nested_nullable = nested_type->is_nullable();

        bool is_nullable = arguments[0]->is_nullable();

        auto return_nested_type = std::make_shared<DataTypeInt64>();

        DataTypePtr return_type = std::make_shared<DataTypeArray>(

                is_nested_nullable ? make_nullable(return_nested_type) : return_nested_type);

        if (is_nullable) {

            return_type = make_nullable(return_type);

        }

        return return_type;

    }

    Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,

                        uint32_t result, size_t input_rows_count) const override {

        auto left_column =

                block.get_by_position(arguments[0]).column->convert_to_full_column_if_const();

        const ColumnArray* array =

                check_and_get_column<ColumnArray>(remove_nullable(left_column->get_ptr()).get());

        if (!array) {

            return Status::RuntimeError(

                    fmt::format("Illegal column {}, of first argument of function {}",

                                left_column->get_name(), get_name()));

        }

        const ColumnArray::Offsets64& offsets = array->get_offsets();

        auto res_nested = ColumnInt64::create();

        ColumnInt64::Container& res_values = res_nested->get_data();

        res_values.resize(array->get_data().size());

        ColumnArray::Offset64 prev_off = 0;

        for (auto off : offsets) {

            for (ColumnArray::Offset64 j = prev_off; j < off; ++j) res_values[j] = j - prev_off + 1;

            prev_off = off;

        }

        ColumnPtr nested_column = res_nested->get_ptr();

        if (array->get_data().is_nullable()) {

            nested_column = ColumnNullable::create(nested_column,

                                                   ColumnUInt8::create(nested_column->size(), 0));

        }

        ColumnPtr res_column =

                ColumnArray::create(std::move(nested_column), array->get_offsets_ptr());

        if (block.get_by_position(arguments[0]).column->is_nullable()) {

            const ColumnNullable* nullable =

                    check_and_get_column<ColumnNullable>(left_column.get());

            res_column = ColumnNullable::create(

                    res_column, nullable->get_null_map_column().clone_resized(nullable->size()));

        }

        block.replace_by_position(result, std::move(res_column));

        return Status::OK();

    }

};

void register_function_array_enumerate(SimpleFunctionFactory& factory) {

    factory.register_function<FunctionArrayEnumerate>();

}

} // namespace doris