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

#pragma once

#include <cstdint>

#include <map>

#include <memory>

#include <mutex>

#include <unordered_map>

#include <vector>

#include "common/status.h"

#include "core/block/block.h"

#include "exec/common/hash_table/phmap_fwd_decl.h"

#include "storage/tablet_info.h"

#include "util/bitmap.h"

#include "util/uid_util.h"

namespace doris {

class AdaptiveRandomBucketState {

public:

    explicit AdaptiveRandomBucketState(UniqueId load_id) : _load_id(load_id) {}

    void init_partition(int32_t sender_id, int64_t partition_id,

                        const std::vector<int64_t>& tablets,

                        const std::vector<int32_t>& bucket_seqs, int32_t start_tablet_idx);

    int64_t current_tablet(int32_t sender_id, int64_t partition_id);

    void rotate_by_tablet(int32_t sender_id, int64_t partition_id, int64_t tablet_id);

private:

    struct PartitionState {

        int64_t partition_id = -1;

        std::vector<int64_t> tablets;

        std::vector<int32_t> bucket_seqs;

        int32_t tablet_pos = 0;

        int64_t current_tablet_id = -1;

    };

    std::mutex _mutex;

    UniqueId _load_id;

    std::unordered_map<int32_t, std::unordered_map<int64_t, PartitionState>>

            _sender_partition_states;

};

class OlapTabletFinder {

public:

    // FIND_TABLET_EVERY_ROW is used for hash distribution info, which indicates that we

    // should compute tablet index for every row

    // FIND_TABLET_EVERY_BATCH is used for random distribution info, which indicates that we should

    // compute tablet index for every row batch

    // FIND_TABLET_EVERY_SINK is used for random distribution info when load_to_single_tablet set to true,

    // which indicates that we should only compute tablet index in the corresponding partition once for the

    // whole time in olap table sink

    // FIND_TABLET_RANDOM_BUCKET is used for V1 receiver-side random bucket mode.

    enum FindTabletMode {

        FIND_TABLET_EVERY_ROW,

        FIND_TABLET_EVERY_BATCH,

        FIND_TABLET_EVERY_SINK,

        FIND_TABLET_RANDOM_BUCKET

    };

    OlapTabletFinder(VOlapTablePartitionParam* vpartition, FindTabletMode mode)

            : _vpartition(vpartition), _find_tablet_mode(mode), _filter_bitmap(1024) {};

    Status find_tablets(RuntimeState* state, Block* block, int rows,

                        std::vector<VOlapTablePartition*>& partitions,

                        std::vector<uint32_t>& tablet_index, std::vector<bool>& skip,

                        std::vector<uint32_t>* miss_rows = nullptr);

    bool is_find_tablet_every_sink() {

        return _find_tablet_mode == FindTabletMode::FIND_TABLET_EVERY_SINK;

    }

    bool is_adaptive_random_bucket() const {

        return _find_tablet_mode == FindTabletMode::FIND_TABLET_RANDOM_BUCKET;

    }

    bool is_single_tablet() {

        return _find_tablet_mode != FindTabletMode::FIND_TABLET_RANDOM_BUCKET &&

               _partition_to_tablet_map.size() == 1;

    }

    // all partitions for multi find-processes of its relative writer.

    const flat_hash_set<int64_t>& partition_ids() { return _partition_ids; }

    int64_t num_filtered_rows() const { return _num_filtered_rows; }

    int64_t num_immutable_partition_filtered_rows() const {

        return _num_immutable_partition_filtered_rows;

    }

    Bitmap& filter_bitmap() { return _filter_bitmap; }

private:

    VOlapTablePartitionParam* _vpartition = nullptr;

    FindTabletMode _find_tablet_mode;

    std::map<VOlapTablePartition*, int64_t> _partition_to_tablet_map;

    flat_hash_set<int64_t> _partition_ids;

    int64_t _num_filtered_rows = 0;

    int64_t _num_immutable_partition_filtered_rows = 0;

    Bitmap _filter_bitmap;

};

} // namespace doris