coderlib 0.1.0

//! WebSocket transport implementation for MCP
//!
//! This module provides WebSocket transport capabilities for MCP connections,
//! supporting both client and server roles with bidirectional communication.

use std::marker::PhantomData;
use std::pin::Pin;
use std::task::{Context, Poll};

use futures::{Sink, Stream};
use pin_project_lite::pin_project;

#[cfg(feature = "websocket")]
use tokio_tungstenite::{
    connect_async, accept_async,
    tungstenite::{self, Message},
    WebSocketStream, MaybeTlsStream,
};

#[cfg(feature = "mcp")]
use rmcp::{
    RoleClient, RoleServer,
    service::{RxJsonRpcMessage, ServiceRole, TxJsonRpcMessage},
};

use super::transport_enhanced::{EnhancedTransport, TransportStats, TransportError};

/// WebSocket transport for MCP connections
#[cfg(feature = "websocket")]
pin_project! {
    pub struct WebSocketTransport<R> {
        #[pin]
        stream: WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>,
        stats: TransportStats,
        url: String,
        marker: PhantomData<R>,
    }
}

#[cfg(feature = "websocket")]
impl<R> WebSocketTransport<R>
where
    R: ServiceRole,
{
    /// Create a new WebSocket client connection
    pub async fn connect(url: &str) -> Result<Self, TransportError> {
        let (stream, response) = connect_async(url).await?;
        
        if response.status() != tungstenite::http::StatusCode::SWITCHING_PROTOCOLS {
            return Err(TransportError::ConnectionFailed(
                format!("Failed to upgrade connection: {}", response.status())
            ));
        }

        Ok(Self {
            stream,
            stats: TransportStats::default(),
            url: url.to_string(),
            marker: PhantomData,
        })
    }

    /// Create a WebSocket server from an accepted TCP stream
    pub async fn from_tcp_stream(stream: tokio::net::TcpStream) -> Result<Self, TransportError> {
        let ws_stream = accept_async(MaybeTlsStream::Plain(stream)).await?;

        Ok(Self {
            stream: ws_stream,
            stats: TransportStats::default(),
            url: "server".to_string(),
            marker: PhantomData,
        })
    }

    /// Create a client-side WebSocket transport
    pub fn new_client(stream: WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>) -> Self {
        Self {
            stream,
            stats: TransportStats::default(),
            url: "client".to_string(),
            marker: PhantomData,
        }
    }

    /// Create a server-side WebSocket transport
    pub fn new_server(stream: WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>) -> Self {
        Self {
            stream,
            stats: TransportStats::default(),
            url: "server".to_string(),
            marker: PhantomData,
        }
    }

    /// Get mutable reference to stats
    fn stats_mut(&mut self) -> &mut TransportStats {
        &mut self.stats
    }
}

#[cfg(feature = "websocket")]
impl<R> EnhancedTransport<R> for WebSocketTransport<R>
where
    R: ServiceRole,
{
    fn description(&self) -> String {
        format!("WebSocket transport ({})", self.url)
    }

    fn is_connected(&self) -> bool {
        // In a real implementation, you'd check the WebSocket state
        // For now, we assume it's connected if the struct exists
        true
    }

    fn stats(&self) -> TransportStats {
        self.stats.clone()
    }
}

