//! Traffic-shaping padding helpers (project §7.2).
//!
//! Aura's outer wire is QUIC dressed up as HTTP/3 (see [`crate::mimicry`]). Real HTTPS/H3 traffic
//! tends to cluster at a handful of record/datagram sizes; padding an Aura payload up to one of
//! those buckets makes a passive size-only classifier less able to single it out. These helpers are
//! deliberately simple, allocation-light, and fully unit-tested so the higher layers can call them
//! with confidence.
//!
//! Note these operate on the *application* payload before it is handed to the proto/QUIC layers.
//! They do not, and cannot, hide QUIC's own framing — they only normalize plaintext lengths so the
//! ciphertext lands in a common size class.

use rand::Rng;

/// The HTTPS/H3-like size buckets a payload is rounded up to, in ascending order.
///
/// The first five (64 / 128 / 256 / 512 / 1024) are common small-record sizes; 1280 is the IPv6
/// minimum-MTU "safe" QUIC datagram size; 1460 is a typical Ethernet TCP/QUIC payload (1500-byte
/// MTU minus IP+UDP headers). Keep this sorted ascending: [`pad_to_https_size`] relies on it.
///
/// This is **profile 0** of [`PADDING_PROFILES`]; the standalone `const` is preserved as a
/// backwards-compatible alias so existing transport tests and direct callers keep working.
pub const HTTPS_SIZE_BUCKETS: [usize; 7] = [64, 128, 256, 512, 1024, 1280, 1460];

/// The largest bucket in [`HTTPS_SIZE_BUCKETS`]; payloads at or above this are left unpadded.
pub const MAX_BUCKET: usize = HTTPS_SIZE_BUCKETS[HTTPS_SIZE_BUCKETS.len() - 1];

/// Padding-profile palettes used by the daily mask rotation
/// ([`aura_crypto::MaskSet::padding_profile_id`]).
///
/// Each profile is a strictly-ascending list of size buckets; [`pad_to_bucket`] rounds a packet up
/// to the smallest bucket of the chosen profile that is `>= len`, and leaves packets at/over the
/// profile's maximum unchanged (same semantics as the original [`pad_to_https_size`]).
///
/// * Profile **0** = [`HTTPS_SIZE_BUCKETS`] — the historical default; identical wire behaviour to
///   pre-rotation so legacy tests pass.
/// * Profiles **1..3** — alternative bucket shapes drawn from common MTU/TLS-record sizes; their
///   total bandwidth overhead is similar but the on-wire size distribution is different.
///
/// The mask layer expects exactly [`aura_crypto::PADDING_PROFILE_COUNT`] entries here; the
/// transport reduces a profile id modulo this length defensively so an out-of-range id is silently
/// reinterpreted instead of panicking.
pub const PADDING_PROFILES: &[&[usize]] = &[
    // 0: backwards-compatible default (same as HTTPS_SIZE_BUCKETS).
    &[64, 128, 256, 512, 1024, 1280, 1460],
    // 1: small-record + larger-MTU profile (rounds the largest payloads to a full 1500 frame).
    &[128, 256, 512, 1024, 1280, 1500],
    // 2: coarser buckets, deliberately fewer steps to obscure fine length variations.
    &[200, 400, 800, 1200, 1500],
    // 3: TLS-record-ish (256 / 512) plus three larger buckets.
    &[256, 512, 768, 1024, 1280, 1500],
];

/// Pad `packet` (in place, appending zero bytes) up to the next HTTPS-like size bucket.
///
/// Behavior:
/// * If `packet.len()` is already exactly a bucket, it is left unchanged (idempotent).
/// * Otherwise it grows to the smallest bucket strictly larger than its current length.
/// * **At or over the largest bucket ([`MAX_BUCKET`] = 1460):** the packet is left unchanged. We do
///   not pad to a multiple of 1460, because a single Aura payload is expected to fit within one
///   datagram; over-MTU payloads are the caller's concern (e.g. they will be split by QUIC anyway),
///   and rounding them up would only waste bandwidth without improving the size-class disguise.
///
/// Padding is appended as zero bytes; this is a length-shaping primitive, not an authenticated
/// framing scheme — the proto layer seals the whole (already-padded) payload with AEAD, so the pad
/// bytes are encrypted on the wire. Callers that need to *recover* the original length must carry it
/// themselves (e.g. an inner length prefix); for Aura's IP-packet payloads the IP total-length field
/// already bounds the real data, so trailing zeros are simply ignored by the receiver.
pub fn pad_to_https_size(packet: &mut Vec<u8>) {
    let len = packet.len();
    // Smallest bucket that can already hold `len` (>=, so an exact-bucket length is a no-op and the
    // operation is idempotent). If `len` exceeds every bucket, no padding is applied.
    if let Some(&target) = HTTPS_SIZE_BUCKETS.iter().find(|&&b| b >= len) {
        packet.resize(target, 0);
    }
}

