use anyhow::Result;
use bytes::Bytes;
use ostp_core::{OstpEvent, ProtocolAction, ProtocolConfig, ProtocolMachine};
use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::{Arc, RwLock};
use std::sync::atomic::Ordering;
use portable_atomic::AtomicU64;

/// Maximum number of concurrent authenticated sessions.
/// Excess handshake attempts are silently dropped -- no response, no state allocated.
const MAX_SESSIONS: usize = 1024;

/// Cap on the anti-replay handshake cache. When reached, expired entries are
/// reclaimed (and if needed the oldest is evicted) rather than rejecting new
/// handshakes globally — see the eviction logic in on_datagram.
const REPLAY_CACHE_MAX: usize = 50_000;

pub enum DispatchOutcome {
    Unauthorized,
    /// Packet matched a registered key's per-key junk marker — drop silently.
    Junk,
    Accepted {
        responses: Vec<Bytes>,
        app_payloads: Vec<(u32, u16, Bytes)>, // session_id, stream_id, payload
        peer_addr: SocketAddr,
    },
}

/// Per-user traffic statistics.
pub struct UserStats {
    pub bytes_up: AtomicU64,
    pub bytes_down: AtomicU64,
    pub connections: AtomicU64,
    pub limit_bytes: Option<u64>,
    pub created_at: std::time::SystemTime,
}

impl UserStats {
    pub fn new(limit: Option<u64>) -> Self {
        Self {
            bytes_up: AtomicU64::new(0),
            bytes_down: AtomicU64::new(0),
            connections: AtomicU64::new(0),
            limit_bytes: limit,
            created_at: std::time::SystemTime::now(),
        }
    }

    pub fn is_over_limit(&self) -> bool {
        if let Some(limit) = self.limit_bytes {
            let total = self.bytes_up.load(Ordering::Relaxed)
                + self.bytes_down.load(Ordering::Relaxed);
            total >= limit
        } else {
            false
        }
    }
}

/// Snapshot of user stats for API responses.
#[derive(Debug, Clone, serde::Serialize)]
pub struct UserStatsSnapshot {
    pub access_key: String,
    pub name: Option<String>,
    pub bytes_up: u64,
    pub bytes_down: u64,
    pub connections: u64,
    pub limit_bytes: Option<u64>,
    pub online: bool,
    pub last_seen: Option<u64>,
}

pub struct PeerState {
    pub machine: ProtocolMachine,
    pub last_addr: SocketAddr,
    pub obfuscation_key: [u8; 8],
    pub last_seen: std::time::Instant,
    pub access_key: String,
}

pub struct Dispatcher {
    peer_machines: HashMap<u32, PeerState>,
    addr_to_session: HashMap<SocketAddr, u32>,
    machine_config: ProtocolConfig,
    access_keys: Arc<RwLock<HashMap<String, crate::api::UserMeta>>>,
    user_stats: Arc<RwLock<HashMap<String, Arc<UserStats>>>>,
    replay_cache: std::collections::HashMap<Vec<u8>, u64>,
    roaming_tokens: f64,
    last_token_regen: std::time::Instant,
    /// Cache of per-key derived secrets (obf key / psk / padding). These are a
    /// pure function of the access key + PROTOCOL_VERSION, so they never change
    /// for a given key — computing the HKDF on every unknown datagram, for every
    /// registered key, was pure waste and an attacker-amplified CPU sink.
    secrets_cache: HashMap<String, ostp_core::crypto::DerivedSecrets>,
    /// Cache of each key's junk markers for the current time window. The marker
    /// rotates every window, so the cached `(window, m_now, m_prev)` is refreshed
    /// when the window rolls; within a window it's an HMAC we compute once, not
    /// twice per key per packet.
    junk_cache: HashMap<String, (u64, [u8; 4], [u8; 4])>,
    /// Token bucket bounding how many expensive new-handshake key-trials we run
    /// per second. The existing-session fast path and roaming path are NOT gated
    /// by this; only the O(N_keys) trial over unknown datagrams is, so a garbage
    /// flood from spoofed sources can't force unbounded per-packet crypto work.
    trial_tokens: f64,
    last_trial_regen: std::time::Instant,
}

