takproto 0.4.2

use crate::error::{Result, TakError};
use crate::framing::{decode_tak_header, encode_tak_message};
use crate::proto::{CotEvent, TakMessage};
use crate::tls::TlsConfig;
use bytes::BytesMut;
use prost::Message;
use rustls::pki_types::ServerName;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio_rustls::TlsConnector;

/// Connection type for TAK client
enum Connection {
    Plain(TcpStream),
    Tls(tokio_rustls::client::TlsStream<TcpStream>),
}

impl Connection {
    async fn read_buf(&mut self, buf: &mut BytesMut) -> std::io::Result<usize> {
        match self {
            Connection::Plain(stream) => stream.read_buf(buf).await,
            Connection::Tls(stream) => stream.read_buf(buf).await,
        }
    }

    async fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
        match self {
            Connection::Plain(stream) => stream.write_all(buf).await,
            Connection::Tls(stream) => stream.write_all(buf).await,
        }
    }

    async fn flush(&mut self) -> std::io::Result<()> {
        match self {
            Connection::Plain(stream) => stream.flush().await,
            Connection::Tls(stream) => stream.flush().await,
        }
    }

    async fn shutdown(&mut self) -> std::io::Result<()> {
        match self {
            Connection::Plain(stream) => stream.shutdown().await,
            Connection::Tls(stream) => stream.shutdown().await,
        }
    }
}

/// TAK Protocol client for streaming connections to TAK servers
///
/// This client implements TAK Protocol Version 1 using Protocol Buffers.
/// It handles the framing protocol for streaming connections and supports
/// both plain TCP and TLS (including mTLS with client certificates).
///
/// # Example with mTLS
///
/// ```no_run
/// use takproto::{TakClient, TlsConfig, proto::CotEvent};
///
/// #[tokio::main]
/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
///     let tls_config = TlsConfig::new_with_client_cert(
///         "ca.pem",
///         "client.pem",
///         "client-key.pem"
///     )?;
///
///     let mut client = TakClient::connect_tls(
///         "takserver.example.com:8089",
///         "takserver.example.com",
///         tls_config
///     ).await?;
///
///     let event = CotEvent {
///         r#type: "a-f-G-U-C".to_string(),
///         uid: "RUST-1".to_string(),
///         // ... other fields
///         ..Default::default()
///     };
///
///     client.send_cot_event(event).await?;
///     Ok(())
/// }
/// ```
pub struct TakClient {
    connection: Connection,
    read_buffer: BytesMut,
}

impl TakClient {
    /// Connect to a TAK server at the specified address using plain TCP
    ///
    /// Note: Most production TAK servers require TLS. Use `connect_tls` instead.
    ///
    /// # Arguments
    ///
    /// * `addr` - The address of the TAK server (e.g., "127.0.0.1:8087")
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use takproto::TakClient;
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let client = TakClient::connect("127.0.0.1:8087").await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn connect(addr: impl tokio::net::ToSocketAddrs) -> Result<Self> {
        let stream = TcpStream::connect(addr).await?;
        Ok(Self {
            connection: Connection::Plain(stream),
            read_buffer: BytesMut::with_capacity(8192),
        })
    }

    /// Connect to a TAK server using TLS with optional client certificate authentication (mTLS)
    ///
    /// # Arguments
    ///
    /// * `addr` - The address of the TAK server (e.g., "takserver.example.com:8089")
    /// * `server_name` - The server name for SNI and certificate validation
    /// * `tls_config` - TLS configuration including optional client certificates
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use takproto::{TakClient, TlsConfig};
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let tls_config = TlsConfig::new_with_client_cert(
    ///     "ca.pem",
    ///     "client.pem",
    ///     "client-key.pem"
    /// )?;
    ///
    /// let client = TakClient::connect_tls(
    ///     "takserver.example.com:8089",
    ///     "takserver.example.com",
    ///     tls_config
    /// ).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn connect_tls(
        addr: impl tokio::net::ToSocketAddrs,
        server_name: &str,
        tls_config: TlsConfig,
    ) -> Result<Self> {
        let tcp_stream = TcpStream::connect(addr).await?;

        let connector = TlsConnector::from(tls_config.config);

        let server_name = ServerName::try_from(server_name.to_owned())
            .map_err(|e| TakError::Tls(format!("Invalid server name: {}", e)))?;

        let tls_stream = connector.connect(server_name, tcp_stream).await?;

        Ok(Self {
            connection: Connection::Tls(tls_stream),
            read_buffer: BytesMut::with_capacity(8192),
        })
    }

