satori-client 0.1.9

// src/lib.rs

use serde;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use tokio::sync::mpsc::{self, Sender};
use tokio::task;
use tokio::time::Duration;
use tokio_tungstenite::{connect_async_with_config, tungstenite::protocol::Message};
use url::Url;
use uuid::Uuid;

pub type Callback = Box<dyn Fn(serde_json::Value) + Send + Sync + 'static>;

#[derive(Clone)]
pub struct Satori {
    username: String,
    password: String,
    url: String,
    sender: Sender<Message>,
    subscriptions: Arc<Mutex<HashMap<String, Callback>>>,
}

impl Satori {
    pub async fn connect(username: String, password: String, url: String) -> anyhow::Result<Self> {
        // Configure WebSocket connection with no timeout
        let url_parsed = Url::parse(&url)?;

        // Create a custom connector with no timeout
        let (ws_stream, _) = connect_async_with_config(url_parsed, None, false).await?;

        // Configure the WebSocket to have no timeout by using keepalive settings
        // The WebSocket connection will maintain the connection without timing out

        let (mut write, read) = ws_stream.split();
        let (sender, mut receiver) = mpsc::channel(100);
        let subscriptions: Arc<Mutex<HashMap<String, Callback>>> =
            Arc::new(Mutex::new(HashMap::new()));

        let _subscriptions_clone = subscriptions.clone();

        task::spawn(async move {
            while let Some(msg) = receiver.recv().await {
                let _ = write.send(msg).await;
            }
        });

        let subscriptions_clone2 = subscriptions.clone();

        // Spawn a task to handle incoming messages and ping responses
        task::spawn(async move {
            let mut read = read;
            while let Some(Ok(msg)) = read.next().await {
                match msg {
                    Message::Text(txt) => {
                        if let Ok(json) = serde_json::from_str::<serde_json::Value>(&txt) {
                            if json.get("type")
                                == Some(&serde_json::Value::String("notification".to_string()))
                            {
                                if let Some(key) = json.get("key").and_then(|k| k.as_str()) {
                                    if let Some(cb) = subscriptions_clone2.lock().unwrap().get(key)
                                    {
                                        cb(json["data"].clone());
                                    }
                                }
                            }
                        }
                    }
                    Message::Ping(_data) => {
                        // Respond to ping with pong to keep connection alive
                        // Note: We can't send pong here since write is moved to another task
                        // The WebSocket library should handle ping/pong automatically
                    }
                    Message::Pong(_) => {
                        // Handle pong messages (optional)
                    }
                    _ => {
                        // Handle other message types if needed
                    }
                }
            }
        });

        // Spawn a task to send periodic pings to keep the connection alive
        let sender_clone = sender.clone();
        task::spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_secs(30));
            loop {
                interval.tick().await;
                if let Err(e) = sender_clone.send(Message::Ping(vec![])).await {
                    eprintln!("Failed to send ping: {}", e);
                    break;
                }
            }
        });

        Ok(Self {
            username,
            password,
            url,
            sender,
            subscriptions,
        })
    }

    pub async fn send(
        &self,
        mut payload: serde_json::Map<String, serde_json::Value>,
    ) -> anyhow::Result<serde_json::Value> {
        let id = Uuid::new_v4().to_string();
        payload.insert("id".into(), id.clone().into());
        payload.insert("username".into(), self.username.clone().into());
        payload.insert("password".into(), self.password.clone().into());

        let msg = Message::Text(serde_json::Value::Object(payload.clone()).to_string());
        self.sender.send(msg).await?;
        Ok(serde_json::json!({"status": "sent", "id": id}))
    }

    pub async fn command(
        &self,
        command: &str,
        args: serde_json::Value,
    ) -> anyhow::Result<serde_json::Value> {
        let mut payload = args.as_object().cloned().unwrap_or_default();
        payload.insert("command".into(), command.into());
        self.send(payload).await
    }

    pub async fn set_notify<F>(&self, key: &str, callback: F) -> anyhow::Result<()>
    where
        F: Fn(serde_json::Value) + Send + Sync + 'static,
    {
        self.subscriptions
            .lock()
            .unwrap()
            .insert(key.to_string(), Box::new(callback));
        let mut payload = serde_json::Map::new();
        payload.insert("command".into(), "NOTIFY".into());
        payload.insert("key".into(), key.into());
        self.send(payload).await?;
        Ok(())
    }

    // Convenience methods
    pub async fn set(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("SET", args).await
    }

    pub async fn query(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("QUERY", args).await
    }

    pub async fn ann(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("ANN", args).await
    }

    pub async fn set_middleware(
        &self,
        args: serde_json::Value,
    ) -> anyhow::Result<serde_json::Value> {
        self.command("SET_MIDDLEWARE", args).await
    }

    pub async fn get_access_frequency(
        &self,
        args: serde_json::Value,
    ) -> anyhow::Result<serde_json::Value> {
        self.command("GET_ACCESS_FREQUENCY", args).await
    }

    pub async fn get_operations(&self) -> anyhow::Result<serde_json::Value> {
        self.command("GET_OPERATIONS", serde_json::Value::Null)
            .await
    }

    pub async fn ask(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("ASK", args).await
    }

    pub async fn get(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("GET", args).await
    }
    pub async fn put(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("PUT", args).await
    }
    pub async fn delete(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("DELETE", args).await
    }

    pub async fn encrypt(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("ENCRYPT", args).await
    }
    pub async fn decrypt(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("DECRYPT", args).await
    }
    pub async fn push(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("PUSH", args).await
    }
    pub async fn pop(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("POP", args).await
    }

    pub async fn splice(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("SPLICE", args).await
    }
    pub async fn remove(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("REMOVE", args).await
    }
    pub async fn dfs(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("DFS", args).await
    }
    pub async fn set_vertex(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("SET_VERTEX", args).await
    }
    pub async fn get_vertex(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("GET_VERTEX", args).await
    }
    pub async fn delete_vertex(
        &self,
        args: serde_json::Value,
    ) -> anyhow::Result<serde_json::Value> {
        self.command("DELETE_VERTEX", args).await
    }

    // ========== Graph Operations ==========

    /// Performs breadth-first search on the graph.
    /// Returns all nodes reachable from a starting node.
    pub async fn graph_bfs(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("GRAPH_BFS", args).await
    }

    /// Performs depth-first search on the graph.
    /// Returns all nodes reachable from a starting node.
    pub async fn graph_dfs(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("GRAPH_DFS", args).await
    }

    /// Finds the shortest path between a start node and end node using Dijkstra's algorithm.
    pub async fn graph_shortest_path(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("GRAPH_SHORTEST_PATH", args).await
    }

    /// Identifies all connected components in the graph.
    pub async fn graph_connected_components(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("GRAPH_CONNECTED_COMPONENTS", args).await
    }

    /// Finds strongly connected components using Tarjan's algorithm.
    pub async fn graph_scc(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("GRAPH_SCC", args).await
    }

    /// Calculates degree centrality (number of connections) for each node in the graph.
    pub async fn graph_degree_centrality(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("GRAPH_DEGREE_CENTRALITY", args).await
    }

    /// Calculates closeness centrality for each node.
    pub async fn graph_closeness_centrality(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("GRAPH_CLOSENESS_CENTRALITY", args).await
    }

    /// Finds the centroid node (node with highest closeness centrality).
    pub async fn graph_centroid(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("GRAPH_CENTROID", args).await
    }

    // ========== AI Operations ==========

    /// Alias for ANN - performs approximate nearest neighbor search.
    pub async fn get_similar(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("GET_SIMILAR", args).await
    }

    // ========== Mindspace Operations ==========

    /// Creates a new mindspace (cognitive context) for AI-powered conversations.
    pub async fn set_mindspace(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("SET_MINDSPACE", args).await
    }

    /// Alias for SET_MINDSPACE.
    pub async fn create_mindspace(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("SET_MINDSPACE", args).await
    }

    /// Deletes a mindspace and all its associated context.
    pub async fn delete_mindspace(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("DELETE_MINDSPACE", args).await
    }

    /// Sends a message to a mindspace and receives an AI-generated response.
    pub async fn chat_mindspace(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("CHAT_MINDSPACE", args).await
    }

    /// Imports text corpus into a mindspace for semantic search and context retrieval.
    pub async fn lecture_mindspace(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("LECTURE_MINDSPACE", args).await
    }
}

use futures_util::SinkExt;
use futures_util::StreamExt;