pjson-rs 0.6.0

//! WebSocket transport layer for real-time PJS streaming
//!
//! Provides WebSocket-based streaming with progressive JSON delivery
//! and backpressure handling for optimal client performance.

use crate::{
    Error as PjsError, Result as PjsResult, StreamFrame, domain::Priority, security::RateLimitGuard,
};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sha2::{Digest, Sha256};
use std::{
    collections::HashMap,
    future::Future,
    sync::Arc,
    time::{Duration, Instant},
};
use tokio::sync::{RwLock, broadcast};
use tracing::{debug, error, info, warn};
use uuid::Uuid;

#[cfg(feature = "websocket-client")]
pub mod client;
pub mod security;
#[cfg(feature = "http-server")]
pub mod server;

#[cfg(feature = "websocket-client")]
pub use client::*;
pub use security::SecureWebSocketHandler;
#[cfg(feature = "http-server")]
pub use server::*;

/// WebSocket message types for PJS streaming
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", content = "data")]
pub enum WsMessage {
    /// Stream initialization request
    StreamInit {
        /// Identifier of the WebSocket session.
        session_id: String,
        /// Source JSON payload to be streamed.
        data: Value,
        /// Per-stream options controlling framing and compression.
        options: StreamOptions,
    },
    /// Stream frame with priority data
    StreamFrame {
        /// Identifier of the WebSocket session.
        session_id: String,
        /// Monotonic frame index within the session.
        frame_id: u32,
        /// Priority assigned to this frame.
        priority: u8,
        /// Payload carried by the frame.
        payload: Value,
        /// Whether this frame completes the stream.
        is_complete: bool,
    },
    /// Client acknowledgment of frame
    FrameAck {
        /// Identifier of the WebSocket session.
        session_id: String,
        /// Index of the frame being acknowledged.
        frame_id: u32,
        /// Time the client took to process the frame, in milliseconds.
        processing_time_ms: u64,
    },
    /// Stream completion signal
    StreamComplete {
        /// Identifier of the WebSocket session.
        session_id: String,
        /// SHA-256 checksum of the concatenated frame payloads.
        checksum: String,
    },
    /// Error message
    Error {
        /// Identifier of the WebSocket session, if known.
        session_id: Option<String>,
        /// Human-readable error description.
        error: String,
        /// Numeric error code.
        code: u16,
    },
    /// Heartbeat/ping message
    Ping {
        /// Wall-clock timestamp at which the ping was sent.
        timestamp: u64,
    },
    /// Heartbeat/pong response
    Pong {
        /// Wall-clock timestamp at which the pong was sent.
        timestamp: u64,
    },
}

/// Stream configuration options
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamOptions {
    /// Maximum frame size in bytes
    pub max_frame_size: usize,
    /// Client processing capability (frames per second)
    pub client_fps: Option<u32>,
    /// Enable compression
    pub compression: bool,
    /// Custom priority mapping
    pub priority_mapping: Option<HashMap<String, u8>>,
}

impl Default for StreamOptions {
    fn default() -> Self {
        Self {
            max_frame_size: 64 * 1024, // 64KB
            client_fps: None,          // Auto-detect
            compression: true,
            priority_mapping: None,
        }
    }
}

/// WebSocket streaming session state.
///
/// This type is intentionally distinct from the domain-layer
/// [`crate::domain::aggregates::StreamSession`] aggregate. The WebSocket
/// transport maintains an ephemeral, transport-local session model with raw
/// `String` identifiers and an in-memory `HashMap` keyed off
/// [`AdaptiveStreamController`]; it does **not** share state with the
/// `StreamRepositoryGat`-backed domain session created by
/// `POST /pjs/sessions`. Sessions created over WebSocket cannot be addressed
/// over HTTP and vice versa, and dictionary training, auth, and rate-limit
/// middleware applied to the HTTP router do not apply to this controller.
///
/// See issue #239 for the full rationale and the deliberate split between
/// the two session models.
#[derive(Debug)]
pub struct WebSocketStreamSession {
    /// Transport-local session identifier.
    pub id: String,
    /// Instant the session was created.
    pub created_at: Instant,
    /// Streaming options negotiated for this session.
    pub options: StreamOptions,
    /// Pre-computed delivery plan as an ordered list of frames.
    pub plan: Vec<StreamFrame>,
    /// Index of the next frame to send.
    pub current_frame: u32,
    /// Frame indices acknowledged by the client so far.
    pub acknowledged_frames: Vec<u32>,
    /// Adaptive streaming metrics derived from client acks.
    pub client_metrics: ClientMetrics,
    /// Rate-limit guard scoped to this session, if installed.
    pub rate_limit_guard: Option<RateLimitGuard>,
}

