satori-client 0.1.6

// src/lib.rs

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::process::Command;
use std::sync::{Arc, Mutex};
use tokio::net::TcpStream;
use tokio::sync::mpsc::{self, Sender};
use tokio::task;
use tokio_tungstenite::{connect_async, 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> {
        let (ws_stream, _) = connect_async(Url::parse(&url)?).await?;
        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();
        task::spawn(async move {
            let mut read = read;
            while let Some(Ok(msg)) = read.next().await {
                if let Message::Text(txt) = msg {
                    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());
                                }
                            }
                        }
                    }
                }
            }
        });

        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 train(&self, args: serde_json::Value) -> anyhow::Result<serde_json::Value> {
        self.command("TRAIN", args).await
    }

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

}

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