#include "rate_limiter.h"

#include <algorithm>

namespace mosis {

//=============================================================================
// CONSTRUCTOR (with default limits)
//=============================================================================

RateLimiter::RateLimiter() {
    // Network operations
    SetLimit("network.request", {10.0, 100.0});      // 10/sec, burst 100
    SetLimit("network.websocket", {2.0, 10.0});      // 2/sec, burst 10
    SetLimit("network.download", {5.0, 20.0});       // 5/sec, burst 20

    // Storage operations
    SetLimit("storage.read", {100.0, 500.0});        // 100/sec, burst 500
    SetLimit("storage.write", {20.0, 100.0});        // 20/sec, burst 100
    SetLimit("storage.delete", {10.0, 50.0});        // 10/sec, burst 50
    SetLimit("database.query", {50.0, 200.0});       // 50/sec, burst 200

    // Hardware access
    SetLimit("camera.capture", {30.0, 30.0});        // 30 fps max
    SetLimit("microphone.record", {1.0, 1.0});       // 1 session at a time
    SetLimit("location.request", {1.0, 5.0});        // 1/sec, burst 5
    SetLimit("sensor.read", {60.0, 60.0});           // 60 Hz max

    // Timers
    SetLimit("timer.create", {10.0, 100.0});         // 10/sec, burst 100

    // Crypto
    SetLimit("crypto.random", {100.0, 1000.0});      // 100/sec, burst 1000
    SetLimit("crypto.hash", {100.0, 1000.0});        // 100/sec, burst 1000
}

//=============================================================================
// CONFIGURATION
//=============================================================================

void RateLimiter::SetLimit(const std::string& operation, const RateLimitConfig& config) {
    std::lock_guard<std::mutex> lock(m_mutex);
    m_configs[operation] = config;
}

const RateLimitConfig* RateLimiter::GetLimit(const std::string& operation) const {
    std::lock_guard<std::mutex> lock(m_mutex);
    auto it = m_configs.find(operation);
    if (it != m_configs.end()) {
        return &it->second;
    }
    return nullptr;
}

//=============================================================================
// CHECKING
//=============================================================================

bool RateLimiter::Check(const std::string& app_id, const std::string& operation) {
    std::lock_guard<std::mutex> lock(m_mutex);

    // Find config
    auto config_it = m_configs.find(operation);
    if (config_it == m_configs.end()) {
        // No limit configured, allow by default
        return true;
    }

    const auto& config = config_it->second;
    auto& bucket = GetBucket(app_id, operation);

    // Refill based on elapsed time
    Refill(bucket, config);

    // Check if we have a token
    if (bucket.tokens >= 1.0) {
        bucket.tokens -= 1.0;
        return true;
    }

    return false;
}

bool RateLimiter::CanProceed(const std::string& app_id, const std::string& operation) const {
    std::lock_guard<std::mutex> lock(m_mutex);

    // Find config
    auto config_it = m_configs.find(operation);
    if (config_it == m_configs.end()) {
        return true;  // No limit
    }

    const auto& config = config_it->second;
    std::string key = MakeKey(app_id, operation);

    auto bucket_it = m_buckets.find(key);
    if (bucket_it == m_buckets.end()) {
        return true;  // New bucket would have full tokens
    }

    // Make a copy to check without modifying
    Bucket bucket = bucket_it->second;
    Refill(bucket, config);

    return bucket.tokens >= 1.0;
}

double RateLimiter::GetTokens(const std::string& app_id, const std::string& operation) const {
    std::lock_guard<std::mutex> lock(m_mutex);

    std::string key = MakeKey(app_id, operation);
    auto bucket_it = m_buckets.find(key);

    if (bucket_it == m_buckets.end()) {
        // Check if there's a config
        auto config_it = m_configs.find(operation);
        if (config_it != m_configs.end()) {
            return config_it->second.max_tokens;  // Would start with full
        }
        return 0.0;
    }

    // Find config to refill
    auto config_it = m_configs.find(operation);
    if (config_it == m_configs.end()) {
        return bucket_it->second.tokens;
    }

    // Make a copy to check without modifying
    Bucket bucket = bucket_it->second;
    Refill(bucket, config_it->second);

    return bucket.tokens;
}

//=============================================================================
// RESET
//=============================================================================

void RateLimiter::ResetApp(const std::string& app_id) {
    std::lock_guard<std::mutex> lock(m_mutex);

    // Find and remove all buckets for this app
    std::string prefix = app_id + ":";
    for (auto it = m_buckets.begin(); it != m_buckets.end(); ) {
        if (it->first.rfind(prefix, 0) == 0) {  // starts with app_id:
            it = m_buckets.erase(it);
        } else {
            ++it;
        }
    }
}

void RateLimiter::ClearAll() {
    std::lock_guard<std::mutex> lock(m_mutex);
    m_buckets.clear();
}

//=============================================================================
// INTERNAL
//=============================================================================

void RateLimiter::Refill(Bucket& bucket, const RateLimitConfig& config) const {
    auto now = std::chrono::steady_clock::now();
    auto elapsed = std::chrono::duration<double>(now - bucket.last_refill);

    // Add tokens based on elapsed time
    double new_tokens = bucket.tokens + (elapsed.count() * config.tokens_per_second);

    // Cap at max
    bucket.tokens = std::min(new_tokens, config.max_tokens);
    bucket.last_refill = now;
}

RateLimiter::Bucket& RateLimiter::GetBucket(const std::string& app_id,
                                             const std::string& operation) {
    std::string key = MakeKey(app_id, operation);

    auto it = m_buckets.find(key);
    if (it != m_buckets.end()) {
        return it->second;
    }

    // Create new bucket with full tokens
    auto config_it = m_configs.find(operation);
    double initial = (config_it != m_configs.end()) ? config_it->second.max_tokens : 1.0;

    m_buckets[key] = Bucket{
        .tokens = initial,
        .last_refill = std::chrono::steady_clock::now()
    };

    return m_buckets[key];
}

std::string RateLimiter::MakeKey(const std::string& app_id, const std::string& operation) {
    return app_id + ":" + operation;
}

//=============================================================================
// GLOBAL INSTANCE
//=============================================================================

RateLimiter& GetRateLimiter() {
    static RateLimiter instance;
    return instance;
}

} // namespace mosis