/// Client performance metrics for adaptive streaming
#[derive(Debug, Default)]
pub struct ClientMetrics {
    /// Exponential moving average of client frame-processing time, in milliseconds.
    pub average_processing_time_ms: f64,
    /// Number of frames the client has acknowledged.
    pub frames_acknowledged: u32,
    /// Instant of the most recent acknowledgement, if any.
    pub last_ack_time: Option<Instant>,
    /// Estimated downlink bandwidth, in kilobits per second, if measured.
    pub estimated_bandwidth_kbps: Option<f64>,
    /// Round-trip time of the WebSocket connection, in milliseconds, if measured.
    pub connection_rtt_ms: Option<u64>,
}

impl ClientMetrics {
    /// Fold a new processing-time observation into the moving average.
    pub fn update_processing_time(&mut self, processing_time_ms: u64) {
        let new_time = processing_time_ms as f64;
        if self.frames_acknowledged == 0 {
            self.average_processing_time_ms = new_time;
        } else {
            // Exponential moving average
            let alpha = 0.3;
            self.average_processing_time_ms =
                alpha * new_time + (1.0 - alpha) * self.average_processing_time_ms;
        }
        self.frames_acknowledged += 1;
        self.last_ack_time = Some(Instant::now());
    }

    /// Returns `true` when average client processing time exceeds the slow-client threshold.
    pub fn is_client_slow(&self) -> bool {
        self.average_processing_time_ms > 100.0 // > 100ms per frame
    }

    /// Recommended delay between frames given the current processing-time average.
    pub fn recommended_frame_delay(&self) -> Duration {
        if self.is_client_slow() {
            Duration::from_millis((self.average_processing_time_ms * 0.5) as u64)
        } else {
            Duration::from_millis(10) // Fast clients get minimal delay
        }
    }
}

/// WebSocket transport trait for different implementations (GAT-based)
pub trait WebSocketTransport: Send + Sync {
    /// Concrete connection type the implementor uses for I/O.
    type Connection: Send + Sync;

    /// Future type for starting stream
    type StartStreamFuture<'a>: Future<Output = PjsResult<String>> + Send + 'a
    where
        Self: 'a;

    /// Future type for sending frame
    type SendFrameFuture<'a>: Future<Output = PjsResult<()>> + Send + 'a
    where
        Self: 'a;

    /// Future type for handling message
    type HandleMessageFuture<'a>: Future<Output = PjsResult<()>> + Send + 'a
    where
        Self: 'a;

    /// Future type for closing stream
    type CloseStreamFuture<'a>: Future<Output = PjsResult<()>> + Send + 'a
    where
        Self: 'a;

    /// Start streaming session
    fn start_stream(
        &self,
        connection: Arc<Self::Connection>,
        data: Value,
        options: StreamOptions,
    ) -> Self::StartStreamFuture<'_>;

    /// Send frame to client
    fn send_frame(
        &self,
        connection: Arc<Self::Connection>,
        message: WsMessage,
    ) -> Self::SendFrameFuture<'_>;

    /// Handle incoming message
    fn handle_message(
        &self,
        connection: Arc<Self::Connection>,
        message: WsMessage,
    ) -> Self::HandleMessageFuture<'_>;

    /// Close streaming session
    fn close_stream(&self, session_id: &str) -> Self::CloseStreamFuture<'_>;
}

/// Adaptive streaming controller
pub struct AdaptiveStreamController {
    sessions: Arc<RwLock<HashMap<String, WebSocketStreamSession>>>,
    frame_tx: broadcast::Sender<(String, WsMessage)>,
}

impl AdaptiveStreamController {
    /// Create an empty controller with no active sessions.
    pub fn new() -> Self {
        let (frame_tx, _) = broadcast::channel(1000);

        Self {
            sessions: Arc::new(RwLock::new(HashMap::new())),
            frame_tx,
        }
    }

    /// Create new streaming session
    pub async fn create_session(&self, data: Value, options: StreamOptions) -> PjsResult<String> {
        let session_id = Uuid::new_v4().to_string();
        let plan = vec![StreamFrame {
            data: data.clone(),
            priority: Priority::HIGH,
            metadata: std::collections::HashMap::new(),
        }]; // Simplified for now

        let session = WebSocketStreamSession {
            id: session_id.clone(),
            created_at: Instant::now(),
            options,
            plan,
            current_frame: 0,
            acknowledged_frames: Vec::new(),
            client_metrics: ClientMetrics::default(),
            rate_limit_guard: None, // Will be set when connection is established
        };

        self.sessions
            .write()
            .await
            .insert(session_id.clone(), session);

        info!("Created streaming session: {}", session_id);
        Ok(session_id)
    }

