mechutil 0.8.0

//
// Copyright (C) 2024 - 2025 Automated Design Corp. All Rights Reserved.
//

//! TCP transport layer with length-prefix framing for IPC messages.
//!
//! This module provides the low-level transport for sending and receiving
//! CommandMessages over TCP connections with proper framing.
//!
//! ## Wire Format
//!
//! Each message is framed as:
//! ```text
//! +----------------+------------+------------------+
//! | Length (4 bytes) | Type (1 byte) | Payload (N bytes) |
//! +----------------+------------+------------------+
//! ```
//!
//! - Length: u32 little-endian, size of payload only (not including header)
//! - Type: u8, MessageType value (redundant with payload but enables quick routing)
//! - Payload: JSON-serialized CommandMessage

use std::sync::Arc;
use tokio::io::{AsyncReadExt, AsyncWriteExt, BufReader, BufWriter};
use tokio::net::TcpStream;
use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
use tokio::sync::Mutex;

use super::command_message::CommandMessage;
use super::error::IpcError;
use super::message::IpcMessageHeader;

/// Maximum message size (16 MB)
const MAX_MESSAGE_SIZE: u32 = 16 * 1024 * 1024;

/// Default buffer size for reading/writing
const BUFFER_SIZE: usize = 64 * 1024;

/// A framed TCP stream for IPC communication.
///
/// This wraps a TcpStream and provides methods to send and receive
/// length-prefixed CommandMessages.
pub struct FramedStream {
    reader: BufReader<OwnedReadHalf>,
    writer: BufWriter<OwnedWriteHalf>,
}

impl FramedStream {
    /// Create a new FramedStream from a TcpStream.
    pub fn new(stream: TcpStream) -> Self {
        let (read_half, write_half) = stream.into_split();
        Self {
            reader: BufReader::with_capacity(BUFFER_SIZE, read_half),
            writer: BufWriter::with_capacity(BUFFER_SIZE, write_half),
        }
    }

    /// Send a CommandMessage over the stream.
    ///
    /// The message is serialized with JSON and framed with a length prefix.
    pub async fn send(&mut self, msg: &CommandMessage) -> Result<(), IpcError> {
        // Serialize the payload using JSON
        let payload = serde_json::to_vec(msg)?;
        let payload_len = payload.len() as u32;

        if payload_len > MAX_MESSAGE_SIZE {
            return Err(IpcError::Framing(format!(
                "Message too large: {} bytes (max {})",
                payload_len, MAX_MESSAGE_SIZE
            )));
        }

        // Write header (includes message_type for quick routing without deserialization)
        let header = IpcMessageHeader::new(msg.message_type, payload_len);
        self.writer.write_all(&header.to_bytes()).await?;

        // Write payload
        self.writer.write_all(&payload).await?;

        // Flush to ensure message is sent
        self.writer.flush().await?;

        Ok(())
    }

    /// Receive a CommandMessage from the stream.
    ///
    /// Blocks until a complete message is received or an error occurs.
    pub async fn recv(&mut self) -> Result<CommandMessage, IpcError> {
        // Read header
        let mut header_bytes = [0u8; IpcMessageHeader::SIZE];
        self.reader.read_exact(&mut header_bytes).await?;

        let header = IpcMessageHeader::from_bytes(&header_bytes);

        if header.length > MAX_MESSAGE_SIZE {
            return Err(IpcError::Framing(format!(
                "Message too large: {} bytes (max {})",
                header.length, MAX_MESSAGE_SIZE
            )));
        }

        // Read payload
        let mut payload = vec![0u8; header.length as usize];
        self.reader.read_exact(&mut payload).await?;

        // Deserialize payload using JSON
        let msg: CommandMessage = serde_json::from_slice(&payload)?;

        Ok(msg)
    }

    /// Try to receive a message with a timeout.
    ///
    /// Returns None if the timeout expires before a message is received.
    pub async fn recv_timeout(
        &mut self,
        timeout: std::time::Duration,
    ) -> Result<Option<CommandMessage>, IpcError> {
        match tokio::time::timeout(timeout, self.recv()).await {
            Ok(result) => result.map(Some),
            Err(_) => Ok(None),
        }
    }