/// Sustained rate (and burst ceiling) of new-handshake trials per second. Legit
/// first-connect packets are rare, so this is generous for real use while still
/// capping flood-driven trial work at TRIAL_RATE × num_keys crypto ops/sec.
const TRIAL_RATE: f64 = 100.0;

/// Short, non-reversible fingerprint of an access key for logs. The access key
/// is a shared secret, so it must never be written to logs verbatim; this lets
/// an operator correlate events without exposing the key itself.
pub(crate) fn key_fp(access_key: &str) -> String {
    use sha2::{Digest, Sha256};
    let h = Sha256::digest(access_key.as_bytes());
    format!("{:02x}{:02x}{:02x}", h[0], h[1], h[2])
}

#[allow(dead_code)]
impl Dispatcher {
    pub fn new(machine_config: ProtocolConfig, access_keys: Arc<RwLock<HashMap<String, crate::api::UserMeta>>>) -> Self {
        let mut initial_stats = HashMap::new();
        for (key, meta) in access_keys.read().unwrap_or_else(|e| e.into_inner()).iter() {
            initial_stats.insert(key.clone(), Arc::new(UserStats::new(meta.limit_bytes)));
        }
        Self {
            peer_machines: HashMap::new(),
            addr_to_session: HashMap::new(),
            machine_config,
            access_keys,
            user_stats: Arc::new(RwLock::new(initial_stats)),
            replay_cache: std::collections::HashMap::new(),
            roaming_tokens: 50.0,
            last_token_regen: std::time::Instant::now(),
            secrets_cache: HashMap::new(),
            junk_cache: HashMap::new(),
            trial_tokens: TRIAL_RATE,
            last_trial_regen: std::time::Instant::now(),
        }
    }

    /// Fetch this key's derived secrets from cache, computing (and caching) them
    /// on first sight. Pure function of the key, so the entry never goes stale.
    fn cached_secrets(&mut self, key: &str) -> ostp_core::crypto::DerivedSecrets {
        if let Some(s) = self.secrets_cache.get(key) {
            return s.clone();
        }
        let s = ostp_core::crypto::derive_all_secrets(key.as_bytes());
        self.secrets_cache.insert(key.to_string(), s.clone());
        s
    }

    /// Fetch this key's `(m_now, m_prev)` junk markers for `window`, recomputing
    /// only when the cached window has rolled.
    fn cached_junk_markers(&mut self, key: &str, window: u64) -> ([u8; 4], [u8; 4]) {
        if let Some(&(w, m_now, m_prev)) = self.junk_cache.get(key) {
            if w == window {
                return (m_now, m_prev);
            }
        }
        let m_now = ostp_core::crypto::derive_junk_marker(key.as_bytes(), window);
        let m_prev = ostp_core::crypto::derive_junk_marker(key.as_bytes(), window.wrapping_sub(1));
        self.junk_cache.insert(key.to_string(), (window, m_now, m_prev));
        (m_now, m_prev)
    }

    /// Returns a shared reference to user stats for the Management API.
    pub fn user_stats_ref(&self) -> Arc<RwLock<HashMap<String, Arc<UserStats>>>> {
        self.user_stats.clone()
    }

    /// Snapshot all user stats for API responses.
    pub fn snapshot_all_users(&self) -> Vec<UserStatsSnapshot> {
        let stats = self.user_stats.read().unwrap_or_else(|e| e.into_inner());
        let mut online_keys: HashMap<String, std::time::Instant> = HashMap::new();
        for ps in self.peer_machines.values() {
            let key = ps.access_key.clone();
            if let Some(existing) = online_keys.get(&key) {
                if ps.last_seen > *existing {
                    online_keys.insert(key, ps.last_seen);
                }
            } else {
                online_keys.insert(key, ps.last_seen);
            }
        }
        
        let now = std::time::Instant::now();
        let current_sys_time = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_secs();

        stats.iter().map(|(key, us)| {
            let last_seen_unix = online_keys.get(key).map(|&instant| {
                let diff = now.duration_since(instant).as_secs();
                current_sys_time.saturating_sub(diff)
            });
            
            UserStatsSnapshot {
                access_key: key.clone(),
                name: None,
                bytes_up: us.bytes_up.load(Ordering::Relaxed),
                bytes_down: us.bytes_down.load(Ordering::Relaxed),
                connections: us.connections.load(Ordering::Relaxed),
                limit_bytes: us.limit_bytes,
                online: online_keys.contains_key(key),
                last_seen: last_seen_unix,
            }
        }).collect()
    }

