pixeluvw_supabase/

realtime.rs

use crate::core::SupabaseClient;
use crate::error::{Result, SupaError};
use futures_util::{SinkExt, StreamExt};
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use serde_json::{json, Value};

use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;
use tokio::sync::mpsc;
use tokio::time::sleep;
use tokio_stream::Stream;
use tokio_tungstenite::tungstenite::Message;

// ============================================================================
//  RealtimeClient
// ============================================================================

#[derive(Clone)]
pub struct RealtimeClient {
    pub(crate) client: SupabaseClient,
}

impl RealtimeClient {
    pub(crate) fn new(client: SupabaseClient) -> Self {
        Self { client }
    }

    /// Create a channel builder to configure a new subscription.
    pub fn channel(&self, topic: &str) -> RealtimeChannelBuilder {
        RealtimeChannelBuilder::new(self.client.clone(), topic)
    }
}

// ============================================================================
//  Realtime Types
// ============================================================================

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PostgresEvent {
    Insert,
    Update,
    Delete,
    All,
}

impl ToString for PostgresEvent {
    fn to_string(&self) -> String {
        match self {
            PostgresEvent::Insert => "INSERT".to_string(),
            PostgresEvent::Update => "UPDATE".to_string(),
            PostgresEvent::Delete => "DELETE".to_string(),
            PostgresEvent::All => "*".to_string(),
        }
    }
}

/// Connection state for realtime channels.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ConnectionState {
    /// Channel is connecting to the server.
    Connecting,
    /// Channel is connected and receiving messages.
    Connected,
    /// Channel is reconnecting after a disconnection.
    Reconnecting,
    /// Channel has been closed.
    Closed,
}

/// Commands sent from the user handle to the connection loop.
enum ChannelCommand {
    Broadcast {
        event: String,
        payload: Value,
    },
    Track {
        payload: Value,
    },
    Untrack,
    /// Close the channel gracefully.
    Close,
}

// ============================================================================
//  RealtimeChannel (Handle)
// ============================================================================

/// A handle to a subscribed Realtime channel.
///
/// Implements `Stream` to receive messages, and provides methods to send broadcasts or track presence.
pub struct RealtimeChannel {
    topic: String,
    rx: mpsc::UnboundedReceiver<Result<RealtimeMessage>>,
    cmd_tx: mpsc::UnboundedSender<ChannelCommand>,
}

impl Stream for RealtimeChannel {
    type Item = Result<RealtimeMessage>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.rx.poll_recv(cx)
    }
}

impl RealtimeChannel {
    /// Get the topic name of this channel.
    pub fn topic(&self) -> &str {
        &self.topic
    }

    /// Send a broadcast message to other clients in the channel.
    pub fn broadcast(&self, event: &str, payload: Value) -> Result<()> {
        self.cmd_tx
            .send(ChannelCommand::Broadcast {
                event: event.to_string(),
                payload,
            })
            .map_err(|_| SupaError::RealtimeError {
                message: "Channel closed".to_string(),
            })
    }

    /// Track user presence.
    pub fn track(&self, payload: Value) -> Result<()> {
        self.cmd_tx
            .send(ChannelCommand::Track { payload })
            .map_err(|_| SupaError::RealtimeError {
                message: "Channel closed".to_string(),
            })
    }

    /// Untrack user presence.
    pub fn untrack(&self) -> Result<()> {
        self.cmd_tx
            .send(ChannelCommand::Untrack)
            .map_err(|_| SupaError::RealtimeError {
                message: "Channel closed".to_string(),
            })
    }

    /// Close the channel gracefully.
    ///
    /// This will disconnect the WebSocket connection and stop receiving messages.
    /// The channel cannot be reused after closing.
    pub fn close(&self) -> Result<()> {
        self.cmd_tx
            .send(ChannelCommand::Close)
            .map_err(|_| SupaError::RealtimeError {
                message: "Channel already closed".to_string(),
            })
    }
}

// ============================================================================
//  RealtimeChannelBuilder
// ============================================================================

