strike48-connector 0.3.6

//! WebSocket Transport for Strike48 Connector SDK.
//!
//! WebSocket over HTTP/1.1 - works through corporate proxies that block HTTP/2.
//! Uses Phoenix Channels with JSON transport and base64-encoded protobuf data.
//!
//! Protocol:
//! - Connect to /socket/connector/websocket?vsn=2.0.0
//! - Join "connector:lobby" channel
//! - Send/receive "proto" events with Base64-encoded protobuf data
//! - Phoenix message format: [join_ref, ref, topic, event, payload]
//! - Heartbeat: Client must send [null, ref, "phoenix", "heartbeat", {}] every 30s

use async_trait::async_trait;
use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64};
use futures_util::{SinkExt, StreamExt};
use prost::Message as ProstMessage;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use tokio::sync::mpsc;
use tokio::time::{Duration, timeout};
use tokio_tungstenite::{Connector, connect_async_tls_with_config, tungstenite::Message};
use tracing::{debug, error, trace, warn};

use super::{Transport, TransportOptions, TransportType, create_unbounded_wrapper};
use crate::error::{ConnectorError, Result};
use strike48_proto::proto::StreamMessage;

/// Default channel capacity for bounded message channels.
const DEFAULT_CHANNEL_CAPACITY: usize = 1024;

/// Strip the port from a `host:port` string when it's the well-known default
/// for the scheme (80 for HTTP/WS, 443 for HTTPS/WSS).
/// This keeps the HTTP Host header clean for hostname-based gateway routing.
fn strip_default_port(host_port: &str, use_tls: bool) -> String {
    let default_port: u16 = if use_tls { 443 } else { 80 };

    if let Some((host, port_str)) = host_port.rsplit_once(':')
        && let Ok(port) = port_str.parse::<u16>()
        && port == default_port
    {
        return host.to_string();
    }

    host_port.to_string()
}

/// Phoenix heartbeat interval in milliseconds.
/// Phoenix default is 30s timeout, we send every 25s to have margin.
const PHOENIX_HEARTBEAT_INTERVAL_MS: u64 = 25000;

/// WebSocket Transport implementation.
///
/// Uses Phoenix Channels over WebSocket for connector communication.
/// Messages are base64-encoded protobuf over JSON transport.
///
/// Endpoint: /socket/connector/websocket?vsn=2.0.0
/// Channel: connector:lobby
/// Event: "proto" with payload `{"data": "<base64-encoded-protobuf>"}`
pub struct WebSocketTransport {
    options: TransportOptions,
    connected: Arc<AtomicBool>,
    /// Channel capacity for backpressure (default: 1024)
    channel_capacity: usize,

    // Phoenix protocol state (atomic for lock-free access)
    join_ref: Arc<AtomicU32>,
    msg_ref: Arc<AtomicU32>,

    /// Handles for spawned background tasks so we can abort them on disconnect.
    task_handles: Vec<tokio::task::JoinHandle<()>>,
}

impl WebSocketTransport {
    /// Create a new WebSocket transport.
    pub fn new(options: TransportOptions) -> Self {
        debug!(
            "WebSocketTransport created: {} (TLS: {})",
            options.host, options.use_tls
        );

        Self {
            channel_capacity: options.channel_capacity.unwrap_or(DEFAULT_CHANNEL_CAPACITY),
            options,
            connected: Arc::new(AtomicBool::new(false)),
            join_ref: Arc::new(AtomicU32::new(0)),
            msg_ref: Arc::new(AtomicU32::new(0)),
            task_handles: Vec::new(),
        }
    }

    /// Build a TLS connector, optionally accepting invalid certificates.
    ///
    /// When `MATRIX_TLS_INSECURE=true` is set, the connector will accept
    /// self-signed and otherwise invalid TLS certificates. This is intended
    /// for local development with self-signed certs only.
    fn build_tls_connector(&self) -> Option<Connector> {
        if !self.options.use_tls {
            return None;
        }

        let tls_insecure = std::env::var("MATRIX_TLS_INSECURE")
            .map(|v| v == "true" || v == "1")
            .unwrap_or(false);

        if tls_insecure {
            warn!(
                "TLS certificate verification DISABLED (MATRIX_TLS_INSECURE=true). Do NOT use in production!"
            );
            let tls_connector = native_tls::TlsConnector::builder()
                .danger_accept_invalid_certs(true)
                .build()
                .expect("Failed to build TLS connector");
            Some(Connector::NativeTls(tls_connector))
        } else {
            // Use default system TLS (validates certificates)
            None
        }
    }

