// 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 <bthread/countdown_event.h>

#include <gen_cpp/BackendService.h>

#include <condition_variable>

#include <cstddef>

#include <deque>

#include <mutex>

#include <string>

#include <thread>

#include <unordered_set>

#include <vector>

#include "cloud/cloud_storage_engine.h"

#include "cloud/cloud_tablet.h"

#include "common/status.h"

#include "util/threadpool.h"

namespace doris {

enum class DownloadType {

    BE,

    S3,

};

// Filter for event-driven warmup jobs.

// nullopt = cluster-level (no table filter, warm up all tables)

// has_value = table-level filter (only warm up tables in the set)

using EventDrivenJobFilter = std::optional<std::unordered_set<int64_t>>;

struct JobReplicaInfo {

    int64_t job_id;

    TReplicaInfo replica;

};

struct JobMeta {

    JobMeta() = default;

    JobMeta(const TJobMeta& meta);

    DownloadType download_type;

    std::string be_ip;

    int32_t brpc_port;

    std::vector<int64_t> tablet_ids;

};

// manager for

// table warm up

// cluster warm up

// balance peer addr cache

class CloudWarmUpManager {

public:

    explicit CloudWarmUpManager(CloudStorageEngine& engine);

    ~CloudWarmUpManager();

    // Set the job id if the id is zero

    Status check_and_set_job_id(int64_t job_id);

    // Set the batch id to record download progress

    Status check_and_set_batch_id(int64_t job_id, int64_t batch_id, bool* retry = nullptr);

    // Add the dowload job

    void add_job(const std::vector<TJobMeta>& job_metas);

#ifdef BE_TEST

    void consumer_job();

#endif

    // Get the job state tuple<cur_job_id, cur_batch_id, pending_job_metas_size, _finish_job_size>

    std::tuple<int64_t, int64_t, int64_t, int64_t> get_current_job_state();

    // Cancel the job

    Status clear_job(int64_t job_id);

    Status set_event(int64_t job_id, TWarmUpEventType::type event, bool clear = false,

                     const std::vector<int64_t>* table_ids = nullptr);

    // If `sync_wait_timeout_ms` <= 0, the function will send the warm-up RPC

    // and return immediately without waiting for the warm-up to complete.

    // If `sync_wait_timeout_ms` > 0, the function will wait for the warm-up

    // to finish or until the specified timeout (in milliseconds) is reached.

    //

    // @param rs_meta Metadata of the rowset to be warmed up.

    // @param sync_wait_timeout_ms Timeout in milliseconds to wait for the warm-up

    //                              to complete. Non-positive value means no waiting.

    void warm_up_rowset(RowsetMeta& rs_meta, int64_t table_id, int64_t sync_wait_timeout_ms = -1);

    void recycle_cache(int64_t tablet_id, const std::vector<RecycledRowsets>& rowsets);

    // Balance warm up cache management methods

    void record_balanced_tablet(int64_t tablet_id, const std::string& host, int32_t brpc_port);

    std::optional<std::pair<std::string, int32_t>> get_balanced_tablet_info(int64_t tablet_id);

    void remove_balanced_tablet(int64_t tablet_id);

    void remove_balanced_tablets(const std::vector<int64_t>& tablet_ids);

    bool is_balanced_tablet_expired(const std::chrono::system_clock::time_point& ctime) const;

    std::unordered_map<int64_t, std::pair<std::string, int32_t>> get_all_balanced_tablets() const;

private:

    struct WarmUpRowsetFailure {

        int code;

        std::string reason;

    };

    static Status _build_warm_up_rowset_result(const std::vector<WarmUpRowsetFailure>& failures,

                                               size_t replica_count, int64_t tablet_id,

                                               int64_t table_id, const std::string& rowset_id);

    void schedule_remove_balanced_tablet(int64_t tablet_id);

    static void clean_up_expired_mappings(void* arg);

    void handle_jobs();

    Status _do_warm_up_rowset(RowsetMeta& rs_meta, int64_t table_id,

                              std::vector<JobReplicaInfo>& replicas, int64_t sync_wait_timeout_ms,

                              bool skip_existence_check);

    std::vector<JobReplicaInfo> get_replica_info(int64_t tablet_id, int64_t table_id,

                                                 bool bypass_cache, bool& cache_hit);

    void _warm_up_rowset(RowsetMeta& rs_meta, int64_t table_id, int64_t sync_wait_timeout_ms);

    void _recycle_cache(int64_t tablet_id, const std::vector<RecycledRowsets>& rowsets);

    void submit_download_tasks(io::Path path, int64_t file_size, io::FileSystemSPtr file_system,

                               int64_t expiration_time,

                               std::shared_ptr<bthread::CountdownEvent> wait, bool is_index = false,

                               std::function<void(Status)> done_cb = nullptr,

                               int64_t tablet_id = -1);

    std::mutex _mtx;

    std::condition_variable _cond;

    int64_t _cur_job_id {0};

    int64_t _cur_batch_id {-1};

    std::deque<std::shared_ptr<JobMeta>> _pending_job_metas;

    std::vector<std::shared_ptr<JobMeta>> _finish_job;

    std::thread _download_thread;

    bool _closed {false};

    // the attribute for compile in ut

    [[maybe_unused]] CloudStorageEngine& _engine;

    // timestamp, info

    using CacheEntry = std::pair<std::chrono::steady_clock::time_point, TReplicaInfo>;

    // tablet_id -> entry

    using Cache = std::unordered_map<int64_t, CacheEntry>;

    // job_id -> cache

    std::unordered_map<int64_t, Cache> _tablet_replica_cache;

    // job_id -> table filter (nullopt = cluster-level, no filter)

    std::unordered_map<int64_t, EventDrivenJobFilter> _event_driven_filters;

    std::unique_ptr<ThreadPool> _thread_pool;

    std::unique_ptr<ThreadPoolToken> _thread_pool_token;

    // Sharded lock for better performance

    static constexpr size_t SHARD_COUNT = 10240;

    struct Shard {

        mutable std::mutex mtx;

        std::unordered_map<int64_t, JobMeta> tablets;

    };

    std::array<Shard, SHARD_COUNT> _balanced_tablets_shards;

    // Helper methods for shard operations

    size_t get_shard_index(int64_t tablet_id) const {

        return std::hash<int64_t> {}(tablet_id) % SHARD_COUNT;

    }

    Shard& get_shard(int64_t tablet_id) {

        return _balanced_tablets_shards[get_shard_index(tablet_id)];

    }

};

} // namespace doris