/*
 * IptController.cpp — Inter-Packet Timing obfuscation for tgnet
 * Ported from telemt/tdlib-obf (MIT License, Copyright 2026 telemt community)
 */

#include "IptController.h"

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

namespace stealth {

// ---------------------------------------------------------------------------
// Internal constants
// ---------------------------------------------------------------------------

static constexpr double kDenom      = 4294967296.0;  // 2^32
static constexpr double kEpsilon    = 1e-9;
static constexpr double kTwoPi      = 6.28318530717958647692;
static constexpr double kMaxExpIn   =  709.78271289338397;
static constexpr double kMinExpIn   = -708.39641853226408;

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

static double safe_exp(double v) {
    if (!std::isfinite(v)) return std::numeric_limits<double>::min();
    v = std::max(kMinExpIn, std::min(kMaxExpIn, v));
    double r = std::exp(v);
    return (r > 0.0 && std::isfinite(r)) ? r : std::numeric_limits<double>::min();
}

static uint64_t to_us(double ms) {
    if (!(ms > 0.0)) return 0;
    long double us = static_cast<long double>(ms) * 1000.0L;
    constexpr long double kMax = static_cast<long double>(
        std::numeric_limits<uint64_t>::max());
    if (!(us < kMax)) return std::numeric_limits<uint64_t>::max();
    auto v = static_cast<uint64_t>(us);
    return v == 0 ? 1 : v;
}

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

IptController::IptController(const IptParams &params) : params_(params) {}

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

uint64_t IptController::next_delay_us(bool has_pending_data, TrafficHint hint) {
    if (hint == TrafficHint::Unknown) hint = TrafficHint::Interactive;
    if (is_bypass(hint)) {
        state_ = State::Idle;
        return 0;
    }

    state_ = transition(has_pending_data);
    if (state_ == State::Burst) {
        double ms = std::min(
            sample_lognormal(params_.burst_mu_ms, params_.burst_sigma),
            params_.burst_max_ms);
        return to_us(ms);
    } else {
        // Idle state: only add delay when there IS pending data
        // (otherwise we'd stall indefinitely with nothing to send)
        if (!has_pending_data) return 0;
        double ms = sample_pareto(sample_uniform01(),
                                  params_.idle_alpha,
                                  params_.idle_scale_ms,
                                  params_.idle_max_ms);
        return to_us(ms);
    }
}

uint64_t IptController::sample_idle_delay_us() {
    double ms = sample_pareto(sample_uniform01(),
                              params_.idle_alpha,
                              params_.idle_scale_ms,
                              params_.idle_max_ms);
    return to_us(ms);
}

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

bool IptController::is_bypass(TrafficHint hint) {
    return hint == TrafficHint::Keepalive
        || hint == TrafficHint::BulkData
        || hint == TrafficHint::AuthHandshake;
}

IptController::State IptController::transition(bool has_pending_data) {
    if (!has_pending_data) {
        state_ = State::Idle;
        return state_;
    }
    double u = sample_uniform01();
    if (state_ == State::Burst) {
        state_ = (u < params_.p_burst_stay) ? State::Burst : State::Idle;
    } else {
        state_ = (u < params_.p_idle_to_burst) ? State::Burst : State::Idle;
    }
    return state_;
}

double IptController::sample_uniform01() {
    uint32_t v;
    RAND_bytes(reinterpret_cast<uint8_t *>(&v), 4);
    double u = static_cast<double>(v) / kDenom;
    return std::max(kEpsilon, std::min(1.0 - kEpsilon, u));
}

double IptController::sample_normal() {
    if (has_spare_normal_) {
        has_spare_normal_ = false;
        return spare_normal_;
    }
    // Box-Muller transform
    double u1 = sample_uniform01();
    double u2 = sample_uniform01();
    double r   = std::sqrt(-2.0 * std::log(u1));
    double th  = kTwoPi * u2;
    spare_normal_     = r * std::sin(th);
    has_spare_normal_ = true;
    return r * std::cos(th);
}

double IptController::sample_lognormal(double mu, double sigma) {
    if (sigma == 0.0) return safe_exp(mu);
    return safe_exp(mu + sigma * sample_normal());
}

double IptController::sample_pareto(double u, double alpha,
                                     double scale, double max_val) const {
    if (scale >= max_val) return max_val;
    double support = 1.0 - std::pow(scale / max_val, alpha);
    double su = std::max(0.0, std::min(1.0 - kEpsilon, u * support));
    double v   = scale / std::pow(1.0 - su, 1.0 / alpha);
    if (v >= max_val) v = std::nextafter(max_val, 0.0);
    return v;
}

}  // namespace stealth