pub struct RealtimeChannelBuilder {
    client: SupabaseClient,
    topic: String,
    postgres_changes: Vec<Value>,
}

impl RealtimeChannelBuilder {
    pub fn new(client: SupabaseClient, topic: &str) -> Self {
        Self {
            client,
            topic: topic.to_string(),
            postgres_changes: Vec::new(),
        }
    }

    /// Listen to Postgres Changes
    pub fn on_postgres_changes<S1, S2, S3>(
        mut self,
        event: PostgresEvent,
        schema: S1,
        table: Option<S2>,
        filter: Option<S3>,
    ) -> Self
    where
        S1: Into<String>,
        S2: Into<String>,
        S3: Into<String>,
    {
        let mut config = json!({
            "event": event.to_string(),
            "schema": schema.into(),
        });

        if let Some(t) = table {
            config
                .as_object_mut()
                .unwrap()
                .insert("table".to_string(), json!(t.into()));
        }
        if let Some(f) = filter {
            config
                .as_object_mut()
                .unwrap()
                .insert("filter".to_string(), json!(f.into()));
        }

        self.postgres_changes.push(config);
        self
    }

    /// Subscribe to the channel.
    /// Returns a `RealtimeChannel` handle.
    pub async fn subscribe(self) -> Result<RealtimeChannel> {
        let (tx, rx) = mpsc::unbounded_channel();
        let (cmd_tx, mut cmd_rx) = mpsc::unbounded_channel();

        let client = self.client.clone();
        let topic = self.topic.clone();

        // Build the config payload for phx_join
        let mut postgres_changes_config = Vec::new();
        for cfg in &self.postgres_changes {
            postgres_changes_config.push(json!({
                "event": cfg["event"],
                "schema": cfg["schema"],
                "table": cfg.get("table"),
                "filter": cfg.get("filter")
            }));
        }

        let mut config = json!({});
        if !postgres_changes_config.is_empty() {
            config.as_object_mut().unwrap().insert(
                "postgres_changes".to_string(),
                json!(postgres_changes_config),
            );
        }

        config.as_object_mut().unwrap().insert(
            "broadcast".to_string(),
            json!({ "ack": false, "self": false }),
        );
        config
            .as_object_mut()
            .unwrap()
            .insert("presence".to_string(), json!({ "key": "" }));

        let config_clone = config.clone();

        tokio::spawn(async move {
            let mut retry_count = 0;
            let base_delay = client.inner.config.retry_base_delay_ms;

            loop {
                // Connection Loop
                // We pass &mut cmd_rx to the connection function
                match connect_and_listen(&client, &topic, &config_clone, &tx, &mut cmd_rx).await {
                    Ok(_) => {
                        retry_count = 0;
                    }
                    Err(e) => {
                        let _ = tx.send(Err(SupaError::RealtimeError {
                            message: format!("Realtime disconnected: {}. Reconnecting...", e),
                        }));
                    }
                }

                retry_count += 1;
                let delay = base_delay * 2u64.pow(retry_count.min(9) as u32);
                sleep(Duration::from_millis(delay)).await;
            }
        });

        Ok(RealtimeChannel {
            topic: self.topic,
            rx,
            cmd_tx,
        })
    }
}

// ============================================================================
//  Internal Connection Logic
// ============================================================================