    /// Send a CoT event to the TAK server
    ///
    /// This wraps the CotEvent in a TakMessage and sends it using the
    /// TAK Protocol streaming message format.
    ///
    /// # Arguments
    ///
    /// * `event` - The CoT event to send
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use takproto::{TakClient, proto::CotEvent};
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// # let mut client = TakClient::connect("127.0.0.1:8087").await?;
    /// let event = CotEvent {
    ///     r#type: "a-f-G-U-C".to_string(),
    ///     uid: "RUST-1".to_string(),
    ///     send_time: 1234567890000,
    ///     start_time: 1234567890000,
    ///     stale_time: 1234567950000,
    ///     how: "m-g".to_string(),
    ///     lat: 37.7749,
    ///     lon: -122.4194,
    ///     hae: 10.0,
    ///     ce: 9.9,
    ///     le: 9.9,
    ///     ..Default::default()
    /// };
    /// client.send_cot_event(event).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn send_cot_event(&mut self, event: CotEvent) -> Result<()> {
        let tak_message = TakMessage {
            tak_control: None,
            cot_event: Some(event),
        };
        self.send_tak_message(tak_message).await
    }

    /// Send a raw TakMessage to the server
    ///
    /// This is a lower-level method that allows sending any TakMessage,
    /// including those with TakControl fields.
    ///
    /// # Arguments
    ///
    /// * `message` - The TakMessage to send
    pub async fn send_tak_message(&mut self, message: TakMessage) -> Result<()> {
        let frame = encode_tak_message(&message)?;
        self.connection.write_all(&frame).await?;
        self.connection.flush().await?;
        Ok(())
    }

    /// Receive a TakMessage from the server
    ///
    /// This reads from the stream and decodes the next TAK Protocol message.
    ///
    /// # Returns
    ///
    /// * `Ok(Some(message))` - A message was received
    /// * `Ok(None)` - Connection was closed gracefully
    /// * `Err(e)` - An error occurred
    pub async fn receive_tak_message(&mut self) -> Result<Option<TakMessage>> {
        loop {
            // Try to decode a message from the buffer
            if let Some(message) = self.try_decode_message()? {
                return Ok(Some(message));
            }

            // Read more data from the stream
            let n = self.connection.read_buf(&mut self.read_buffer).await?;

            if n == 0 {
                // Connection closed
                if self.read_buffer.is_empty() {
                    return Ok(None);
                } else {
                    return Err(TakError::ConnectionClosed);
                }
            }
        }
    }

    /// Try to decode a message from the read buffer
    fn try_decode_message(&mut self) -> Result<Option<TakMessage>> {
        // Try to decode the header
        let payload_len = match decode_tak_header(&mut self.read_buffer)? {
            Some(len) => len,
            None => return Ok(None), // Need more data
        };

        // Check if we have the complete payload
        if self.read_buffer.len() < payload_len {
            return Ok(None); // Need more data
        }

        // Split off the payload
        let payload = self.read_buffer.split_to(payload_len);

        // Decode the protobuf message
        let message = TakMessage::decode(&payload[..])?;

        Ok(Some(message))
    }

    /// Send a CoT event as XML (Protocol Version 0)
    ///
    /// This is used for initial communication before protocol negotiation,
    /// or when the server doesn't support protobuf.
    pub async fn send_cot_event_xml(&mut self, event: CotEvent) -> Result<()> {
        let xml = crate::xml::encode_cot_event_xml(&event);
        self.connection.write_all(xml.as_bytes()).await?;
        self.connection.flush().await?;
        Ok(())
    }

    /// Negotiate protocol version with TAK server
    ///
    /// This performs the full protocol negotiation handshake:
    /// 1. Waits for server's TakProtocolSupport message
    /// 2. Sends TakRequest for the specified version
    /// 3. Waits for TakResponse confirmation
    ///
    /// After successful negotiation, the connection is ready for protobuf messages.
    ///
    /// # Arguments
    ///
    /// * `version` - Protocol version to negotiate (typically 1)
    /// * `timeout` - Maximum time to wait for negotiation (in seconds)
    ///
    /// # Returns
    ///
    /// Returns `Ok(())` if negotiation succeeds, or an error if it fails or times out.
    pub async fn negotiate_protocol(&mut self, version: u32, timeout_secs: u64) -> Result<()> {
        use tokio::time::{timeout, Duration};

        // Wait for server's protocol support announcement
        let xml_buffer = timeout(
            Duration::from_secs(timeout_secs),
            self.read_xml_messages_until(|msg| crate::xml::is_protocol_support(msg)),
        )
        .await
        .map_err(|_| TakError::NegotiationFailed("Timeout waiting for protocol support".to_string()))??;

        if xml_buffer.is_empty() {
            return Err(TakError::NegotiationFailed(
                "Server did not advertise protocol support".to_string(),
            ));
        }

        // Send protocol request
        let request_xml = crate::xml::create_protocol_request(version);
        self.connection.write_all(request_xml.as_bytes()).await?;
        self.connection.flush().await?;

        // Wait for server's response
        let response_xml = timeout(
            Duration::from_secs(timeout_secs),
            self.read_xml_messages_until(|msg| crate::xml::is_protocol_response(msg)),
        )
        .await
        .map_err(|_| TakError::NegotiationFailed("Timeout waiting for protocol response".to_string()))??;

        if response_xml.is_empty() {
            return Err(TakError::NegotiationFailed(
                "No protocol response received".to_string(),
            ));
        }

        // Check if response is successful
        if !crate::xml::is_protocol_response_success(&response_xml) {
            return Err(TakError::NegotiationFailed(
                "Server rejected protocol request".to_string(),
            ));
        }

        Ok(())
    }

    /// Read XML messages from the connection until a condition is met
    async fn read_xml_messages_until<F>(&mut self, mut condition: F) -> Result<String>
    where
        F: FnMut(&str) -> bool,
    {
        let mut buffer = String::new();
        let mut temp_buf = BytesMut::with_capacity(4096);

        loop {
            // Read more data
            temp_buf.clear();
            let n = self.connection.read_buf(&mut temp_buf).await?;

            if n == 0 {
                return Err(TakError::ConnectionClosed);
            }

            // Append to buffer
            buffer.push_str(&String::from_utf8_lossy(&temp_buf[..]));

            // Check for complete XML message(s)
            while let Some(end_pos) = buffer.find("</event>") {
                let message = &buffer[..end_pos + 8]; // Include </event>
                if condition(message) {
                    return Ok(message.to_string());
                }
                // Remove the processed message from buffer
                buffer = buffer[end_pos + 8..].to_string();
            }
        }
    }

    /// Close the connection to the TAK server
    pub async fn close(mut self) -> Result<()> {
        self.connection.shutdown().await?;
        Ok(())
    }
}

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

    #[test]
    fn test_cot_event_creation() {
        let event = CotEvent {
            r#type: "a-f-G-U-C".to_string(),
            uid: "TEST-1".to_string(),
            send_time: 1000,
            start_time: 1000,
            stale_time: 2000,
            how: "m-g".to_string(),
            lat: 37.7749,
            lon: -122.4194,
            hae: 10.0,
            ce: 9.9,
            le: 9.9,
            ..Default::default()
        };

        assert_eq!(event.r#type, "a-f-G-U-C");
        assert_eq!(event.uid, "TEST-1");
    }
}