stochastic-routing-extended 1.0.2

//! Transport policy: environment-aware selection of optimal transports.
//!
//! Different network environments favor different transports:
//! - **Wifi**: low latency, prefer QUIC/UDP
//! - **Cellular**: unreliable, prefer TCP/QUIC (retransmission built-in)
//! - **Corporate**: restrictive firewalls, prefer HTTP/gRPC tunnels
//! - **Unknown**: balanced default order

use crate::transport::TransportKind;

/// Network environment hint for transport selection tuning.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NetworkEnvironment {
    Wifi,
    Cellular,
    Corporate,
    Unknown,
}

/// Per-environment transport preference (first = most preferred).
fn preference_for(env: NetworkEnvironment) -> &'static [TransportKind] {
    match env {
        NetworkEnvironment::Wifi => &[
            TransportKind::Quic,
            TransportKind::Udp,
            TransportKind::Tcp,
            TransportKind::WebSocket,
            TransportKind::Grpc,
            TransportKind::Http2,
        ],
        NetworkEnvironment::Cellular => &[
            TransportKind::Tcp,
            TransportKind::Quic,
            TransportKind::WebSocket,
            TransportKind::Http2,
            TransportKind::Grpc,
            TransportKind::Udp,
        ],
        NetworkEnvironment::Corporate => &[
            TransportKind::Http2,
            TransportKind::Grpc,
            TransportKind::WebSocket,
            TransportKind::Tcp,
            TransportKind::Quic,
            TransportKind::Udp,
        ],
        NetworkEnvironment::Unknown => &[
            TransportKind::Quic,
            TransportKind::Tcp,
            TransportKind::WebSocket,
            TransportKind::Udp,
            TransportKind::Grpc,
            TransportKind::Http2,
        ],
    }
}

/// Selects the optimal set of transports based on network environment and
/// optional health scores.
pub struct TransportPolicy {
    env: NetworkEnvironment,
}

impl TransportPolicy {
    pub fn new(env: NetworkEnvironment) -> Self {
        Self { env }
    }

    /// Current network environment.
    pub fn environment(&self) -> NetworkEnvironment {
        self.env
    }

    /// Update the detected network environment.
    pub fn set_environment(&mut self, env: NetworkEnvironment) {
        self.env = env;
    }

    /// Return `available` transports sorted by environment-specific preference;
    /// unknown kinds appended at the end.
    pub fn recommend(&self, available: &[TransportKind]) -> Vec<TransportKind> {
        let pref = preference_for(self.env);
        let mut out: Vec<TransportKind> = pref
            .iter()
            .copied()
            .filter(|k| available.contains(k))
            .collect();
        for k in available {
            if !out.contains(k) {
                out.push(*k);
            }
        }
        out
    }

    /// Return `available` transports sorted by health score (descending),
    /// with ties broken by environment-specific preference.
    ///
    /// `scores` maps each transport to a health score (0.0 = dead, 1.0 = perfect).
    /// Transports without a score entry are treated as score 0.5 (unknown).
    pub fn recommend_with_health(
        &self,
        available: &[TransportKind],
        scores: &[(TransportKind, f64)],
    ) -> Vec<TransportKind> {
        let pref = preference_for(self.env);

        let pref_rank =
            |k: &TransportKind| -> usize { pref.iter().position(|p| p == k).unwrap_or(pref.len()) };
        let score_of = |k: &TransportKind| -> f64 {
            scores
                .iter()
                .find(|(sk, _)| sk == k)
                .map(|(_, s)| *s)
                .unwrap_or(0.5)
        };

        let mut ranked: Vec<TransportKind> = available.to_vec();
        ranked.sort_by(|a, b| {
            let sa = score_of(a);
            let sb = score_of(b);
            // Higher score first; on tie, lower preference rank first.
            sb.partial_cmp(&sa)
                .unwrap_or(std::cmp::Ordering::Equal)
                .then_with(|| pref_rank(a).cmp(&pref_rank(b)))
        });
        ranked
    }
}

impl Default for TransportPolicy {
    fn default() -> Self {
        Self::new(NetworkEnvironment::Unknown)
    }
}

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

    #[test]
    fn wifi_prefers_quic_udp() {
        let p = TransportPolicy::new(NetworkEnvironment::Wifi);
        let v = p.recommend(&[TransportKind::Tcp, TransportKind::Udp, TransportKind::Quic]);
        assert_eq!(v[0], TransportKind::Quic);
        assert_eq!(v[1], TransportKind::Udp);
        assert_eq!(v[2], TransportKind::Tcp);
    }

    #[test]
    fn cellular_prefers_tcp() {
        let p = TransportPolicy::new(NetworkEnvironment::Cellular);
        let v = p.recommend(&[TransportKind::Udp, TransportKind::Tcp, TransportKind::Quic]);
        assert_eq!(v[0], TransportKind::Tcp);
    }

    #[test]
    fn corporate_prefers_http_grpc() {
        let p = TransportPolicy::new(NetworkEnvironment::Corporate);
        let v = p.recommend(&[
            TransportKind::Tcp,
            TransportKind::Http2,
            TransportKind::Grpc,
        ]);
        assert_eq!(v[0], TransportKind::Http2);
        assert_eq!(v[1], TransportKind::Grpc);
        assert_eq!(v[2], TransportKind::Tcp);
    }

    #[test]
    fn unknown_matches_original_order() {
        let p = TransportPolicy::new(NetworkEnvironment::Unknown);
        let v = p.recommend(&[TransportKind::Udp, TransportKind::Tcp, TransportKind::Quic]);
        assert_eq!(
            v,
            vec![TransportKind::Quic, TransportKind::Tcp, TransportKind::Udp]
        );
    }

    #[test]
    fn health_scores_override_preference() {
        let p = TransportPolicy::new(NetworkEnvironment::Wifi);
        // UDP has a higher health score than QUIC despite Wifi preferring QUIC.
        let scores = vec![
            (TransportKind::Quic, 0.3),
            (TransportKind::Udp, 0.9),
            (TransportKind::Tcp, 0.5),
        ];
        let v = p.recommend_with_health(
            &[TransportKind::Quic, TransportKind::Udp, TransportKind::Tcp],
            &scores,
        );
        assert_eq!(v[0], TransportKind::Udp);
        assert_eq!(v[1], TransportKind::Tcp);
        assert_eq!(v[2], TransportKind::Quic);
    }

    #[test]
    fn equal_scores_break_tie_by_env_preference() {
        let p = TransportPolicy::new(NetworkEnvironment::Wifi);
        let scores = vec![(TransportKind::Quic, 0.8), (TransportKind::Udp, 0.8)];
        let v = p.recommend_with_health(&[TransportKind::Udp, TransportKind::Quic], &scores);
        // Wifi prefers QUIC over UDP.
        assert_eq!(v[0], TransportKind::Quic);
        assert_eq!(v[1], TransportKind::Udp);
    }
}