async fn connect_and_listen(
    client: &SupabaseClient,
    topic: &str,
    config: &Value,
    tx: &mpsc::UnboundedSender<Result<RealtimeMessage>>,
    user_cmd_rx: &mut mpsc::UnboundedReceiver<ChannelCommand>,
) -> Result<()> {
    // 1. WebSocket Handshake
    let url = client.inner.url.clone();
    let scheme = match url.scheme() {
        "https" => "wss",
        "http" => "ws",
        _ => "wss",
    };
    let host = url.host_str().unwrap_or_default();
    let port = url.port_or_known_default().unwrap_or(443);

    let ws_url = format!(
        "{}://{}:{}/realtime/v1/websocket?apikey={}&vsn=1.0.0",
        scheme, host, port, client.inner.key
    );

    let (ws_stream, _) = tokio_tungstenite::connect_async(&ws_url)
        .await
        .map_err(|e| SupaError::RealtimeError {
            message: format!("Connection failed: {}", e),
        })?;

    let (mut write, mut read) = ws_stream.split();
    let (internal_cmd_tx, mut internal_cmd_rx) = mpsc::channel::<Message>(10);

    // 2. Writer Task (Proxies messages from internal loop to WebSocket)
    let writer_handle = tokio::spawn(async move {
        while let Some(msg) = internal_cmd_rx.recv().await {
            if let Err(_) = write.send(msg).await {
                break;
            }
        }
    });

    // 3. Join Channel (phx_join)
    let join_ref = format!("{}", rand::random::<u64>());
    let access_token = {
        let lock = client.inner.session.read().unwrap();
        // Fallback to anon key if no session (Realtime usually uses public key or user token, logic varies)
        // If we want RLS, we need user token. If public, anon key.
        // Let's use get_access_token (internal clone) logic? No it's async and we are in async context but wrapped.
        // Just read from session or fallback to key.
        lock.as_ref()
            .map(|s| s.access_token.clone())
            .unwrap_or_else(|| client.inner.key.clone())
    };

    let join_msg = json!({
        "topic": topic,
        "event": "phx_join",
        "payload": {
            "config": config,
            "access_token": access_token
        },
        "ref": join_ref
    });

    internal_cmd_tx
        .send(Message::Text(join_msg.to_string()))
        .await
        .map_err(|e| SupaError::RealtimeError {
            message: format!("Failed to send join: {}", e),
        })?;

    // 4. Heartbeat Task
    let hb_cmd_tx = internal_cmd_tx.clone();
    let hb_handle = tokio::spawn(async move {
        loop {
            sleep(Duration::from_secs(30)).await;
            let msg = json!({
                "topic": "phoenix",
                "event": "heartbeat",
                "payload": {},
                "ref": format!("{}", rand::random::<u64>())
            });
            if hb_cmd_tx
                .send(Message::Text(msg.to_string()))
                .await
                .is_err()
            {
                break;
            }
        }
    });

    // 5. Main Select Loop
    loop {
        tokio::select! {
            // Incoming WebSocket Message
            msg_res = read.next() => {
                match msg_res {
                    Some(Ok(msg)) => {
                        match msg {
                            Message::Text(text) => {
                                if let Ok(parsed) = serde_json::from_str::<RealtimeMessage>(&text) {
                                    if parsed.event == "phx_reply" {
                                        // Heartbeat reply or join reply, ignore for now
                                        continue;
                                    }
                                    if parsed.event == "phx_close" {
                                        // Server closed channel
                                        break;
                                    }
                                    if parsed.event == "phx_error" {
                                         // Error
                                         break;
                                    }
                                    if tx.send(Ok(parsed)).is_err() {
                                        break;
                                    }
                                }
                            }
                            Message::Close(_) => break,
                            _ => {}
                        }
                    }
                    Some(Err(_)) => break, // WS Error
                    None => break, // Stream ended
                }
            }

            // Outgoing User Command (Broadcast / Presence)
            cmd = user_cmd_rx.recv() => {
                match cmd {
                    Some(ChannelCommand::Broadcast { event, payload }) => {
                        let msg = json!({
                            "topic": topic,
                            "event": "broadcast",
                            "payload": {
                                "event": event,
                                "payload": payload
                            },
                            "ref": format!("{}", rand::random::<u64>())
                        });
                        if internal_cmd_tx.send(Message::Text(msg.to_string())).await.is_err() {
                             break;
                        }
                    }
                    Some(ChannelCommand::Track { payload }) => {
                         let msg = json!({
                            "topic": topic,
                            "event": "presence",
                            "payload": {
                                "type": "track",
                                "event": "track",
                                "payload": payload
                            },
                             "ref": format!("{}", rand::random::<u64>())
                        });
                        if internal_cmd_tx.send(Message::Text(msg.to_string())).await.is_err() {
                             break;
                        }
                    }
                     Some(ChannelCommand::Untrack) => {
                         let msg = json!({
                            "topic": topic,
                            "event": "presence",
                            "payload": {
                                "type": "untrack",
                                "event": "untrack"
                            },
                             "ref": format!("{}", rand::random::<u64>())
                        });
                        if internal_cmd_tx.send(Message::Text(msg.to_string())).await.is_err() {
                             break;
                        }
                    }
                    Some(ChannelCommand::Close) => {
                        // Send phx_leave to cleanly disconnect
                        let leave_msg = json!({
                            "topic": topic,
                            "event": "phx_leave",
                            "payload": {},
                            "ref": format!("{}", rand::random::<u64>())
                        });
                        let _ = internal_cmd_tx.send(Message::Text(leave_msg.to_string())).await;
                        // Return Ok to signal intentional close (not error)
                        return Ok(());
                    }
                    None => break // User dropped channel handle
                }
            }
        }
    }

    // Cleanup
    hb_handle.abort();
    writer_handle.abort();

    Err(SupaError::RealtimeError {
        message: "Connection ended".into(),
    })
}