    /// Get or create stats entry for a user key.
    fn get_or_create_user_stats(&self, key: &str) -> Arc<UserStats> {
        let stats = self.user_stats.read().unwrap_or_else(|e| e.into_inner());
        if let Some(existing) = stats.get(key) {
            return existing.clone();
        }
        drop(stats);
        
        let limit_bytes = self.access_keys.read().unwrap_or_else(|e| e.into_inner()).get(key).and_then(|m| m.limit_bytes);
        
        let mut stats = self.user_stats.write().unwrap_or_else(|e| e.into_inner());
        stats.entry(key.to_string())
            .or_insert_with(|| Arc::new(UserStats::new(limit_bytes)))
            .clone()
    }

    /// Set traffic limit for a user.
    pub fn set_user_limit(&self, key: &str, limit: Option<u64>) {
        let mut stats = self.user_stats.write().unwrap_or_else(|e| e.into_inner());
        let entry = stats.entry(key.to_string())
            .or_insert_with(|| Arc::new(UserStats::new(limit)));
        // Replace the entry with new limit (stats reset)
        *entry = Arc::new(UserStats {
            bytes_up: AtomicU64::new(entry.bytes_up.load(Ordering::Relaxed)),
            bytes_down: AtomicU64::new(entry.bytes_down.load(Ordering::Relaxed)),
            connections: AtomicU64::new(entry.connections.load(Ordering::Relaxed)),
            limit_bytes: limit,
            created_at: entry.created_at,
        });
    }

    /// Active session count.
    pub fn active_sessions(&self) -> usize {
        self.peer_machines.len()
    }

    /// Per-session download-direction congestion headroom, in packets:
    /// `(session_id, available)` where `available = clamped cwnd - in_flight`.
    ///
    /// Consumed by the relay's per-target-connection reader tasks (see
    /// `relay::handle_relay_message`'s Connect handler) to throttle how fast
    /// they pull bytes from the upstream target and forward them to the
    /// client's OSTP session. Without this, a fast target (e.g. a CDN) gets
    /// read and forwarded as fast as the target can serve, completely
    /// ignoring the client-facing session's real congestion window - on a
    /// lossy/jittery client path that self-inflicts a loss burst, which
    /// wrecks the RTT/RTO estimate and can stall the session hard enough to
    /// trip the client's keepalive reconnect. Same clamp(16, 16384) the
    /// client uses for its own analogous uplink gate, for symmetry.
    pub fn snapshot_backpressure(&self) -> Vec<(u32, i64)> {
        self.peer_machines
            .iter()
            .map(|(&sid, ps)| {
                // Ceiling matches MAX_CWND_PACKETS in ostp-core. The old 16384
                // allowed ~20 MB outstanding toward one client — on a mobile
                // downlink that is standing queue, not throughput, and it is the
                // download direction that carries video.
                let cwnd = (ps.machine.cwnd_packets() as i64).clamp(16, 1024);
                let in_flight = ps.machine.in_flight_count() as i64;
                // Pacing gates the RATE, cwnd only the outstanding amount. With
                // the pacing bucket empty, report no headroom so the relay
                // reader pauses instead of handing over another chunk that would
                // leave back-to-back.
                if !ps.machine.can_pace_packet() {
                    return (sid, 0);
                }
                (sid, cwnd - in_flight)
            })
            .collect()
    }

