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

#include <array>

#include <atomic>

#include <chrono>

#include <map>

#include <memory>

#include <mutex>

#include <shared_mutex>

#include <string>

#include <string_view>

#include <unordered_map>

#include <vector>

#include "cloud/cloud_throttle_state_machine.h"

#include "util/countdown_latch.h"

#include "util/thread.h"

namespace doris::cloud {

// Strict QPS limiter that doesn't allow burst

// Unlike token bucket, it strictly enforces fixed intervals between requests

class StrictQpsLimiter {

public:

    using Clock = std::chrono::steady_clock;

    explicit StrictQpsLimiter(double qps);

    // Returns the time point when the request is allowed to execute

    // Caller should sleep until this time point

    Clock::time_point reserve();

    // Dynamically update the QPS limit, optionally discarding queued reservations.

    void update_qps(double new_qps, bool reset_reservation = false);

    // Get current QPS limit

    double get_qps() const;

private:

    mutable std::mutex _mtx;

    int64_t _interval_ns;

    Clock::time_point _next_allowed_time;

};

// QPS counter for a single (table, RPC type) pair using bvar

class TableRpcQpsCounter {

public:

    TableRpcQpsCounter(int64_t table_id, LoadRelatedRpc rpc_type, int window_sec);

    ~TableRpcQpsCounter() = default;

    // Record one RPC call

    void increment();

    // Get current QPS (average over the configured time window)

    double get_qps() const;

    int64_t last_record_time_us() const {

        return _last_record_time_us.load(std::memory_order_relaxed);

    }

    int64_t table_id() const { return _table_id; }

    LoadRelatedRpc rpc_type() const { return _rpc_type; }

private:

    int64_t _table_id;

    LoadRelatedRpc _rpc_type;

    std::unique_ptr<bvar::Adder<int64_t>> _counter;

    std::unique_ptr<bvar::PerSecond<bvar::Adder<int64_t>>> _qps;

    std::atomic<int64_t> _last_record_time_us {0};

};

// Registry managing QPS counters for all tables

class TableRpcQpsRegistry {

public:

    TableRpcQpsRegistry();

    TableRpcQpsRegistry(std::chrono::milliseconds cleanup_interval,

                        std::chrono::milliseconds inactive_timeout);

    ~TableRpcQpsRegistry();

    // Record one RPC call for the given table

    void record(LoadRelatedRpc rpc_type, int64_t table_id);

    // Get the top-k tables with highest QPS for the given RPC type

    // Returns: [(table_id, qps), ...] sorted by qps in descending order

    std::vector<std::pair<int64_t, double>> get_top_k_tables(LoadRelatedRpc rpc_type, int k) const;

    // Get QPS for a specific table on a specific RPC type

    double get_qps(LoadRelatedRpc rpc_type, int64_t table_id) const;

    // Clean up counters for tables that have been inactive for a long time

    size_t cleanup_inactive_tables();

    size_t get_tracked_table_count(LoadRelatedRpc rpc_type) const;

private:

    void _cleanup_thread_callback();

    mutable std::shared_mutex _mutex;

    // rpc_type -> (table_id -> counter)

    std::array<std::unordered_map<int64_t, std::unique_ptr<TableRpcQpsCounter>>,

               static_cast<size_t>(LoadRelatedRpc::COUNT)>

            _counters;

    const std::chrono::milliseconds _cleanup_interval;

    const std::chrono::milliseconds _inactive_timeout;

    std::shared_ptr<Thread> _cleanup_thread;

    CountDownLatch _cleanup_stop_latch;

};

struct TableRpcThrottleDecision {

    std::chrono::steady_clock::time_point wait_until;

    double qps_limit {0};

    bool dry_run {false};

};

// Table-level throttler managing StrictQpsLimiter for each (RPC type, table) pair

class TableRpcThrottler {

public:

    TableRpcThrottler();

    ~TableRpcThrottler() = default;

    // Called before RPC execution, returns the time point when execution is allowed

    // Returns now if no limit is set

    std::chrono::steady_clock::time_point throttle(LoadRelatedRpc rpc_type, int64_t table_id);

    TableRpcThrottleDecision throttle(LoadRelatedRpc rpc_type, int64_t table_id, bool dry_run);

    // Log suppression is independent for every RPC type.

    bool should_log(LoadRelatedRpc rpc_type, int64_t now_us);

    // Set or update the QPS limit for a table

    void set_qps_limit(LoadRelatedRpc rpc_type, int64_t table_id, double qps_limit,

                       bool reset_reservation = false);

