/*
 * DrsEngine.h — Dynamic Record Sizing for tgnet
 * Ported from telemt/tdlib-obf (MIT License, Copyright 2026 telemt community)
 * Adapted for tgnet: C++17, no tdlib dependencies.
 *
 * DRS shapes TLS record payload sizes across three phases to match
 * empirically-captured browser traffic patterns, defeating DPI classifiers
 * that identify MTProto by its characteristic fixed-size records.
 *
 * Phase progression:
 *   SlowStart     — small records (80-494 bytes), first N records after connect
 *   CongestionOpen — medium records (200-2941 bytes), until M bytes sent
 *   SteadyState   — full browser-range records (200-8192 bytes)
 *
 * After idle (default 250-1200ms) resets to SlowStart.
 */

#ifndef DRS_ENGINE_H
#define DRS_ENGINE_H

#include <cstdint>
#include <vector>

namespace stealth {

struct RecordSizeBin {
    int32_t lo;
    int32_t hi;
    uint32_t weight;
};

struct DrsPhaseModel {
    std::vector<RecordSizeBin> bins;
    int32_t local_jitter{0};   // ± jitter added to sampled value
    int32_t max_repeat_run{3}; // penalise picking same size >N times in a row
};

struct DrsPolicy {
    DrsPhaseModel slow_start;
    DrsPhaseModel congestion_open;
    DrsPhaseModel steady_state;
    int32_t min_payload_cap{80};
    int32_t max_payload_cap{16408};
    int32_t slow_start_records{8};   // records before leaving SlowStart
    int32_t congestion_bytes{32768}; // bytes before leaving CongestionOpen
    int32_t idle_reset_ms_min{250};
    int32_t idle_reset_ms_max{1200};
};

// Returns the default policy with capture-aligned bins from tdlib-obf baselines
DrsPolicy default_drs_policy();

// Traffic hint — tells DRS which sampling path to use
enum class TrafficHint : uint8_t {
    Unknown       = 0,
    Interactive   = 1,  // normal messaging — uses phase model
    BulkData      = 2,  // file upload/download — uses steady_state bins directly
    Keepalive     = 3,  // ping — uses min_payload_cap
    AuthHandshake = 4,  // key exchange — uses min_payload_cap
};

class DrsEngine {
public:
    enum class Phase : uint8_t { SlowStart = 0, CongestionOpen = 1, SteadyState = 2 };

    explicit DrsEngine(const DrsPolicy &policy);

    // Returns the max payload bytes for the next TLS record.
    // Call before each send(); pass result as the record size cap.
    int32_t next_payload_cap(TrafficHint hint = TrafficHint::Interactive);

    // Call after a record is actually sent with `bytes` payload bytes.
    void notify_bytes_written(size_t bytes);

    // Call when the connection goes idle (no sends for idle_reset_ms).
    void notify_idle();

    Phase current_phase() const noexcept { return phase_; }

    // Returns true if idle_ms exceeds the sampled idle reset threshold
    bool should_reset_after_idle(int64_t idle_ms) const noexcept;

private:
    DrsPolicy policy_;
    Phase phase_{Phase::SlowStart};
    size_t records_in_phase_{0};
    size_t bytes_in_phase_{0};
    int32_t sampled_idle_reset_ms_{0};
    int32_t transition_anchor_cap_{-1};
    int32_t previous_cap_{-1};
    int32_t last_cap_{-1};
    int32_t last_cap_run_{0};
    int32_t monotonic_run_{0};
    int8_t  last_direction_{0};

    // Bulk-isolated run state
    int32_t bulk_previous_cap_{-1};
    int32_t bulk_last_cap_{-1};
    int32_t bulk_last_cap_run_{0};
    int32_t bulk_monotonic_run_{0};
    int8_t  bulk_last_direction_{0};

    const DrsPhaseModel &phase_model() const noexcept;
    int32_t sample_from_phase(const DrsPhaseModel &model);
    int32_t sample_weighted_bin_value(const DrsPhaseModel &model);
    int32_t score_candidate(const DrsPhaseModel &model, int32_t candidate) const noexcept;
    int32_t smooth_transition(int32_t candidate) noexcept;
    void maybe_advance_phase();
    void note_selected_cap(int32_t cap) noexcept;
    void reset_run_state() noexcept;

    static uint32_t rand_bounded(uint32_t n);
};

}  // namespace stealth

#endif  // DRS_ENGINE_H