    /// Start streaming frames for session
    pub async fn start_streaming(&self, session_id: &str) -> PjsResult<()> {
        let mut sessions = self.sessions.write().await;
        let session = sessions
            .get_mut(session_id)
            .ok_or_else(|| PjsError::InvalidSession(session_id.to_string()))?;

        // Start streaming task
        let session_id = session_id.to_string();
        let frame_tx = self.frame_tx.clone();
        let plan = session.plan.clone();

        tokio::spawn(async move {
            if let Err(e) = Self::stream_frames(session_id, plan, frame_tx).await {
                error!("Error streaming frames: {}", e);
            }
        });

        Ok(())
    }

    async fn stream_frames(
        session_id: String,
        plan: Vec<StreamFrame>, // Simplified for now
        frame_tx: broadcast::Sender<(String, WsMessage)>,
    ) -> Result<(), PjsError> {
        let mut frames_data = Vec::new();

        for (frame_id, frame) in plan.iter().enumerate() {
            // Collect frame payload for checksum calculation
            let payload_bytes =
                serde_json::to_vec(&frame.data).map_err(|e| PjsError::Other(e.to_string()))?;
            frames_data.push(payload_bytes);

            let ws_message = WsMessage::StreamFrame {
                session_id: session_id.clone(),
                frame_id: frame_id as u32,
                priority: frame.priority.value(),
                payload: frame.data.clone(),
                is_complete: frame_id == (plan.len() - 1),
            };

            if let Err(e) = frame_tx.send((session_id.clone(), ws_message)) {
                error!("Failed to send frame {}: {}", frame_id, e);
                break;
            }

            // TODO: Add adaptive delay based on client metrics
            tokio::time::sleep(Duration::from_millis(10)).await;
        }

        // Send completion message with calculated checksum
        let complete_message = WsMessage::StreamComplete {
            session_id: session_id.clone(),
            checksum: calculate_stream_checksum(&frames_data),
        };

        let _ = frame_tx.send((session_id, complete_message));
        Ok(())
    }

    /// Handle frame acknowledgment
    pub async fn handle_frame_ack(
        &self,
        session_id: &str,
        frame_id: u32,
        processing_time_ms: u64,
    ) -> PjsResult<()> {
        let mut sessions = self.sessions.write().await;
        let session = sessions
            .get_mut(session_id)
            .ok_or_else(|| PjsError::InvalidSession(session_id.to_string()))?;

        session.acknowledged_frames.push(frame_id);
        session
            .client_metrics
            .update_processing_time(processing_time_ms);

        debug!(
            "Frame {} acknowledged for session {} (processing: {}ms, avg: {:.1}ms)",
            frame_id,
            session_id,
            processing_time_ms,
            session.client_metrics.average_processing_time_ms
        );

        if session.client_metrics.is_client_slow() {
            warn!(
                "Client {} is processing slowly (avg: {:.1}ms)",
                session_id, session.client_metrics.average_processing_time_ms
            );
        }

        Ok(())
    }

    /// Get subscriber for frame events
    pub fn subscribe_frames(&self) -> broadcast::Receiver<(String, WsMessage)> {
        self.frame_tx.subscribe()
    }

    /// Set rate limit guard for a session
    pub async fn set_rate_limit_guard(
        &self,
        session_id: &str,
        guard: RateLimitGuard,
    ) -> PjsResult<()> {
        let mut sessions = self.sessions.write().await;
        let session = sessions
            .get_mut(session_id)
            .ok_or_else(|| PjsError::InvalidSession(session_id.to_string()))?;

        session.rate_limit_guard = Some(guard);
        Ok(())
    }

    /// Validate message against rate limits
    pub async fn validate_message(&self, session_id: &str, frame_size: usize) -> PjsResult<()> {
        let sessions = self.sessions.read().await;
        let session = sessions
            .get(session_id)
            .ok_or_else(|| PjsError::InvalidSession(session_id.to_string()))?;

        if let Some(guard) = &session.rate_limit_guard {
            guard
                .check_message(frame_size)
                .map_err(|e| PjsError::SecurityError(format!("Rate limit violation: {}", e)))?;
        }

        Ok(())
    }

