/*
 * ChaffScheduler.cpp — Idle chaff injection for tgnet
 * Ported from telemt/tdlib-obf (MIT License, Copyright 2026 telemt community)
 */

#include "ChaffScheduler.h"

#include <openssl/rand.h>
#include <algorithm>
#include <cmath>
#include <limits>

namespace stealth {

ChaffPolicy default_chaff_policy() {
    ChaffPolicy p;
    p.enabled             = false; // opt-in per connection
    p.idle_threshold_ms   = 15000;
    p.min_interval_ms     = 5000.0;
    p.max_bytes_per_minute = 4096;
    p.record_size_lo      = 50;
    p.record_size_hi      = 800;
    return p;
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

static uint32_t rand_u32() {
    uint32_t v;
    RAND_bytes(reinterpret_cast<uint8_t *>(&v), 4);
    return v;
}

static uint32_t rand_bounded(uint32_t n) {
    if (n <= 1) return 0;
    uint32_t t = static_cast<uint32_t>(-n) % n;
    for (;;) { uint32_t v = rand_u32(); if (v >= t) return v % n; }
}

// ---------------------------------------------------------------------------
// Constructor
// ---------------------------------------------------------------------------

ChaffScheduler::ChaffScheduler(const ChaffPolicy &policy,
                                IptController &ipt, double now_sec)
    : policy_(policy), ipt_(ipt) {
    if (policy_.enabled) {
        schedule_after_activity(now_sec);
    }
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

void ChaffScheduler::note_activity(double now_sec) {
    if (!policy_.enabled || !std::isfinite(now_sec)) {
        disarm();
        return;
    }
    prune_budget(now_sec);
    schedule_after_activity(now_sec);
}

void ChaffScheduler::note_chaff_emitted(double now_sec, size_t bytes) {
    if (!policy_.enabled || !std::isfinite(now_sec)) {
        disarm();
        return;
    }
    prune_budget(now_sec);
    budget_window_.push_back({now_sec, bytes});
    schedule_after_chaff(now_sec);
}

bool ChaffScheduler::should_emit(double now_sec,
                                  bool has_pending_data, bool can_write) const {
    if (!policy_.enabled || has_pending_data || !can_write
        || pending_size_ <= 0 || !std::isfinite(now_sec)
        || !std::isfinite(next_send_at_) || next_send_at_ == 0.0) {
        return false;
    }
    if (now_sec + 1e-9 < next_send_at_) return false;
    return budget_allows(now_sec, static_cast<size_t>(pending_size_));
}

double ChaffScheduler::get_wakeup(double now_sec,
                                   bool has_pending_data, bool can_write) const {
    if (!policy_.enabled || has_pending_data || !can_write
        || next_send_at_ == 0.0 || !std::isfinite(now_sec)
        || !std::isfinite(next_send_at_)) {
        return 0.0;
    }
    double wakeup = next_send_at_;
    if (pending_size_ > 0) {
        double resume = budget_resume_at(now_sec, static_cast<size_t>(pending_size_));
        if (resume != 0.0) wakeup = std::max(wakeup, resume);
    }
    return wakeup;
}

int32_t ChaffScheduler::next_record_size() {
    return pending_size_ > 0 ? pending_size_ : sample_record_size();
}

// ---------------------------------------------------------------------------
// Private
// ---------------------------------------------------------------------------

void ChaffScheduler::schedule_after_activity(double now_sec) {
    pending_size_ = sample_record_size();
    double idle_thresh = static_cast<double>(policy_.idle_threshold_ms) / 1000.0;
    double interval    = sample_interval_sec();
    if (!std::isfinite(interval) || interval < 0.0) { disarm(); return; }
    next_send_at_ = now_sec + idle_thresh + interval;
    if (!std::isfinite(next_send_at_)) disarm();
}

void ChaffScheduler::schedule_after_chaff(double now_sec) {
    pending_size_ = sample_record_size();
    double interval = sample_interval_sec();
    if (!std::isfinite(interval) || interval < 0.0) { disarm(); return; }
    next_send_at_ = now_sec + interval;
    if (!std::isfinite(next_send_at_)) disarm();
}

void ChaffScheduler::disarm() {
    budget_window_.clear();
    next_send_at_ = 0.0;
    pending_size_ = 0;
}

void ChaffScheduler::prune_budget(double now_sec) {
    while (!budget_window_.empty()
           && budget_window_.front().at + kBudgetWindow <= now_sec) {
        budget_window_.pop_front();
    }
}

double ChaffScheduler::budget_resume_at(double now_sec, size_t target) const {
    const size_t limit = policy_.max_bytes_per_minute;

    if (target > limit) {
        double d = now_sec + kBudgetWindow;
        return std::isfinite(d) ? d : std::numeric_limits<double>::max();
    }

    size_t used = 0;
    double earliest = 0.0, latest = 0.0;
    for (const auto &s : budget_window_) {
        if (s.at + kBudgetWindow <= now_sec) continue;
        used += s.bytes;
        if (earliest == 0.0) earliest = s.at + kBudgetWindow;
        latest = s.at + kBudgetWindow;
    }

    if (used + target <= limit) return 0.0; // budget fine

    size_t need_release = used + target - limit;
    size_t released = 0;
    for (const auto &s : budget_window_) {
        double expire = s.at + kBudgetWindow;
        if (expire <= now_sec) continue;
        released += s.bytes;
        if (released >= need_release) return expire;
    }
    return latest;
}

bool ChaffScheduler::budget_allows(double now_sec, size_t target) const {
    return budget_resume_at(now_sec, target) == 0.0;
}

double ChaffScheduler::sample_interval_sec() {
    uint64_t delay_us = ipt_.sample_idle_delay_us();
    double min_sec = policy_.min_interval_ms / 1000.0;
    double sampled = static_cast<double>(delay_us) / 1e6;
    return std::max(min_sec, sampled);
}

int32_t ChaffScheduler::sample_record_size() {
    if (policy_.record_size_lo >= policy_.record_size_hi)
        return policy_.record_size_lo;
    uint32_t range = static_cast<uint32_t>(policy_.record_size_hi
                                          - policy_.record_size_lo + 1);
    return policy_.record_size_lo + static_cast<int32_t>(rand_bounded(range));
}

}  // namespace stealth