    /// Send a message and wait for a response with the matching transaction_id.
    pub async fn send_recv(
        &mut self,
        msg: &CommandMessage,
        timeout: std::time::Duration,
    ) -> Result<CommandMessage, IpcError> {
        let transaction_id = msg.transaction_id;
        self.send(msg).await?;

        let deadline = tokio::time::Instant::now() + timeout;

        loop {
            let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
            if remaining.is_zero() {
                return Err(IpcError::Timeout(format!(
                    "Timeout waiting for response to transaction {}",
                    transaction_id
                )));
            }

            match self.recv_timeout(remaining).await? {
                Some(response) => {
                    if response.transaction_id == transaction_id {
                        return Ok(response);
                    }
                    // Not our response, keep waiting
                    // In a real implementation, you might queue this for another handler
                }
                None => {
                    return Err(IpcError::Timeout(format!(
                        "Timeout waiting for response to transaction {}",
                        transaction_id
                    )));
                }
            }
        }
    }

    /// Close the stream gracefully.
    pub async fn close(mut self) -> Result<(), IpcError> {
        self.writer.flush().await?;
        self.writer.shutdown().await?;
        Ok(())
    }
}

/// Thread-safe wrapper around FramedStream for concurrent access.
///
/// This allows multiple tasks to send and receive messages on the same
/// connection, with proper synchronization.
pub struct IpcTransport {
    inner: Arc<Mutex<FramedStream>>,
}

impl IpcTransport {
    /// Create a new IpcTransport from a TcpStream.
    pub fn new(stream: TcpStream) -> Self {
        Self {
            inner: Arc::new(Mutex::new(FramedStream::new(stream))),
        }
    }

    /// Create from an existing FramedStream.
    pub fn from_framed(stream: FramedStream) -> Self {
        Self {
            inner: Arc::new(Mutex::new(stream)),
        }
    }

    /// Send a message.
    pub async fn send(&self, msg: &CommandMessage) -> Result<(), IpcError> {
        let mut guard = self.inner.lock().await;
        guard.send(msg).await
    }

    /// Receive a message.
    pub async fn recv(&self) -> Result<CommandMessage, IpcError> {
        let mut guard = self.inner.lock().await;
        guard.recv().await
    }

    /// Receive with timeout.
    pub async fn recv_timeout(
        &self,
        timeout: std::time::Duration,
    ) -> Result<Option<CommandMessage>, IpcError> {
        let mut guard = self.inner.lock().await;
        guard.recv_timeout(timeout).await
    }

    /// Send and receive with matching transaction_id.
    pub async fn send_recv(
        &self,
        msg: &CommandMessage,
        timeout: std::time::Duration,
    ) -> Result<CommandMessage, IpcError> {
        let mut guard = self.inner.lock().await;
        guard.send_recv(msg, timeout).await
    }

    /// Clone the transport handle.
    pub fn clone_handle(&self) -> Self {
        Self {
            inner: Arc::clone(&self.inner),
        }
    }
}

impl Clone for IpcTransport {
    fn clone(&self) -> Self {
        self.clone_handle()
    }
}

/// Split transport for separate read/write tasks.
///
/// This is useful when you want one task to handle incoming messages
/// and another task to send outgoing messages.
pub struct SplitTransport {
    reader: Arc<Mutex<BufReader<OwnedReadHalf>>>,
    writer: Arc<Mutex<BufWriter<OwnedWriteHalf>>>,
}

impl std::fmt::Debug for SplitTransport {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SplitTransport").finish_non_exhaustive()
    }
}

impl SplitTransport {
    /// Create a new SplitTransport from a TcpStream.
    pub fn new(stream: TcpStream) -> Self {
        let (read_half, write_half) = stream.into_split();
        Self {
            reader: Arc::new(Mutex::new(BufReader::with_capacity(BUFFER_SIZE, read_half))),
            writer: Arc::new(Mutex::new(BufWriter::with_capacity(BUFFER_SIZE, write_half))),
        }
    }

    /// Send a message.
    pub async fn send(&self, msg: &CommandMessage) -> Result<(), IpcError> {
        let payload = serde_json::to_vec(msg)?;
        let payload_len = payload.len() as u32;

        if payload_len > MAX_MESSAGE_SIZE {
            return Err(IpcError::Framing(format!(
                "Message too large: {} bytes (max {})",
                payload_len, MAX_MESSAGE_SIZE
            )));
        }

        let header = IpcMessageHeader::new(msg.message_type, payload_len);

        let mut writer = self.writer.lock().await;
        writer.write_all(&header.to_bytes()).await?;
        writer.write_all(&payload).await?;
        writer.flush().await?;

        Ok(())
    }