// ============================================================================
//  Message Types
// ============================================================================

#[derive(Debug, Serialize, Deserialize)]
pub struct RealtimeMessage {
    pub topic: String,
    pub event: String,
    pub payload: Value,
    #[serde(rename = "ref")]
    pub ref_: Option<String>,
}

impl RealtimeMessage {
    /// Check if this is a Postgres Change event
    pub fn is_postgres_change(&self) -> bool {
        self.event == "postgres_changes"
            || self.event == "INSERT"
            || self.event == "UPDATE"
            || self.event == "DELETE"
    }

    /// Check if this is a Presence event
    pub fn is_presence(&self) -> bool {
        self.event == "presence_state" || self.event == "presence_diff"
    }

    /// Check if this is a Broadcast event
    pub fn is_broadcast(&self) -> bool {
        self.event == "broadcast"
    }

    /// Parse as Postgres change record
    pub fn as_insert<T: DeserializeOwned>(&self) -> Result<T> {
        self.extract_record("INSERT")
    }

    pub fn as_update<T: DeserializeOwned>(&self) -> Result<T> {
        self.extract_record("UPDATE")
    }

    pub fn as_delete<T: DeserializeOwned>(&self) -> Result<T> {
        self.extract_record("DELETE")
    }

    // Helper to extract record from payload, handling Supabase's wrapper structure
    fn extract_record<T: DeserializeOwned>(&self, expected_type: &str) -> Result<T> {
        // Payload for postgres_changes usually looks like:
        // { "type": "INSERT", "table": "users", "schema": "public", "record": { ... }, "old_record": null }

        let type_ = self
            .payload
            .get("type")
            .and_then(|v| v.as_str())
            .unwrap_or_default();

        // Note: Sometimes strict checking fails if event name differs from internal type.
        // We match loosely if type is empty or matches.
        if !type_.is_empty() && type_ != expected_type {
            return Err(SupaError::RealtimeError {
                message: format!("Expected type {}, got {}", expected_type, type_),
            });
        }

        let record_key = if expected_type == "DELETE" {
            "old_record"
        } else {
            "record"
        };
        let record = self.payload.get(record_key);

        match record {
            Some(val) if !val.is_null() => {
                serde_json::from_value(val.clone()).map_err(|e| SupaError::RealtimeError {
                    message: format!("Deserialization failed: {}", e),
                })
            }
            _ => {
                // Fallback: maybe it's the other key
                let fallback = self
                    .payload
                    .get("record")
                    .or_else(|| self.payload.get("old_record"));
                if let Some(val) = fallback {
                    if !val.is_null() {
                        return serde_json::from_value(val.clone()).map_err(|e| {
                            SupaError::RealtimeError {
                                message: format!("Deserialization failed (fallback): {}", e),
                            }
                        });
                    }
                }
                Err(SupaError::RealtimeError {
                    message: format!("No {} found in payload", record_key),
                })
            }
        }
    }
}