    /// Build WebSocket URL from options.
    ///
    /// Uses `self.options.host` as-is (already resolved by the URL parser).
    /// Omits the port when it matches the scheme default so that the Host header
    /// stays clean for hostname-based gateway routing (e.g. `*.example.com`).
    fn build_ws_url(&self) -> String {
        let scheme = if self.options.use_tls { "wss" } else { "ws" };

        let authority = strip_default_port(&self.options.host, self.options.use_tls);

        format!("{scheme}://{authority}/socket/connector/websocket?vsn=2.0.0")
    }

    /// Wait for Phoenix channel join confirmation with proper timeout.
    async fn wait_for_join_confirmation<S>(read: &mut S, timeout_ms: u64) -> Result<()>
    where
        S: StreamExt<Item = std::result::Result<Message, tokio_tungstenite::tungstenite::Error>>
            + Unpin,
    {
        let deadline = Duration::from_millis(timeout_ms);

        loop {
            match timeout(deadline, read.next()).await {
                Ok(Some(Ok(Message::Text(text)))) => {
                    // Parse Phoenix message: [join_ref, ref, topic, event, payload]
                    if let Ok(parsed) = serde_json::from_str::<Vec<serde_json::Value>>(&text)
                        && parsed.len() >= 5
                    {
                        let event = parsed[3].as_str().unwrap_or("");
                        let payload = &parsed[4];

                        if event == "phx_reply"
                            && let Some(status) = payload
                                .as_object()
                                .and_then(|p| p.get("status"))
                                .and_then(|s| s.as_str())
                        {
                            if status == "ok" {
                                debug!("Phoenix join confirmed");
                                return Ok(());
                            } else if status == "error" {
                                return Err(ConnectorError::ConnectionError(format!(
                                    "Phoenix join failed: {payload:?}"
                                )));
                            }
                        }
                    }
                }
                Ok(Some(Ok(_))) => {
                    // Ignore other message types during join
                    continue;
                }
                Ok(Some(Err(e))) => {
                    return Err(ConnectorError::ConnectionError(format!(
                        "WebSocket error during join: {e}"
                    )));
                }
                Ok(None) => {
                    return Err(ConnectorError::ConnectionError(
                        "WebSocket closed during join".to_string(),
                    ));
                }
                Err(_) => {
                    return Err(ConnectorError::Timeout(
                        "Timeout waiting for Phoenix join confirmation".to_string(),
                    ));
                }
            }
        }
    }

    /// Decode base64-encoded protobuf data to StreamMessage.
    fn decode_proto_message(data: &str) -> Option<StreamMessage> {
        // Try base64 decode first
        let proto_bytes = match BASE64.decode(data) {
            Ok(bytes) => bytes,
            Err(_) => {
                // Fallback to hex decode for backwards compatibility
                match hex::decode(data) {
                    Ok(bytes) => bytes,
                    Err(e) => {
                        error!("Failed to decode proto data: {}", e);
                        return None;
                    }
                }
            }
        };

        // Parse protobuf
        match StreamMessage::decode(proto_bytes.as_slice()) {
            Ok(msg) => Some(msg),
            Err(e) => {
                error!("Failed to parse protobuf message: {}", e);
                None
            }
        }
    }
}

#[async_trait]
impl Transport for WebSocketTransport {
    #[allow(dead_code)]
    fn transport_type(&self) -> TransportType {
        TransportType::WebSocket
    }

    async fn connect(&mut self) -> Result<()> {
        let url = self.build_ws_url();
        let connect_timeout = self.options.connect_timeout_ms.unwrap_or(10000);
        let tls_connector = self.build_tls_connector();

        debug!("Connecting to WebSocket at {}", url);

        // Connect with timeout (using custom TLS connector if configured)
        let (ws_stream, _) = timeout(
            Duration::from_millis(connect_timeout),
            connect_async_tls_with_config(&url, None, false, tls_connector),
        )
        .await
        .map_err(|_| ConnectorError::Timeout("WebSocket connection timeout".to_string()))?
        .map_err(|e| ConnectorError::ConnectionError(format!("WebSocket connect failed: {e}")))?;

        let (mut write, mut read) = ws_stream.split();

        // Join the Phoenix channel with proper confirmation
        let jr = self.join_ref.fetch_add(1, Ordering::SeqCst) + 1;
        let mr = self.msg_ref.fetch_add(1, Ordering::SeqCst) + 1;

        let join_msg = serde_json::json!([
            jr.to_string(),
            mr.to_string(),
            "connector:lobby",
            "phx_join",
            {}
        ]);

        debug!("Sending phx_join: {}", join_msg);
        write
            .send(Message::Text(join_msg.to_string().into()))
            .await
            .map_err(|e| ConnectorError::ConnectionError(format!("Failed to send join: {e}")))?;

        // Wait for proper join confirmation (not just a sleep!)
        Self::wait_for_join_confirmation(&mut read, 5000).await?;

        self.connected.store(true, Ordering::SeqCst);
        debug!("WebSocket transport connected and joined connector:lobby");

        Ok(())
    }