    /// Remove a single session by id.
    ///
    /// Returns `true` if the session existed and was removed, `false` if the id
    /// was not present. Callers may safely invoke this multiple times — the
    /// second call is a no-op.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # use pjson_rs::infrastructure::websocket::{AdaptiveStreamController, StreamOptions};
    /// # use serde_json::json;
    /// # #[tokio::main] async fn main() {
    /// let controller = AdaptiveStreamController::new();
    /// let id = controller.create_session(json!({}), StreamOptions::default()).await.unwrap();
    /// assert!(controller.remove_session(&id).await);
    /// // Idempotent — second call is a no-op:
    /// assert!(!controller.remove_session(&id).await);
    /// # }
    /// ```
    pub async fn remove_session(&self, session_id: &str) -> bool {
        let mut sessions = self.sessions.write().await;
        let removed = sessions.remove(session_id).is_some();
        if removed {
            info!("Removed streaming session: {}", session_id);
        } else {
            debug!("remove_session called on unknown id: {}", session_id);
        }
        removed
    }

    /// Clean up expired sessions
    pub async fn cleanup_expired_sessions(&self, max_age: Duration) {
        let mut sessions = self.sessions.write().await;
        let now = Instant::now();

        sessions.retain(|id, session| {
            let expired = now.duration_since(session.created_at) > max_age;
            if expired {
                info!("Cleaning up expired session: {}", id);
            }
            !expired
        });
    }
}

impl Default for AdaptiveStreamController {
    fn default() -> Self {
        Self::new()
    }
}

/// Calculate SHA-256 checksum for stream completion verification
fn calculate_stream_checksum(frames_data: &[Vec<u8>]) -> String {
    let mut hasher = Sha256::new();

    // Hash each frame's data
    for frame_data in frames_data {
        hasher.update(frame_data);
    }

    // Hash frame count to ensure integrity
    hasher.update((frames_data.len() as u64).to_le_bytes());

    let result = hasher.finalize();
    format!("sha256:{result:x}")
}

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

    #[tokio::test]
    async fn test_create_session() {
        let controller = AdaptiveStreamController::new();
        let data = json!({
            "critical": {"id": 1, "status": "active"},
            "details": {"name": "test", "description": "test data"}
        });

        let session_id = controller
            .create_session(data, StreamOptions::default())
            .await
            .unwrap();

        assert!(!session_id.is_empty());

        let sessions = controller.sessions.read().await;
        assert!(sessions.contains_key(&session_id));
    }

    #[tokio::test]
    async fn test_frame_acknowledgment() {
        let controller = AdaptiveStreamController::new();
        let data = json!({"test": "data"});

        let session_id = controller
            .create_session(data, StreamOptions::default())
            .await
            .unwrap();

        controller
            .handle_frame_ack(&session_id, 0, 50)
            .await
            .unwrap();

        let sessions = controller.sessions.read().await;
        let session = sessions.get(&session_id).unwrap();
        assert_eq!(session.acknowledged_frames, vec![0]);
        assert_eq!(session.client_metrics.average_processing_time_ms, 50.0);
    }

    #[test]
    fn test_client_metrics() {
        let mut metrics = ClientMetrics::default();

        metrics.update_processing_time(100);
        assert_eq!(metrics.average_processing_time_ms, 100.0);

        metrics.update_processing_time(200);
        // Should be exponential moving average: 0.3 * 200 + 0.7 * 100 = 130
        assert!((metrics.average_processing_time_ms - 130.0).abs() < 0.1);

        assert!(metrics.is_client_slow());
    }

    #[test]
    fn test_checksum_calculation() {
        // Test empty frames
        let empty_frames: Vec<Vec<u8>> = vec![];
        let checksum = calculate_stream_checksum(&empty_frames);
        assert!(checksum.starts_with("sha256:"));

        // Test single frame
        let single_frame = vec![vec![1, 2, 3, 4]];
        let checksum1 = calculate_stream_checksum(&single_frame);
        assert!(checksum1.starts_with("sha256:"));

        // Test multiple frames
        let multi_frames = vec![vec![1, 2], vec![3, 4], vec![5, 6]];
        let checksum2 = calculate_stream_checksum(&multi_frames);
        assert!(checksum2.starts_with("sha256:"));

        // Same data should produce same checksum
        let same_frames = vec![vec![1, 2], vec![3, 4], vec![5, 6]];
        let checksum3 = calculate_stream_checksum(&same_frames);
        assert_eq!(checksum2, checksum3);

        // Different data should produce different checksum
        let diff_frames = vec![vec![1, 2], vec![3, 4], vec![5, 7]]; // Last byte different
        let checksum4 = calculate_stream_checksum(&diff_frames);
        assert_ne!(checksum2, checksum4);

        // Different order should produce different checksum
        let reordered_frames = vec![vec![3, 4], vec![1, 2], vec![5, 6]];
        let checksum5 = calculate_stream_checksum(&reordered_frames);
        assert_ne!(checksum2, checksum5);
    }
}