/*
 * DrsEngine.cpp — Dynamic Record Sizing for tgnet
 * Ported from telemt/tdlib-obf (MIT License, Copyright 2026 telemt community)
 */

#include "DrsEngine.h"

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

namespace stealth {

// ---------------------------------------------------------------------------
// Default policy — capture-aligned bins from tdlib-obf StealthRecordSizeBaselines.h
// ---------------------------------------------------------------------------

DrsPolicy default_drs_policy() {
    DrsPolicy p;
    p.min_payload_cap    = 80;
    p.max_payload_cap    = 16408;
    p.slow_start_records = 8;
    p.congestion_bytes   = 32768;
    p.idle_reset_ms_min  = 250;
    p.idle_reset_ms_max  = 1200;

    // SlowStart: small records matching initial browser GET bursts
    p.slow_start.bins = {
        {80,  494,  4},
        {495, 1255, 2},
        {1256,2941, 1},
    };
    p.slow_start.local_jitter   = 8;
    p.slow_start.max_repeat_run = 2;

    // CongestionOpen: medium records
    p.congestion_open.bins = {
        {200,  494,  2},
        {495,  1255, 3},
        {1256, 2941, 3},
        {2942, 5394, 2},
        {5395, 8192, 1},
    };
    p.congestion_open.local_jitter   = 16;
    p.congestion_open.max_repeat_run = 3;

    // SteadyState: full browser range
    p.steady_state.bins = {
        {200,  494,  2},
        {495,  1255, 3},
        {1256, 2941, 3},
        {2942, 5394, 2},
        {5395, 8192, 1},
    };
    p.steady_state.local_jitter   = 32;
    p.steady_state.max_repeat_run = 3;

    return p;
}

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

uint32_t DrsEngine::rand_bounded(uint32_t n) {
    if (n <= 1) return 0;
    uint32_t threshold = static_cast<uint32_t>(-n) % n;
    for (;;) {
        uint32_t v;
        RAND_bytes(reinterpret_cast<uint8_t *>(&v), 4);
        if (v >= threshold) return v % n;
    }
}

static int8_t direction_of(int32_t from, int32_t to) {
    if (to > from) return  1;
    if (to < from) return -1;
    return 0;
}

static int64_t transition_lower(int32_t anchor) {
    return static_cast<int64_t>(anchor) / 3 + 1;
}

static int64_t transition_upper(int32_t anchor) {
    return static_cast<int64_t>(anchor) * 3 - 1;
}

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

DrsEngine::DrsEngine(const DrsPolicy &policy) : policy_(policy) {
    uint32_t width = static_cast<uint32_t>(
        policy_.idle_reset_ms_max - policy_.idle_reset_ms_min + 1);
    sampled_idle_reset_ms_ =
        policy_.idle_reset_ms_min + static_cast<int32_t>(rand_bounded(width));
}

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

int32_t DrsEngine::next_payload_cap(TrafficHint hint) {
    if (hint == TrafficHint::Unknown) hint = TrafficHint::Interactive;

    if (hint == TrafficHint::Keepalive || hint == TrafficHint::AuthHandshake) {
        return policy_.min_payload_cap;
    }

    if (hint == TrafficHint::BulkData) {
        // Bulk uses steady_state bins but isolated run-state so it doesn't
        // corrupt interactive phase transition anchors
        auto save_anchor    = transition_anchor_cap_;
        auto save_prev      = previous_cap_;
        auto save_last      = last_cap_;
        auto save_run       = last_cap_run_;
        auto save_mono      = monotonic_run_;
        auto save_dir       = last_direction_;

        transition_anchor_cap_ = -1;
        previous_cap_   = bulk_previous_cap_;
        last_cap_       = bulk_last_cap_;
        last_cap_run_   = bulk_last_cap_run_;
        monotonic_run_  = bulk_monotonic_run_;
        last_direction_ = bulk_last_direction_;

        auto sampled = sample_from_phase(policy_.steady_state);

        bulk_previous_cap_   = previous_cap_;
        bulk_last_cap_       = last_cap_;
        bulk_last_cap_run_   = last_cap_run_;
        bulk_monotonic_run_  = monotonic_run_;
        bulk_last_direction_ = last_direction_;

        transition_anchor_cap_ = save_anchor;
        previous_cap_   = save_prev;
        last_cap_       = save_last;
        last_cap_run_   = save_run;
        monotonic_run_  = save_mono;
        last_direction_ = save_dir;

        return sampled;
    }

    return sample_from_phase(phase_model());
}

void DrsEngine::notify_bytes_written(size_t bytes) {
    if (bytes == 0) return;
    if (records_in_phase_ < std::numeric_limits<size_t>::max())
        records_in_phase_++;
    auto max_sz = std::numeric_limits<size_t>::max();
    bytes_in_phase_ = (bytes_in_phase_ > max_sz - bytes)
                    ? max_sz
                    : bytes_in_phase_ + bytes;
    maybe_advance_phase();
}

void DrsEngine::notify_idle() {
    phase_            = Phase::SlowStart;
    records_in_phase_ = 0;
    bytes_in_phase_   = 0;
    transition_anchor_cap_ = -1;
    reset_run_state();
    bulk_previous_cap_   = -1;
    bulk_last_cap_       = -1;
    bulk_last_cap_run_   = 0;
    bulk_monotonic_run_  = 0;
    bulk_last_direction_ = 0;
}

bool DrsEngine::should_reset_after_idle(int64_t idle_ms) const noexcept {
    return idle_ms >= static_cast<int64_t>(sampled_idle_reset_ms_);
}

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

const DrsPhaseModel &DrsEngine::phase_model() const noexcept {
    switch (phase_) {
        case Phase::SlowStart:     return policy_.slow_start;
        case Phase::CongestionOpen:return policy_.congestion_open;
        case Phase::SteadyState:   return policy_.steady_state;
    }
    return policy_.steady_state;
}

int32_t DrsEngine::sample_from_phase(const DrsPhaseModel &model) {
    constexpr int kAttempts = 32;
    auto sampled    = sample_weighted_bin_value(model);
    auto best_score = score_candidate(model, sampled);

    for (int i = 1; i < kAttempts && best_score > 0; i++) {
        auto c = sample_weighted_bin_value(model);
        auto s = score_candidate(model, c);
        if (s < best_score) { sampled = c; best_score = s; }
    }

    // If still penalised, scan bin boundaries for an escape
    if (best_score > 0) {
        auto try_candidate = [&](int32_t c) -> bool {
            auto s = score_candidate(model, c);
            if (s < best_score) { sampled = c; best_score = s; }
            return best_score == 0;
        };
        for (const auto &bin : model.bins) {
            auto lo = std::max(bin.lo, policy_.min_payload_cap);
            auto hi = std::min(bin.hi, policy_.max_payload_cap);
            if (lo > hi) continue;
            if (try_candidate(lo) || try_candidate(hi)) break;
            if (last_cap_ >= 0) {
                auto anch = std::max(lo, std::min(last_cap_, hi));
                if (try_candidate(anch)) break;
                if (anch > lo && try_candidate(anch - 1)) break;
                if (anch < hi && try_candidate(anch + 1)) break;
            }
        }
    }

    sampled = smooth_transition(sampled);
    note_selected_cap(sampled);
    return sampled;
}

int32_t DrsEngine::sample_weighted_bin_value(const DrsPhaseModel &model) {
    uint32_t total = 0;
    for (const auto &b : model.bins) total += b.weight;

    uint32_t sel = rand_bounded(total);
    const RecordSizeBin *chosen = &model.bins.back();
    for (const auto &b : model.bins) {
        if (sel < b.weight) { chosen = &b; break; }
        sel -= b.weight;
    }

    uint32_t width = static_cast<uint32_t>(chosen->hi - chosen->lo + 1);
    int32_t v = chosen->lo + static_cast<int32_t>(rand_bounded(width));
    if (model.local_jitter > 0) {
        uint32_t jw = static_cast<uint32_t>(model.local_jitter * 2 + 1);
        v += static_cast<int32_t>(rand_bounded(jw)) - model.local_jitter;
    }
    v = std::max(chosen->lo, std::min(v, chosen->hi));
    v = std::max(policy_.min_payload_cap, std::min(v, policy_.max_payload_cap));
    return v;
}

int32_t DrsEngine::score_candidate(const DrsPhaseModel &model, int32_t c) const noexcept {
    int32_t score = 0;
    if (model.max_repeat_run > 0 && last_cap_ >= 0 &&
        c == last_cap_ && last_cap_run_ >= model.max_repeat_run)
        score += 10000;

    if (last_cap_ >= 0) {
        auto dir = direction_of(last_cap_, c);
        if (dir != 0 && dir == last_direction_ && monotonic_run_ >= 2)
            score += 2000 * monotonic_run_;
    }
    if (transition_anchor_cap_ >= 0) {
        auto lo64 = transition_lower(transition_anchor_cap_);
        auto hi64 = transition_upper(transition_anchor_cap_);
        auto c64  = static_cast<int64_t>(c);
        if (c64 > hi64) score += static_cast<int32_t>(c64 - hi64);
        else if (c64 < lo64) score += static_cast<int32_t>(lo64 - c64);
    }
    return score;
}

int32_t DrsEngine::smooth_transition(int32_t candidate) noexcept {
    if (transition_anchor_cap_ < 0) return candidate;
    auto anchor = transition_anchor_cap_;
    transition_anchor_cap_ = -1;
    auto lo = transition_lower(anchor);
    auto hi = transition_upper(anchor);
    auto c  = static_cast<int64_t>(candidate);
    if (c >= lo && c <= hi) return candidate;
    auto smoothed = static_cast<int64_t>(anchor) + (c - static_cast<int64_t>(anchor)) / 2;
    auto bounded  = std::max(lo, std::min(smoothed, hi));
    return static_cast<int32_t>(
        std::max<int64_t>(policy_.min_payload_cap,
        std::min<int64_t>(bounded, policy_.max_payload_cap)));
}

void DrsEngine::maybe_advance_phase() {
    if (phase_ == Phase::SlowStart &&
        records_in_phase_ >= static_cast<size_t>(policy_.slow_start_records)) {
        transition_anchor_cap_ = last_cap_;
        phase_ = Phase::CongestionOpen;
        records_in_phase_ = bytes_in_phase_ = 0;
        reset_run_state();
    } else if (phase_ == Phase::CongestionOpen &&
               bytes_in_phase_ >= static_cast<size_t>(policy_.congestion_bytes)) {
        transition_anchor_cap_ = last_cap_;
        phase_ = Phase::SteadyState;
        records_in_phase_ = bytes_in_phase_ = 0;
        reset_run_state();
    }
}

void DrsEngine::note_selected_cap(int32_t cap) noexcept {
    auto dir = last_cap_ >= 0 ? direction_of(last_cap_, cap) : 0;
    if (dir != 0) {
        monotonic_run_ = (dir == last_direction_) ? monotonic_run_ + 1 : 1;
        last_direction_ = dir;
    } else {
        monotonic_run_  = 0;
        last_direction_ = 0;
    }
    previous_cap_ = last_cap_;
    if (cap == last_cap_) { last_cap_run_++; }
    else                  { last_cap_ = cap; last_cap_run_ = 1; }
}

void DrsEngine::reset_run_state() noexcept {
    previous_cap_   = -1;
    last_cap_       = -1;
    last_cap_run_   = 0;
    monotonic_run_  = 0;
    last_direction_ = 0;
}

}  // namespace stealth