    pub fn on_datagram(&mut self, peer: SocketAddr, packet: Bytes) -> Result<DispatchOutcome> {
        if packet.len() < 4 {
            return Ok(DispatchOutcome::Unauthorized);
        }

        let mut session_id_opt = None;

        if let Some(&sid) = self.addr_to_session.get(&peer) {
            if let Some(peer_state) = self.peer_machines.get(&sid) {
                let mut header = [0u8; 12];
                if packet.len() >= 12 {
                    header.copy_from_slice(&packet[0..12]);
                    let ciphertext = &packet[12..];
                    ostp_core::crypto::deobfuscate_header_inplace(&mut header, ciphertext, &peer_state.obfuscation_key, false);
                    let candidate_sid = u32::from_be_bytes([header[0], header[1], header[2], header[3]]);
                    if candidate_sid == sid {
                        session_id_opt = Some(sid);
                    }
                }
            }
        }

        if session_id_opt.is_none() {
            // Token Bucket rate limiter: mitigate seamless roaming CPU DoS vector
            let now = std::time::Instant::now();
            let elapsed = now.duration_since(self.last_token_regen).as_secs_f64();
            self.last_token_regen = now;
            self.roaming_tokens = (self.roaming_tokens + elapsed * 50.0).min(50.0);

            if self.roaming_tokens >= 1.0 {
                self.roaming_tokens -= 1.0;

                // Try seamless roaming over all peers
                for (&sid, peer_state) in &self.peer_machines {
                    if packet.len() >= 12 {
                        let mut header = [0u8; 12];
                        header.copy_from_slice(&packet[0..12]);
                        let ciphertext = &packet[12..];
                        ostp_core::crypto::deobfuscate_header_inplace(&mut header, ciphertext, &peer_state.obfuscation_key, false);
                        let candidate_sid = u32::from_be_bytes([header[0], header[1], header[2], header[3]]);
                        if candidate_sid == sid {
                            session_id_opt = Some(sid);
                            break;
                        }
                    }
                }
            }
        }

        if let Some(session_id) = session_id_opt {
            let key_opt = self.peer_machines.get(&session_id).map(|ps| ps.access_key.clone());
            if let Some(access_key) = key_opt {
                // Check if key is still valid and not over limit
                let key_valid = self.access_keys.read().unwrap_or_else(|e| e.into_inner()).contains_key(&access_key);
                let user_stats = self.get_or_create_user_stats(&access_key);
                if !key_valid || user_stats.is_over_limit() {
                    tracing::info!("Dropping session {} for key {} (valid={}, over_limit={})",
                        session_id, key_fp(&access_key), key_valid, user_stats.is_over_limit());
                    self.drop_session(session_id);
                    return Ok(DispatchOutcome::Unauthorized);
                }
            }

            if let Some(peer_state) = self.peer_machines.get_mut(&session_id) {

                if peer_state.last_addr != peer {
                    tracing::info!("Client roamed: session {} from {} to {}", session_id, peer_state.last_addr, peer);
                    self.addr_to_session.remove(&peer_state.last_addr);
                }
                peer_state.last_addr = peer;
                peer_state.last_seen = std::time::Instant::now();
                self.addr_to_session.insert(peer, session_id);
                
                // Track inbound bytes per user
                let key = peer_state.access_key.clone();
                track_user_bytes_up(&self.user_stats, &self.access_keys, &key, packet.len() as u64);

                let action = match peer_state.machine.on_event(OstpEvent::Inbound(packet)) {
                    Ok(a) => a,
                    Err(e) => {
                        tracing::warn!("Protocol error for session {}: {}", session_id, e);
                        return Ok(DispatchOutcome::Unauthorized);
                    }
                };

                let mut responses = Vec::new();
                let mut app_payloads = Vec::new();

                fn collect_action(
                    act: ProtocolAction,
                    sid: u32,
                    resps: &mut Vec<Bytes>,
                    loads: &mut Vec<(u32, u16, Bytes)>,
                ) {
                    match act {
                        ProtocolAction::SendDatagram(frame) => {
                            resps.push(frame);
                        }
                        ProtocolAction::DeliverApp(stream_id, data) => {
                            loads.push((sid, stream_id, data));
                        }
                        ProtocolAction::Multiple(list) => {
                            for item in list {
                                collect_action(item, sid, resps, loads);
                            }
                        }
                        _ => {}
                    }
                }

                collect_action(action, session_id, &mut responses, &mut app_payloads);

                return Ok(DispatchOutcome::Accepted {
                    responses,
                    app_payloads,
                    peer_addr: peer,
                });
            }
        }

        // Not an existing session — this is the expensive O(N_keys) trial path.
        // Gate it behind a token bucket so a garbage/spoofed-source flood cannot
        // force unbounded per-packet crypto work. Existing sessions (fast path
        // above) and roaming are unaffected. Regenerate at TRIAL_RATE/sec.
        {
            let now = std::time::Instant::now();
            let elapsed = now.duration_since(self.last_trial_regen).as_secs_f64();
            self.last_trial_regen = now;
            self.trial_tokens = (self.trial_tokens + elapsed * TRIAL_RATE).min(TRIAL_RATE);
            if self.trial_tokens < 1.0 {
                // Out of budget: drop silently (no response, no state, no log spam).
                return Ok(DispatchOutcome::Unauthorized);
            }
            self.trial_tokens -= 1.0;
        }

        let keys_snapshot: Vec<String> = self.access_keys.read().unwrap_or_else(|e| e.into_inner()).keys().cloned().collect();

        // Junk marker rotates per time window; check the current and previous
        // window so a client whose clock is up to ~1 window behind/ahead is still
        // recognised. Computed once per datagram, not per candidate key.
        let junk_window = ostp_core::crypto::current_junk_window();

        for candidate_key in keys_snapshot {
            let secrets = self.cached_secrets(&candidate_key);

            // Junk frames carry this key's time-rotating marker (no global
            // constant, no static per-user signature). Drop silently.
            if packet.len() >= 4 {
                let (m_now, m_prev) = self.cached_junk_markers(&candidate_key, junk_window);
                if packet[0..4] == m_now || packet[0..4] == m_prev {
                    return Ok(DispatchOutcome::Junk);
                }
            }

            // Decode the session_id using this key's obfuscation
            // The handshake mask is derived from the Noise payload at bytes [6..],
            // so we must deobfuscate the full packet, not just the header.
            if packet.len() < 7 { continue; }
            let mut trial = packet.to_vec();
            ostp_core::crypto::deobfuscate_packet_inplace(&mut trial, &secrets.obfuscation_key, true);
            let candidate_session_id = u32::from_be_bytes([trial[0], trial[1], trial[2], trial[3]]);

            let mut cfg = self.machine_config.clone();
            cfg.session_id = candidate_session_id;
            cfg.psk = secrets.psk;
            cfg.handshake_payload = vec![];
            cfg.obfuscation_key = secrets.obfuscation_key;
            cfg.handshake_pad_min = secrets.handshake_pad_min;
            cfg.handshake_pad_max = secrets.handshake_pad_max;

            let mut machine = match ProtocolMachine::new(cfg) {
                Ok(m) => m,
                Err(e) => {
                    tracing::warn!("Failed to create protocol machine for key trial: {}", e);
                    continue;
                }
            };
            let action = match machine.on_event(OstpEvent::Inbound(packet.clone())) {
                Ok(a) => a,
                Err(_) => continue,
            };

            if let ProtocolAction::HandshakePayload(payload, response_opt) = action {
                if payload.len() >= 12 {
                    let mut ts_bytes = [0_u8; 8];
                    ts_bytes.copy_from_slice(&payload[..8]);
                    let ts = u64::from_be_bytes(ts_bytes);

                    let mut sid_bytes = [0_u8; 4];
                    sid_bytes.copy_from_slice(&payload[8..12]);
                    let sid_from_payload = u32::from_be_bytes(sid_bytes);

                    if sid_from_payload != candidate_session_id {
                        continue;
                    }

                    let key_bytes = &payload[12..];
                    if let Ok(key_from_payload) = std::str::from_utf8(key_bytes) {
                        // The key embedded in the payload must match the candidate key we decoded with
                        if key_from_payload != candidate_key {
                            continue;
                        }

                        let now = std::time::SystemTime::now()
                            .duration_since(std::time::UNIX_EPOCH)
                            .unwrap_or_default()
                            .as_secs();

                        let drift = (now as i64 - ts as i64).abs();
                        if drift > 300 {
                            tracing::warn!("Handshake rejected: timestamp drift {}s exceeds 300s limit (peer={})", drift, peer);
                            continue;
                        }

                        if !self.replay_cache.contains_key(&payload.to_vec()) {
                            if self.replay_cache.len() >= REPLAY_CACHE_MAX {
                                // Don't globally reject new handshakes when full —
                                // that would let one flooding key-holder deny
                                // service to everyone. Reclaim space instead:
                                // first drop entries already past the drift
                                // window, then, if still full, evict the single
                                // oldest. A replay is still caught because it can
                                // only be accepted while within the 300s drift
                                // window, and an entry that young is never the
                                // one evicted before the cache genuinely holds
                                // 50k sub-300s handshakes.
                                self.replay_cache.retain(|_, &mut cached_ts| {
                                    (now as i64 - cached_ts as i64).abs() <= 300
                                });
                                if self.replay_cache.len() >= REPLAY_CACHE_MAX {
                                    if let Some(oldest) = self.replay_cache
                                        .iter()
                                        .min_by_key(|(_, &ts)| ts)
                                        .map(|(k, _)| k.clone())
                                    {
                                        self.replay_cache.remove(&oldest);
                                    }
                                    tracing::warn!("Replay cache full ({} entries), evicting oldest", REPLAY_CACHE_MAX);
                                }
                            }
                            if self.peer_machines.len() >= MAX_SESSIONS {
                                tracing::warn!("Max sessions reached ({}), rejecting handshake from {}", MAX_SESSIONS, peer);
                                return Ok(DispatchOutcome::Unauthorized);
                            }

                            self.replay_cache.insert(payload.to_vec(), ts);

                            machine.set_session_keys(candidate_session_id, secrets.obfuscation_key);

                            // Track per-user connection count
                            let user_stats = self.get_or_create_user_stats(&candidate_key);
                            user_stats.connections.fetch_add(1, Ordering::Relaxed);

                            // Check traffic limit before accepting
                            if user_stats.is_over_limit() {
                                tracing::warn!("User {} exceeded traffic limit, rejecting handshake from {}", key_fp(&candidate_key), peer);
                                return Ok(DispatchOutcome::Unauthorized);
                            }

                            self.peer_machines.insert(candidate_session_id, PeerState {
                                machine,
                                last_addr: peer,
                                obfuscation_key: secrets.obfuscation_key,
                                last_seen: std::time::Instant::now(),
                                access_key: candidate_key.clone(),
                            });
                            self.addr_to_session.insert(peer, candidate_session_id);

                            tracing::info!("New session authenticated: sid={} peer={} (active_sessions={}, replay_cache={})",
                                candidate_session_id, peer, self.peer_machines.len(), self.replay_cache.len()
                            );

                            return Ok(DispatchOutcome::Accepted {
                                responses: response_opt.into_iter().collect(),
                                app_payloads: Vec::new(),
                                peer_addr: peer,
                            });
                        }
                    }
                }
            }
        }

        Ok(DispatchOutcome::Unauthorized)
    }