#[cfg(all(feature = "websocket", feature = "mcp"))]
impl<R> Stream for WebSocketTransport<R>
where
    R: ServiceRole,
{
    type Item = RxJsonRpcMessage<R>;

    fn poll_next(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Self::Item>> {
        let this = self.as_mut().project();
        
        match this.stream.poll_next(cx) {
            Poll::Ready(Some(Ok(message))) => {
                let message_text = match message {
                    Message::Text(text) => text,
                    Message::Binary(data) => {
                        // Try to convert binary to text
                        match String::from_utf8(data) {
                            Ok(text) => text,
                            Err(_) => {
                                tracing::warn!("Received non-UTF8 binary message");
                                return self.poll_next(cx);
                            }
                        }
                    }
                    Message::Close(_) => {
                        tracing::info!("WebSocket connection closed");
                        return Poll::Ready(None);
                    }
                    Message::Ping(_) | Message::Pong(_) => {
                        // Handle ping/pong automatically, continue polling
                        return self.poll_next(cx);
                    }
                    Message::Frame(_) => {
                        tracing::warn!("Received raw frame message");
                        return self.poll_next(cx);
                    }
                };

                // Update stats
                this.stats.messages_received += 1;

                // Parse JSON-RPC message
                match serde_json::from_str::<RxJsonRpcMessage<R>>(&message_text) {
                    Ok(message) => Poll::Ready(Some(message)),
                    Err(e) => {
                        tracing::warn!(error = %e, message = %message_text, "Failed to parse JSON-RPC message");
                        this.stats.connection_errors += 1;
                        this.stats.last_error = Some(e.to_string());
                        self.poll_next(cx)
                    }
                }
            }
            Poll::Ready(Some(Err(e))) => {
                tracing::warn!(error = %e, "WebSocket error");
                let this = self.as_mut().project();
                this.stats.connection_errors += 1;
                this.stats.last_error = Some(e.to_string());
                self.poll_next(cx)
            }
            Poll::Ready(None) => {
                tracing::info!("WebSocket stream ended");
                Poll::Ready(None)
            }
            Poll::Pending => Poll::Pending,
        }
    }
}

#[cfg(all(feature = "websocket", feature = "mcp"))]
impl<R> Sink<TxJsonRpcMessage<R>> for WebSocketTransport<R>
where
    R: ServiceRole,
{
    type Error = tungstenite::Error;

    fn poll_ready(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), Self::Error>> {
        let this = self.project();
        this.stream.poll_ready(cx)
    }

    fn start_send(
        mut self: Pin<&mut Self>,
        item: TxJsonRpcMessage<R>,
    ) -> Result<(), Self::Error> {
        // Update stats
        self.stats_mut().messages_sent += 1;

        let this = self.project();
        let message_text = serde_json::to_string(&item)
            .expect("JSON-RPC message should be valid JSON");
        let message = Message::Text(message_text);
        this.stream.start_send(message)
    }

    fn poll_flush(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), Self::Error>> {
        let this = self.project();
        this.stream.poll_flush(cx)
    }

    fn poll_close(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), Self::Error>> {
        let this = self.project();
        this.stream.poll_close(cx)
    }
}

/// WebSocket server for accepting MCP connections
#[cfg(feature = "websocket")]
pub struct WebSocketServer {
    listener: tokio::net::TcpListener,
}

#[cfg(feature = "websocket")]
impl WebSocketServer {
    /// Create a new WebSocket server
    pub async fn bind(addr: &str) -> Result<Self, TransportError> {
        let listener = tokio::net::TcpListener::bind(addr).await?;
        Ok(Self { listener })
    }

    /// Accept a new WebSocket connection
    pub async fn accept<R>(&self) -> Result<WebSocketTransport<R>, TransportError>
    where
        R: ServiceRole,
    {
        let (stream, _addr) = self.listener.accept().await?;
        WebSocketTransport::from_tcp_stream(stream).await
    }

    /// Get the local address the server is bound to
    pub fn local_addr(&self) -> Result<std::net::SocketAddr, std::io::Error> {
        self.listener.local_addr()
    }
}

/// WebSocket client for connecting to MCP servers
#[cfg(feature = "websocket")]
pub struct WebSocketClient;

#[cfg(feature = "websocket")]
impl WebSocketClient {
    /// Connect to a WebSocket MCP server
    pub async fn connect<R>(url: &str) -> Result<WebSocketTransport<R>, TransportError>
    where
        R: ServiceRole,
    {
        WebSocketTransport::connect(url).await
    }

    /// Connect with custom headers
    pub async fn connect_with_headers<R>(
        url: &str,
        headers: std::collections::HashMap<String, String>,
    ) -> Result<WebSocketTransport<R>, TransportError>
    where
        R: ServiceRole,
    {
        // For now, we ignore headers - in a full implementation,
        // you'd use a custom connector with headers
        let _ = headers;
        WebSocketTransport::connect(url).await
    }
}

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

    #[cfg(feature = "websocket")]
    #[test]
    fn test_websocket_transport_creation() {
        // This test just verifies the types compile correctly
        // Actual connection tests would require a running WebSocket server
    }

    #[test]
    fn test_transport_stats() {
        let stats = TransportStats::default();
        assert_eq!(stats.messages_sent, 0);
        assert_eq!(stats.messages_received, 0);
        assert_eq!(stats.connection_errors, 0);
        assert!(stats.last_error.is_none());
    }
}