/// Return the bucket `len` would be padded *up to* by [`pad_to_https_size`], or `len` itself if it
/// is at/over [`MAX_BUCKET`]. Useful for sizing buffers ahead of time and for tests.
#[must_use]
pub fn next_https_bucket(len: usize) -> usize {
    HTTPS_SIZE_BUCKETS
        .iter()
        .copied()
        .find(|&b| b >= len)
        .unwrap_or(len)
}

/// Pad `packet` up to the next bucket of [`PADDING_PROFILES`]`[profile_id]` (same semantics as
/// [`pad_to_https_size`], but selecting the bucket palette by profile id).
///
/// `profile_id` is reduced modulo `PADDING_PROFILES.len()` so an out-of-range id silently maps
/// back into the valid range instead of panicking (the higher mask layer derives it from HKDF and
/// has its own range guard, but the transport-layer defence here keeps misuse non-fatal).
///
/// Profile `0` reproduces the original [`pad_to_https_size`] palette exactly, so a deployment that
/// happens to derive profile 0 is byte-identical to the pre-rotation wire behaviour.
pub fn pad_to_bucket(packet: &mut Vec<u8>, profile_id: u8) {
    let buckets = profile_buckets(profile_id);
    let len = packet.len();
    if let Some(&target) = buckets.iter().find(|&&b| b >= len) {
        packet.resize(target, 0);
    }
}

/// Return the bucket `len` would be padded *up to* by [`pad_to_bucket`] for the given
/// `profile_id`, or `len` itself if it is at/over the profile's maximum bucket.
#[must_use]
pub fn next_bucket_for_profile(len: usize, profile_id: u8) -> usize {
    let buckets = profile_buckets(profile_id);
    buckets.iter().copied().find(|&b| b >= len).unwrap_or(len)
}

/// Look up the bucket palette for `profile_id`, reducing the id modulo the palette count so an
/// out-of-range id is harmless.
fn profile_buckets(profile_id: u8) -> &'static [usize] {
    let idx = (profile_id as usize) % PADDING_PROFILES.len();
    PADDING_PROFILES[idx]
}