    pub fn outbound_to_session(&mut self, session_id: u32, stream_id: u16, payload: Bytes) -> Result<Option<(Bytes, SocketAddr)>> {
        let peer_state = if let Some(existing) = self.peer_machines.get_mut(&session_id) {
            existing
        } else {
            return Ok(None);
        };

        let addr = peer_state.last_addr;
        let key = peer_state.access_key.clone();
        match peer_state.machine.on_event(OstpEvent::Outbound(stream_id, payload))? {
            ProtocolAction::SendDatagram(frame) => {
                // Track outbound bytes per user
                track_user_bytes_down(&self.user_stats, &self.access_keys, &key, frame.len() as u64);
                Ok(Some((frame, addr)))
            }
            _ => Ok(None),
        }
    }

    pub fn on_tick(&mut self) -> (Vec<(Bytes, SocketAddr)>, Vec<u32>) {
        // Purge expired handshakes from replay cache (older than 5 min drift allowance)
        let current_sys_time = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs();
        self.replay_cache.retain(|_, &mut ts| (current_sys_time as i64 - ts as i64).abs() <= 300);

        // Drop cached secrets/junk-markers for keys that have been deleted, so the
        // caches can't grow without bound as keys churn.
        {
            let keys = self.access_keys.read().unwrap_or_else(|e| e.into_inner());
            self.secrets_cache.retain(|k, _| keys.contains_key(k));
            self.junk_cache.retain(|k, _| keys.contains_key(k));
        }

        let mut frames = Vec::new();
        let mut expired = Vec::new();
        let now = std::time::Instant::now();
        let timeout_dur = std::time::Duration::from_secs(600); // 10 minute session timeout (mobile NAT can be up to 5-10min)

        // Gather expired or invalid sessions
        for (&sid, peer_state) in &self.peer_machines {
            let key_valid = self.access_keys.read().unwrap_or_else(|e| e.into_inner()).contains_key(&peer_state.access_key);
            let user_stats = self.get_or_create_user_stats(&peer_state.access_key);
            if now.duration_since(peer_state.last_seen) > timeout_dur || !key_valid || user_stats.is_over_limit() {
                expired.push(sid);
            }
        }

        // Clear expired/invalid sessions from internal state
        for sid in &expired {
            let peer_state_opt = self.peer_machines.get(sid);
            let reason = if let Some(ps) = peer_state_opt {
                let key_valid = self.access_keys.read().unwrap_or_else(|e| e.into_inner()).contains_key(&ps.access_key);
                let user_stats = self.get_or_create_user_stats(&ps.access_key);
                if now.duration_since(ps.last_seen) > timeout_dur {
                    "inactive >5min"
                } else if !key_valid {
                    "key deleted"
                } else if user_stats.is_over_limit() {
                    "traffic limit exceeded"
                } else {
                    "unknown"
                }
            } else {
                "unknown"
            };
            tracing::info!("Session {} closed ({}), releasing", sid, reason);
            self.drop_session(*sid);
        }

        // Drive ticks for remaining active sessions
        for peer_state in self.peer_machines.values_mut() {
            let action = match peer_state.machine.on_event(OstpEvent::Tick) {
                Ok(a) => a,
                Err(e) => {
                    tracing::warn!("Tick error for session: {}", e);
                    continue;
                }
            };

            let mut queue = vec![action];
            while let Some(current) = queue.pop() {
                match current {
                    ProtocolAction::Multiple(list) => {
                        for item in list {
                            queue.push(item);
                        }
                    }
                    ProtocolAction::SendDatagram(frame) => {
                        frames.push((frame, peer_state.last_addr));
                    }
                    _ => {}
                }
            }
        }

        (frames, expired)
    }

