// 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.

// This file is copied from

// https://github.com/ClickHouse/ClickHouse/blob/master/AggregateFunctionRetention.h

// and modified by Doris

#pragma once

#include <stddef.h>

#include <stdint.h>

#include <boost/iterator/iterator_facade.hpp>

#include <memory>

#include "common/exception.h"

#include "common/status.h"

#include "core/assert_cast.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_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 "util/var_int.h"

namespace doris {

class Arena;

class BufferReadable;

class BufferWritable;

} // namespace doris

namespace doris {

struct RetentionState {

    static constexpr size_t MAX_EVENTS = 32;

    uint8_t events[MAX_EVENTS] = {0};

    RetentionState() = default;

    void reset() {

        for (int64_t i = 0; i < MAX_EVENTS; i++) {

            events[i] = 0;

        }

    }

    void set(int event) { events[event] = 1; }

    void merge(const RetentionState& other) {

        for (int64_t i = 0; i < MAX_EVENTS; i++) {

            events[i] |= other.events[i];

        }

    }

    void write(BufferWritable& out) const {

        int64_t serialized_events = 0;

        for (int64_t i = 0; i < MAX_EVENTS; i++) {

            serialized_events |= events[i];

            serialized_events <<= 1;

        }

        write_var_int(serialized_events, out);

    }

    void read(BufferReadable& in) {

        int64_t serialized_events = 0;

        uint64_t u_serialized_events = 0;

        read_var_int(serialized_events, in);

        u_serialized_events = serialized_events;

        u_serialized_events >>= 1;

        for (int64_t i = MAX_EVENTS - 1; i >= 0; i--) {

            events[i] = (uint8_t)(1 & u_serialized_events);

            u_serialized_events >>= 1;

        }

    }

    void insert_result_into(IColumn& to, size_t events_size, const uint8_t* arg_events) const {

        auto& data_to = assert_cast<ColumnUInt8&, TypeCheckOnRelease::DISABLE>(to).get_data();

        ColumnArray::Offset64 current_offset = data_to.size();

        data_to.resize(current_offset + events_size);

        bool first_flag = arg_events[0];

        data_to[current_offset] = first_flag;

        ++current_offset;

        for (size_t i = 1; i < events_size; ++i) {

            data_to[current_offset] = (first_flag && arg_events[i]);

            ++current_offset;

        }

    }

};

class AggregateFunctionRetention final

        : public IAggregateFunctionDataHelper<RetentionState, AggregateFunctionRetention>,

          VarargsExpression,

          NullableAggregateFunction {

public:

    AggregateFunctionRetention(const DataTypes& argument_types_)

            : IAggregateFunctionDataHelper<RetentionState, AggregateFunctionRetention>(

                      argument_types_) {

        // RetentionState only has room for MAX_EVENTS(32) events (fixed-size events[] array,

        // plus an int64 serialized bitmap). More params would overflow events[] in add()/

        // insert_result_into() and corrupt the heap, so reject it at construction time.

        if (argument_types_.size() > RetentionState::MAX_EVENTS) {

            throw Exception(ErrorCode::INVALID_ARGUMENT,

                            "retention function can accept at most {} params, but got {}",

                            RetentionState::MAX_EVENTS, argument_types_.size());

        }

    }

    String get_name() const override { return "retention"; }

    DataTypePtr get_return_type() const override {

        return std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeUInt8>()));

    }

    void reset(AggregateDataPtr __restrict place) const override { this->data(place).reset(); }

    void add(AggregateDataPtr __restrict place, const IColumn** columns, const ssize_t row_num,

             Arena&) const override {

        for (int i = 0; i < get_argument_types().size(); i++) {

            auto event = assert_cast<const ColumnUInt8*, TypeCheckOnRelease::DISABLE>(columns[i])

                                 ->get_data()[row_num];

            if (event) {

                this->data(place).set(i);

            }

        }

    }

    void check_input_columns_type(const IColumn** columns) const override {

        for (size_t i = 0; i < get_argument_types().size(); ++i) {

            this->template check_argument_column_type<ColumnUInt8>(columns[i]);

        }

    }

    void merge(AggregateDataPtr __restrict place, ConstAggregateDataPtr rhs,

               Arena&) const override {

        this->data(place).merge(this->data(rhs));

    }

    void serialize(ConstAggregateDataPtr __restrict place, BufferWritable& buf) const override {

        this->data(place).write(buf);

    }

    void deserialize(AggregateDataPtr __restrict place, BufferReadable& buf,

                     Arena&) const override {

        this->data(place).read(buf);

    }

    void insert_result_into(ConstAggregateDataPtr __restrict place, IColumn& to) const override {

        auto& to_arr = assert_cast<ColumnArray&, TypeCheckOnRelease::DISABLE>(to);

        auto& to_nested_col = to_arr.get_data();

        if (is_column_nullable(to_nested_col)) {

            auto col_null = reinterpret_cast<ColumnNullable*>(&to_nested_col);

            this->data(place).insert_result_into(col_null->get_nested_column(),

                                                 get_argument_types().size(),

                                                 this->data(place).events);

            col_null->get_null_map_data().resize_fill(col_null->get_nested_column().size(), 0);

        } else {

            this->data(place).insert_result_into(to_nested_col, get_argument_types().size(),

                                                 this->data(place).events);

        }

        to_arr.get_offsets().push_back(to_nested_col.size());

    }

};

} // namespace doris