/// Best-effort random padding: append between `min_pad` and `max_pad` (inclusive) zero bytes to
/// `packet`, capping the result at `max_total` bytes so a hard size ceiling is never exceeded.
///
/// Signature rationale: callers want jitter on the wire without blowing a datagram budget, so the
/// knobs are an inclusive `[min_pad, max_pad]` range plus an absolute `max_total` clamp. Returns the
/// number of pad bytes actually appended (which may be **less** than `min_pad` if `max_total` was
/// already nearly reached, including `0` when `packet.len() >= max_total`).
///
/// "Best-effort" = if the requested padding does not fit under `max_total`, as much as fits is added
/// rather than erroring. If `min_pad > max_pad` the arguments are swapped so the call still does
/// something sensible instead of panicking.
pub fn inject_padding_frames(
    packet: &mut Vec<u8>,
    min_pad: usize,
    max_pad: usize,
    max_total: usize,
) -> usize {
    let (lo, hi) = if min_pad <= max_pad {
        (min_pad, max_pad)
    } else {
        (max_pad, min_pad)
    };

    let headroom = max_total.saturating_sub(packet.len());
    if headroom == 0 {
        return 0;
    }

    // Pick a random amount in [lo, hi], then clamp to whatever headroom remains.
    let want = if lo == hi {
        lo
    } else {
        rand::thread_rng().gen_range(lo..=hi)
    };
    let pad = want.min(headroom);
    packet.resize(packet.len() + pad, 0);
    pad
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn pads_up_to_each_bucket() {
        // One below / one into each bucket lands on that bucket.
        let cases = [
            (0usize, 64usize),
            (1, 64),
            (63, 64),
            (65, 128),
            (127, 128),
            (200, 256),
            (300, 512),
            (513, 1024),
            (1025, 1280),
            (1281, 1460),
        ];
        for (input, expected) in cases {
            let mut v = vec![0xAB; input];
            pad_to_https_size(&mut v);
            assert_eq!(v.len(), expected, "padding {input} should reach {expected}");
        }
    }

    #[test]
    fn exact_bucket_is_unchanged_and_idempotent() {
        for &b in &HTTPS_SIZE_BUCKETS {
            let mut v = vec![1u8; b];
            pad_to_https_size(&mut v);
            assert_eq!(v.len(), b, "exact bucket {b} must not grow");
            // Idempotence: padding again changes nothing.
            pad_to_https_size(&mut v);
            assert_eq!(v.len(), b, "re-padding bucket {b} must be a no-op");
        }
    }

    #[test]
    fn at_or_over_max_bucket_is_left_alone() {
        for len in [MAX_BUCKET, MAX_BUCKET + 1, 2000, 9000] {
            let mut v = vec![7u8; len];
            pad_to_https_size(&mut v);
            assert_eq!(v.len(), len, "len {len} >= MAX_BUCKET must be unchanged");
        }
    }

    #[test]
    fn padding_preserves_original_prefix() {
        let mut v: Vec<u8> = (0..50u8).collect();
        let original = v.clone();
        pad_to_https_size(&mut v);
        assert_eq!(v.len(), 64);
        assert_eq!(&v[..50], &original[..], "real bytes must be preserved");
        assert!(v[50..].iter().all(|&b| b == 0), "pad must be zero bytes");
    }

    #[test]
    fn next_bucket_matches_padding() {
        for len in [0usize, 1, 64, 65, 1024, 1459, 1460, 5000] {
            let predicted = next_https_bucket(len);
            let mut v = vec![0u8; len];
            pad_to_https_size(&mut v);
            assert_eq!(predicted, v.len(), "next_https_bucket disagrees at {len}");
        }
    }

    #[test]
    fn inject_padding_respects_range_and_cap() {
        // Plenty of headroom: result grows by something in [4, 8].
        let mut v = vec![0u8; 10];
        let added = inject_padding_frames(&mut v, 4, 8, 1000);
        assert!((4..=8).contains(&added), "added {added} outside [4,8]");
        assert_eq!(v.len(), 10 + added);
    }

    #[test]
    fn inject_padding_clamps_to_max_total() {
        // Only 3 bytes of headroom even though we ask for 10..=10.
        let mut v = vec![0u8; 97];
        let added = inject_padding_frames(&mut v, 10, 10, 100);
        assert_eq!(added, 3, "should add only what fits under max_total");
        assert_eq!(v.len(), 100);
    }

    #[test]
    fn inject_padding_zero_headroom_is_noop() {
        let mut v = vec![0u8; 100];
        let added = inject_padding_frames(&mut v, 1, 50, 100);
        assert_eq!(added, 0);
        assert_eq!(v.len(), 100);

        // Already over the cap: still a no-op, never truncates.
        let mut v2 = vec![0u8; 120];
        let added2 = inject_padding_frames(&mut v2, 1, 50, 100);
        assert_eq!(added2, 0);
        assert_eq!(v2.len(), 120);
    }

    #[test]
    fn inject_padding_swapped_bounds_dont_panic() {
        let mut v = vec![0u8; 10];
        let added = inject_padding_frames(&mut v, 8, 4, 1000); // min > max
        assert!((4..=8).contains(&added));
    }

    /// Profile 0 must exactly reproduce the historical [`pad_to_https_size`] behaviour, so any
    /// deployment that derives profile 0 is wire-compatible with pre-rotation Aura.
    #[test]
    fn profile_zero_matches_pad_to_https_size() {
        for len in [0usize, 1, 63, 64, 65, 200, 1280, 1459, 1460, 1461, 9999] {
            let mut a = vec![0u8; len];
            let mut b = vec![0u8; len];
            pad_to_https_size(&mut a);
            pad_to_bucket(&mut b, 0);
            assert_eq!(
                a.len(),
                b.len(),
                "profile 0 must match the legacy palette at {len}"
            );
        }
    }

    /// Every profile must be strictly ascending and pad correctly to its first-`>= len` bucket
    /// (matching `next_bucket_for_profile` and leaving overlarge payloads unchanged).
    #[test]
    fn pad_to_bucket_respects_each_profile() {
        for (pid, palette) in PADDING_PROFILES.iter().enumerate() {
            // Ascending?
            for w in palette.windows(2) {
                assert!(w[0] < w[1], "profile {pid} buckets must be ascending");
            }
            let max = *palette.last().unwrap();
            for &target_len in *palette {
                // Just-under-bucket lands on bucket; exactly-bucket is unchanged.
                let mut v = vec![0u8; target_len.saturating_sub(1)];
                pad_to_bucket(&mut v, pid as u8);
                assert_eq!(v.len(), target_len);
                let mut v2 = vec![0u8; target_len];
                pad_to_bucket(&mut v2, pid as u8);
                assert_eq!(
                    v2.len(),
                    target_len,
                    "exact bucket {target_len} unchanged in profile {pid}"
                );
            }
            // Over the max stays unchanged.
            let mut v = vec![0u8; max + 1];
            pad_to_bucket(&mut v, pid as u8);
            assert_eq!(v.len(), max + 1, "over-max unchanged in profile {pid}");
        }
    }

    /// Out-of-range `profile_id` is reduced modulo the palette count instead of panicking.
    #[test]
    fn pad_to_bucket_handles_out_of_range_id() {
        let mut a = vec![0u8; 100];
        let mut b = vec![0u8; 100];
        // 0 == PADDING_PROFILES.len() * k for any k; e.g. 8 wraps to 0 for a 4-entry palette.
        pad_to_bucket(&mut a, 0);
        pad_to_bucket(&mut b, PADDING_PROFILES.len() as u8 * 2);
        assert_eq!(a.len(), b.len(), "id wraps modulo palette count");
        let predicted = next_bucket_for_profile(100, 0);
        assert_eq!(a.len(), predicted);
    }
}