    pub fn drop_session(&mut self, session_id: u32) {
        if let Some(state) = self.peer_machines.remove(&session_id) {
            self.addr_to_session.remove(&state.last_addr);
        }
    }
}

// Free functions to avoid borrow-checker conflicts when tracking stats
// while holding a mutable reference to peer_machines.

fn get_or_create_stats(
    user_stats: &Arc<RwLock<HashMap<String, Arc<UserStats>>>>,
    access_keys: &Arc<RwLock<HashMap<String, crate::api::UserMeta>>>,
    key: &str,
) -> Arc<UserStats> {
    {
        let stats = user_stats.read().unwrap_or_else(|e| e.into_inner());
        if let Some(existing) = stats.get(key) {
            return existing.clone();
        }
    }
    
    let limit_bytes = access_keys.read().unwrap_or_else(|e| e.into_inner()).get(key).and_then(|m| m.limit_bytes);
    
    let mut stats = user_stats.write().unwrap_or_else(|e| e.into_inner());
    stats.entry(key.to_string())
        .or_insert_with(|| Arc::new(UserStats::new(limit_bytes)))
        .clone()
}

fn track_user_bytes_up(
    user_stats: &Arc<RwLock<HashMap<String, Arc<UserStats>>>>,
    access_keys: &Arc<RwLock<HashMap<String, crate::api::UserMeta>>>,
    key: &str,
    bytes: u64,
) {
    let stats = get_or_create_stats(user_stats, access_keys, key);
    stats.bytes_up.fetch_add(bytes, Ordering::Relaxed);
}

fn track_user_bytes_down(
    user_stats: &Arc<RwLock<HashMap<String, Arc<UserStats>>>>,
    access_keys: &Arc<RwLock<HashMap<String, crate::api::UserMeta>>>,
    key: &str,
    bytes: u64,
) {
    let stats = get_or_create_stats(user_stats, access_keys, key);
    stats.bytes_down.fetch_add(bytes, Ordering::Relaxed);
}