    // Remove the QPS limit for a table

    void remove_qps_limit(LoadRelatedRpc rpc_type, int64_t table_id);

    // Get current QPS limit (returns 0 if not set)

    double get_qps_limit(LoadRelatedRpc rpc_type, int64_t table_id) const;

    // Check if a limit exists for the given (rpc_type, table_id)

    bool has_limit(LoadRelatedRpc rpc_type, int64_t table_id) const;

    // Get the number of throttled tables for a given RPC type

    size_t get_throttled_table_count(LoadRelatedRpc rpc_type) const;

    // Get all currently throttled entries: (rpc_type, table_id, qps_limit)

    struct ThrottleEntry {

        LoadRelatedRpc rpc_type;

        int64_t table_id;

        double qps_limit;

    };

    std::vector<ThrottleEntry> get_all_throttled_entries() const;

private:

    mutable std::shared_mutex _mutex;

    std::map<std::pair<LoadRelatedRpc, int64_t>, std::unique_ptr<StrictQpsLimiter>> _limiters;

    std::array<std::atomic<int64_t>, static_cast<size_t>(LoadRelatedRpc::COUNT)> _next_log_time_us;

    // bvar: current throttled table count per RPC type

    std::array<std::unique_ptr<bvar::Status<size_t>>, static_cast<size_t>(LoadRelatedRpc::COUNT)>

            _throttled_table_counts;

};

// MS backpressure handler that coordinates QPS statistics, throttle upgrade and downgrade

// Uses state machine for decisions, providing better testability

class MSBackpressureHandler {

public:

    MSBackpressureHandler(TableRpcQpsRegistry* qps_registry, TableRpcThrottler* throttler);

    ~MSBackpressureHandler();

    // Called when receiving MS_BUSY response

    // Returns true if throttle upgrade was triggered

    bool on_ms_busy();

    // Called before RPC execution and returns the actual or dry-run throttle decision.

    TableRpcThrottleDecision before_rpc(LoadRelatedRpc rpc_type, int64_t table_id);

    bool should_log_throttle(LoadRelatedRpc rpc_type, int64_t now_us);

    double get_current_qps(LoadRelatedRpc rpc_type, int64_t table_id) const;

    // Called after RPC execution, records QPS statistics

    void after_rpc(LoadRelatedRpc rpc_type, int64_t table_id);

    // Runtime update parameters

    void update_throttle_params(RpcThrottleParams params);

    void update_coordinator_params(ThrottleCoordinatorParams params);

    // Get seconds since last MS_BUSY (for monitoring)

    int64_t seconds_since_last_ms_busy() const;

    // Query current state

    size_t upgrade_level() const;

    int64_t ticks_since_last_ms_busy() const;

    int64_t ticks_since_last_upgrade() const;

private:

    // Background thread that periodically advances time

    void _tick_thread_callback();

    // Advance time by specified ticks, handle any triggered events (e.g., downgrade)

    void _advance_time(int64_t ticks);

    // Apply actions to the throttler

    void _apply_actions(const std::vector<RpcThrottleAction>& actions);

    // Build QPS snapshot from registry

    std::vector<RpcQpsSnapshot> _build_qps_snapshot() const;

    TableRpcQpsRegistry* _qps_registry;

    TableRpcThrottler* _throttler;

    mutable std::mutex _transition_mutex;

    // State machine components

    std::unique_ptr<RpcThrottleStateMachine> _state_machine;

    std::unique_ptr<RpcThrottleCoordinator> _coordinator;

    // Background thread for periodic tick

    std::shared_ptr<Thread> _tick_thread;

    CountDownLatch _stop_latch;

    // For bvar compatibility only - track approximate seconds since last MS_BUSY

    mutable std::mutex _mutex;

    std::chrono::steady_clock::time_point _last_ms_busy_time;

};

// Global bvar metrics for backpressure handling

extern bvar::Adder<uint64_t> g_backpressure_upgrade_count;

extern bvar::Window<bvar::Adder<uint64_t>> g_backpressure_upgrade_60s;

extern bvar::Adder<uint64_t> g_backpressure_downgrade_count;

extern bvar::Window<bvar::Adder<uint64_t>> g_backpressure_downgrade_60s;

extern bvar::Adder<uint64_t> g_ms_busy_count;

extern bvar::Window<bvar::Adder<uint64_t>> g_ms_busy_60s;

// Per-RPC-type throttle wait latency recorders

bvar::LatencyRecorder* get_throttle_wait_recorder(LoadRelatedRpc rpc);

} // namespace doris::cloud