    async fn start_stream(
        &mut self,
        initial_message: Option<StreamMessage>,
    ) -> Result<(
        mpsc::UnboundedSender<StreamMessage>,
        mpsc::UnboundedReceiver<StreamMessage>,
    )> {
        let url = self.build_ws_url();
        let connect_timeout_ms = self.options.connect_timeout_ms.unwrap_or(10000);

        debug!("Starting WebSocket stream at {}", url);

        // Connect to WebSocket with the same timeout used in connect() so a
        // stalled TCP handshake (server mid-restart) doesn't block the reconnect
        // loop indefinitely.
        let tls_connector = self.build_tls_connector();
        let (ws_stream, _) = timeout(
            Duration::from_millis(connect_timeout_ms),
            connect_async_tls_with_config(&url, None, false, tls_connector),
        )
        .await
        .map_err(|_| ConnectorError::Timeout("WebSocket stream connection timeout".to_string()))?
        .map_err(|e| ConnectorError::ConnectionError(format!("WebSocket connect failed: {e}")))?;

        let (mut write, mut read) = ws_stream.split();

        // Create BOUNDED channels for backpressure
        let (proto_tx, mut proto_rx) = mpsc::channel::<StreamMessage>(self.channel_capacity);
        let (response_tx, response_rx) = mpsc::channel::<StreamMessage>(self.channel_capacity);

        let join_ref = self.join_ref.clone();
        let msg_ref = self.msg_ref.clone();
        let connected = self.connected.clone();

        // Join the channel first and wait for confirmation
        let jr = join_ref.fetch_add(1, Ordering::SeqCst) + 1;
        let mr = msg_ref.fetch_add(1, Ordering::SeqCst) + 1;

        let join_msg = serde_json::json!([
            jr.to_string(),
            mr.to_string(),
            "connector:lobby",
            "phx_join",
            {}
        ]);

        debug!("Sending phx_join: {}", join_msg);
        write
            .send(Message::Text(join_msg.to_string().into()))
            .await
            .map_err(|e| ConnectorError::ConnectionError(format!("Failed to send join: {e}")))?;

        // Wait for join confirmation properly
        Self::wait_for_join_confirmation(&mut read, 5000).await?;

        // Create internal channel for WebSocket writes (proto + heartbeat)
        // This allows multiple senders (proto forwarder, heartbeat) to write to the WebSocket
        let (ws_tx, mut ws_rx) = mpsc::channel::<String>(64);

        // Send initial message if provided
        if let Some(msg) = initial_message {
            debug!("WS TX: Sending initial message");
            let proto_bytes = msg.encode_to_vec();
            let b64_data = BASE64.encode(&proto_bytes);

            let mr = msg_ref.fetch_add(1, Ordering::SeqCst) + 1;
            let jr = join_ref.load(Ordering::SeqCst);

            let phoenix_msg = serde_json::json!([
                jr.to_string(),
                mr.to_string(),
                "connector:lobby",
                "proto",
                {"data": b64_data}
            ]);

            write
                .send(Message::Text(phoenix_msg.to_string().into()))
                .await
                .map_err(|e| {
                    ConnectorError::StreamError(format!("Failed to send initial message: {e}"))
                })?;
        }

        // Abort any lingering tasks from a previous stream before spawning new ones.
        self.abort_tasks();

        // Spawn WebSocket writer task - reads from internal channel and sends to WebSocket
        let connected_writer = connected.clone();
        let writer_handle = tokio::spawn(async move {
            while let Some(msg_text) = ws_rx.recv().await {
                if let Err(e) = write.send(Message::Text(msg_text.into())).await {
                    error!("WebSocket write failed: {}", e);
                    connected_writer.store(false, Ordering::SeqCst);
                    break;
                }
            }
            debug!("WebSocket writer task ended");
        });

        // Spawn proto forwarder task - converts proto messages to Phoenix format
        let join_ref_proto = join_ref.clone();
        let msg_ref_proto = msg_ref.clone();
        let ws_tx_proto = ws_tx.clone();
        let forwarder_handle = tokio::spawn(async move {
            while let Some(msg) = proto_rx.recv().await {
                // Encode message to protobuf binary
                let proto_bytes = msg.encode_to_vec();
                let b64_data = BASE64.encode(&proto_bytes);

                // Increment ref atomically
                let mr = msg_ref_proto.fetch_add(1, Ordering::SeqCst) + 1;
                let jr = join_ref_proto.load(Ordering::SeqCst);

                // Phoenix message: [join_ref, ref, topic, event, payload]
                let phoenix_msg = serde_json::json!([
                    jr.to_string(),
                    mr.to_string(),
                    "connector:lobby",
                    "proto",
                    {"data": b64_data}
                ]);

                trace!("WS TX proto ({} bytes)", proto_bytes.len());
                if ws_tx_proto.send(phoenix_msg.to_string()).await.is_err() {
                    debug!("Proto forwarder: ws_tx closed");
                    break;
                }
            }
            debug!("Proto forwarder task ended");
        });

        // Spawn Phoenix heartbeat task - sends heartbeat every 25 seconds
        let msg_ref_heartbeat = msg_ref.clone();
        let ws_tx_heartbeat = ws_tx.clone();
        let connected_heartbeat = connected.clone();
        let heartbeat_handle = tokio::spawn(async move {
            let mut interval =
                tokio::time::interval(Duration::from_millis(PHOENIX_HEARTBEAT_INTERVAL_MS));
            interval.tick().await; // Skip first immediate tick

            loop {
                interval.tick().await;

                if !connected_heartbeat.load(Ordering::SeqCst) {
                    trace!("Heartbeat task: not connected, stopping");
                    break;
                }

                let mr = msg_ref_heartbeat.fetch_add(1, Ordering::SeqCst) + 1;

                // Phoenix heartbeat: [null, ref, "phoenix", "heartbeat", {}]
                let heartbeat = serde_json::json!([
                    serde_json::Value::Null,
                    mr.to_string(),
                    "phoenix",
                    "heartbeat",
                    {}
                ]);

                trace!("Sending Phoenix heartbeat (ref: {})", mr);
                if ws_tx_heartbeat.send(heartbeat.to_string()).await.is_err() {
                    trace!("Heartbeat task: ws_tx closed");
                    break;
                }
            }
            trace!("Heartbeat task ended");
        });

        // Spawn reader task - parses Phoenix format to proto messages
        let ws_tx_reader = ws_tx; // Move remaining ws_tx to reader for heartbeat replies
        let reader_handle = tokio::spawn(async move {
            while let Some(msg) = read.next().await {
                match msg {
                    Ok(Message::Text(text)) => {
                        trace!("WS RX: {}", text);

                        // Parse Phoenix message: [join_ref, ref, topic, event, payload]
                        if let Ok(parsed) = serde_json::from_str::<Vec<serde_json::Value>>(&text)
                            && parsed.len() >= 5
                        {
                            let topic = parsed[2].as_str().unwrap_or("");
                            let event = parsed[3].as_str().unwrap_or("");
                            let payload = &parsed[4];

                            match event {
                                // Handle server-initiated heartbeat - respond immediately
                                "heartbeat" if topic == "phoenix" => {
                                    trace!("Responding to server heartbeat");
                                    let join_ref = &parsed[0];
                                    let ref_val = &parsed[1];

                                    let reply = serde_json::json!([
                                        join_ref,
                                        ref_val,
                                        "phoenix",
                                        "phx_reply",
                                        {"status": "ok", "response": {}}
                                    ]);

                                    if ws_tx_reader.send(reply.to_string()).await.is_err() {
                                        trace!("Reader: ws_tx closed");
                                        break;
                                    }
                                }
                                "phx_reply" => {
                                    if let Some(payload_obj) = payload.as_object() {
                                        let status = payload_obj
                                            .get("status")
                                            .and_then(|s| s.as_str())
                                            .unwrap_or("");

                                        if status == "ok" {
                                            trace!("Phoenix reply OK");

                                            // Check if reply contains proto data
                                            if let Some(response) = payload_obj.get("response")
                                                && let Some(data) = response
                                                    .as_object()
                                                    .and_then(|r| r.get("data"))
                                                    .and_then(|d| d.as_str())
                                                && let Some(proto_msg) =
                                                    WebSocketTransport::decode_proto_message(data)
                                            {
                                                trace!("Decoded proto from phx_reply");
                                                // Use blocking send for backpressure
                                                if response_tx.send(proto_msg).await.is_err() {
                                                    warn!("Response channel closed");
                                                    break;
                                                }
                                            }
                                        } else if status == "error" {
                                            let is_channel_dead = payload_obj
                                                .get("response")
                                                .and_then(|r| r.as_object())
                                                .and_then(|r| r.get("reason"))
                                                .and_then(|r| r.as_str())
                                                .is_some_and(|reason| reason == "unmatched topic");

                                            if is_channel_dead {
                                                warn!(
                                                    "Channel process no longer exists \
                                                     (unmatched topic). Closing stream \
                                                     to trigger reconnect."
                                                );
                                                connected.store(false, Ordering::SeqCst);
                                                break;
                                            }
                                            error!("Phoenix reply error: {:?}", payload_obj);
                                        }
                                    }
                                }
                                "proto" => {
                                    // Incoming proto message from server
                                    if let Some(data) = payload
                                        .as_object()
                                        .and_then(|p| p.get("data"))
                                        .and_then(|d| d.as_str())
                                        && let Some(proto_msg) =
                                            WebSocketTransport::decode_proto_message(data)
                                    {
                                        trace!("Received proto message via 'proto' event");
                                        // Use blocking send for backpressure
                                        if response_tx.send(proto_msg).await.is_err() {
                                            warn!("Response channel closed");
                                            break;
                                        }
                                    }
                                }
                                "phx_error" => {
                                    warn!(
                                        "Phoenix channel process terminated (phx_error): {:?}. \
                                         Closing stream to trigger reconnect.",
                                        payload
                                    );
                                    connected.store(false, Ordering::SeqCst);
                                    break;
                                }
                                "phx_close" => {
                                    warn!("Phoenix channel closed");
                                    connected.store(false, Ordering::SeqCst);
                                    break;
                                }
                                _ => {
                                    debug!("Unhandled Phoenix event: {}", event);
                                }
                            }
                        }
                    }
                    Ok(Message::Binary(data)) => {
                        trace!("Received binary message ({} bytes)", data.len());
                        // Try to decode as raw protobuf
                        if let Ok(proto_msg) = StreamMessage::decode(&data[..])
                            && response_tx.send(proto_msg).await.is_err()
                        {
                            warn!("Response channel closed");
                            break;
                        }
                    }
                    Ok(Message::Close(frame)) => {
                        if let Some(cf) = &frame {
                            warn!(
                                "WebSocket closed by server: code={}, reason={}",
                                cf.code, cf.reason
                            );
                        } else {
                            debug!("WebSocket closed (no close frame)");
                        }
                        connected.store(false, Ordering::SeqCst);
                        break;
                    }
                    Ok(Message::Ping(_)) => {
                        trace!("Received ping");
                    }
                    Ok(Message::Pong(_)) => {
                        trace!("Received pong");
                    }
                    Ok(Message::Frame(_)) => {
                        // Raw frame, ignore
                    }
                    Err(e) => {
                        error!("WebSocket error: {}", e);
                        connected.store(false, Ordering::SeqCst);
                        break;
                    }
                }
            }
            debug!("WebSocket reader task ended");
        });

        self.task_handles = vec![
            writer_handle,
            forwarder_handle,
            heartbeat_handle,
            reader_handle,
        ];

        debug!("WebSocket bidirectional stream started");

        let (unbounded_tx, unbounded_rx, wrapper_handles) =
            create_unbounded_wrapper(proto_tx, response_rx);
        self.task_handles.extend(wrapper_handles);

        Ok((unbounded_tx, unbounded_rx))
    }

    #[allow(dead_code)]
    fn is_connected(&self) -> bool {
        self.connected.load(Ordering::SeqCst)
    }

    async fn disconnect(&mut self) -> Result<()> {
        debug!("Disconnecting WebSocket transport");
        self.connected.store(false, Ordering::SeqCst);
        self.abort_tasks();
        Ok(())
    }
}

impl WebSocketTransport {
    /// Abort all spawned background tasks (writer, forwarder, heartbeat, reader).
    /// Called on disconnect and before starting a new stream to prevent leaked tasks
    /// from a previous connection from spamming "Response channel closed" warnings.
    fn abort_tasks(&mut self) {
        for handle in self.task_handles.drain(..) {
            handle.abort();
        }
    }
}

impl Drop for WebSocketTransport {
    fn drop(&mut self) {
        self.abort_tasks();
    }
}