    /// Receive a message.
    pub async fn recv(&self) -> Result<CommandMessage, IpcError> {
        let mut reader = self.reader.lock().await;

        // Read header
        let mut header_bytes = [0u8; IpcMessageHeader::SIZE];
        reader.read_exact(&mut header_bytes).await?;

        let header = IpcMessageHeader::from_bytes(&header_bytes);

        if header.length > MAX_MESSAGE_SIZE {
            return Err(IpcError::Framing(format!(
                "Message too large: {} bytes (max {})",
                header.length, MAX_MESSAGE_SIZE
            )));
        }

        // Read payload
        let mut payload = vec![0u8; header.length as usize];
        reader.read_exact(&mut payload).await?;

        let msg: CommandMessage = serde_json::from_slice(&payload)?;

        Ok(msg)
    }

    /// Clone the reader handle only.
    pub fn clone_reader(&self) -> Arc<Mutex<BufReader<OwnedReadHalf>>> {
        Arc::clone(&self.reader)
    }

    /// Clone the writer handle only.
    pub fn clone_writer(&self) -> Arc<Mutex<BufWriter<OwnedWriteHalf>>> {
        Arc::clone(&self.writer)
    }
}

impl Clone for SplitTransport {
    fn clone(&self) -> Self {
        Self {
            reader: Arc::clone(&self.reader),
            writer: Arc::clone(&self.writer),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ipc::MessageType;
    use tokio::net::TcpListener;

    #[tokio::test]
    async fn test_framed_stream_roundtrip() {
        // Create a listener
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        // Spawn server task
        let server_task = tokio::spawn(async move {
            let (stream, _) = listener.accept().await.unwrap();
            let mut framed = FramedStream::new(stream);

            // Receive message
            let msg = framed.recv().await.unwrap();
            assert_eq!(msg.message_type, MessageType::Request);
            assert_eq!(msg.topic, "test.function");

            // Send response
            let response = msg.into_response(serde_json::json!({"status": "ok"}));
            framed.send(&response).await.unwrap();
        });

        // Client
        let stream = TcpStream::connect(addr).await.unwrap();
        let mut framed = FramedStream::new(stream);

        // Send request
        let request = CommandMessage::request("test.function", serde_json::json!({"key": "value"}));
        let request_id = request.transaction_id;
        framed.send(&request).await.unwrap();

        // Receive response
        let response = framed.recv().await.unwrap();
        assert_eq!(response.message_type, MessageType::Response);
        assert_eq!(response.transaction_id, request_id);
        assert!(response.success);

        server_task.await.unwrap();
    }

    #[tokio::test]
    async fn test_send_recv_with_timeout() {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        let server_task = tokio::spawn(async move {
            let (stream, _) = listener.accept().await.unwrap();
            let mut framed = FramedStream::new(stream);

            // Receive and respond
            let msg = framed.recv().await.unwrap();
            let response = msg.into_response(serde_json::json!(42));
            framed.send(&response).await.unwrap();
        });

        let stream = TcpStream::connect(addr).await.unwrap();
        let mut framed = FramedStream::new(stream);

        let request = CommandMessage::read("test.value");
        let response = framed
            .send_recv(&request, std::time::Duration::from_secs(5))
            .await
            .unwrap();

        assert!(response.success);
        assert_eq!(response.data, serde_json::json!(42));

        server_task.await.unwrap();
    }

    #[tokio::test]
    async fn test_large_message_rejection() {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        let _server_task = tokio::spawn(async move {
            let (stream, _) = listener.accept().await.unwrap();
            let mut framed = FramedStream::new(stream);

            // Try to receive - should fail with framing error
            let result = framed.recv().await;
            assert!(result.is_err());
        });

        // Client sends invalid header with huge length
        let mut stream = TcpStream::connect(addr).await.unwrap();

        // Write a header with length > MAX_MESSAGE_SIZE
        let bad_header = IpcMessageHeader::new(MessageType::Request, MAX_MESSAGE_SIZE + 1);
        stream.write_all(&bad_header.to_bytes()).await.unwrap();
    }

    #[tokio::test]
    async fn test_split_transport() {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        let server_task = tokio::spawn(async move {
            let (stream, _) = listener.accept().await.unwrap();
            let transport = SplitTransport::new(stream);

            let msg = transport.recv().await.unwrap();
            assert_eq!(msg.topic, "split.test");

            let response = msg.into_response(serde_json::json!("pong"));
            transport.send(&response).await.unwrap();
        });

        let stream = TcpStream::connect(addr).await.unwrap();
        let transport = SplitTransport::new(stream);

        let request = CommandMessage::request("split.test", serde_json::json!("ping"));
        transport.send(&request).await.unwrap();

        let response = transport.recv().await.unwrap();
        assert!(response.success);

        server_task.await.unwrap();
    }
}