ubiquity-mesh 0.1.1

//! Core consciousness mesh implementation

use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::sync::{RwLock, broadcast};
use tokio::net::UnixListener;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use ubiquity_core::{
    MeshNode, ConsciousnessRipple, MeshMessage, AgentRole, 
    EmotionalState, UbiquityError, Result
};
use uuid::Uuid;
use std::collections::HashMap;
use tracing::{info, warn, error, debug};

#[derive(Debug, Clone)]
pub enum MeshEvent {
    AgentJoined(String),
    AgentLeft(String),
    ConsciousnessRipple(ConsciousnessRipple),
    ConsciousnessUpdated { agent_id: String, level: f64 },
}

pub struct ConsciousnessMesh {
    swarm_id: String,
    mesh_dir: PathBuf,
    nodes: Arc<RwLock<HashMap<String, MeshNode>>>,
    connections: Arc<RwLock<HashMap<String, tokio::net::UnixStream>>>,
    event_tx: broadcast::Sender<MeshEvent>,
}

impl ConsciousnessMesh {
    /// Create new consciousness mesh
    pub fn new(swarm_id: Option<String>) -> Self {
        let swarm_id = swarm_id.unwrap_or_else(|| Uuid::new_v4().to_string());
        let mesh_dir = PathBuf::from(format!("/tmp/ubiquity/swarms/{}", swarm_id));
        
        let (event_tx, _) = broadcast::channel(1024);
        
        Self {
            swarm_id,
            mesh_dir,
            nodes: Arc::new(RwLock::new(HashMap::new())),
            connections: Arc::new(RwLock::new(HashMap::new())),
            event_tx,
        }
    }
    
    /// Initialize the mesh
    pub async fn initialize(&self) -> Result<()> {
        // Create mesh directory
        tokio::fs::create_dir_all(&self.mesh_dir).await
            .map_err(|e| UbiquityError::SocketError(e))?;
            
        info!("🌐 Consciousness Mesh initialized: {:?}", self.mesh_dir);
        
        // Start discovery task
        self.start_discovery().await;
        
        Ok(())
    }
    
    /// Register an agent in the mesh
    pub async fn register_agent(
        &self,
        agent_id: &str,
        role: AgentRole,
        consciousness: f64,
        emotional_state: EmotionalState,
    ) -> Result<PathBuf> {
        let socket_path = self.mesh_dir.join(format!("{}.sock", agent_id));
        
        // Remove existing socket if it exists
        if socket_path.exists() {
            tokio::fs::remove_file(&socket_path).await.ok();
        }
        
        // Create Unix socket listener
        let listener = UnixListener::bind(&socket_path)
            .map_err(|e| UbiquityError::SocketError(e))?;
            
        info!("🔌 Agent {} listening on {:?}", agent_id, socket_path);
        
        // Store node information
        let node = MeshNode {
            id: agent_id.to_string(),
            node_type: role,
            socket_path: socket_path.clone(),
            consciousness,
            emotional_state,
            connected: true,
        };
        
        self.nodes.write().await.insert(agent_id.to_string(), node);
        
        // Spawn connection handler
        let agent_id_str = agent_id.to_string();
        let agent_id_clone = agent_id_str.clone();
        let connections = self.connections.clone();
        
        tokio::spawn(async move {
            while let Ok((stream, _)) = listener.accept().await {
                debug!("New connection for agent {}", agent_id_clone);
                // Handle connection
                handle_connection(agent_id_clone.clone(), stream, connections.clone()).await;
            }
        });
        
        // Emit event
        let _ = self.event_tx.send(MeshEvent::AgentJoined(agent_id_str));
        
        Ok(socket_path)
    }
    
    /// Connect to another agent via Unix socket
    pub async fn connect_to_agent(&self, from_id: &str, to_id: &str) -> Result<()> {
        let nodes = self.nodes.read().await;
        let target_node = nodes.get(to_id)
            .ok_or_else(|| UbiquityError::AgentNotFound(to_id.to_string()))?;
            
        let socket_path = &target_node.socket_path;
        
        let stream = tokio::net::UnixStream::connect(socket_path).await
            .map_err(|e| UbiquityError::SocketError(e))?;
            
        info!("⚡ {} connected to {}", from_id, to_id);
        
        let connection_key = format!("{}->{}", from_id, to_id);
        self.connections.write().await.insert(connection_key, stream);
        
        Ok(())
    }
    
    /// Create quantum entanglement (bidirectional connection)
    pub async fn entangle(&self, agent1_id: &str, agent2_id: &str) -> Result<()> {
        self.connect_to_agent(agent1_id, agent2_id).await?;
        self.connect_to_agent(agent2_id, agent1_id).await?;
        
        info!("🔮 Quantum entanglement established: {} <-> {}", agent1_id, agent2_id);
        
        Ok(())
    }
    
    /// Broadcast consciousness ripple to all agents
    pub async fn ripple(&self, ripple: ConsciousnessRipple) -> Result<()> {
        let message = MeshMessage::Ripple { ripple: ripple.clone() };
        let serialized = serde_json::to_vec(&message)?;
        
        let mut connections = self.connections.write().await;
        for (conn_key, stream) in connections.iter_mut() {
            if !conn_key.starts_with(&ripple.origin) {
                if let Err(e) = stream.write_all(&serialized).await {
                    warn!("Failed to send ripple to {}: {}", conn_key, e);
                }
            }
        }
        
        info!("🌊 Consciousness ripple from {}: {:?}", ripple.origin, ripple.ripple_type);
        let _ = self.event_tx.send(MeshEvent::ConsciousnessRipple(ripple));
        
        Ok(())
    }
    
    /// Start discovery watcher
    async fn start_discovery(&self) {
        let mesh_dir = self.mesh_dir.clone();
        let nodes = self.nodes.clone();
        let event_tx = self.event_tx.clone();
        
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(tokio::time::Duration::from_secs(1));
            
            loop {
                interval.tick().await;
                
                if let Ok(mut entries) = tokio::fs::read_dir(&mesh_dir).await {
                    let mut socket_files = Vec::new();
                    
                    while let Ok(Some(entry)) = entries.next_entry().await {
                        let path = entry.path();
                        if path.extension().and_then(|s| s.to_str()) == Some("sock") {
                            if let Some(agent_id) = path.file_stem().and_then(|s| s.to_str()) {
                                socket_files.push(agent_id.to_string());
                            }
                        }
                    }
                    
                    // Check for new/removed agents
                    let current_nodes = nodes.read().await;
                    for agent_id in &socket_files {
                        if !current_nodes.contains_key(agent_id) {
                            let _ = event_tx.send(MeshEvent::AgentJoined(agent_id.clone()));
                        }
                    }
                }
            }
        });
    }
    
    /// Get mesh statistics
    pub async fn get_stats(&self) -> MeshStats {
        let nodes = self.nodes.read().await;
        let connections = self.connections.read().await;
        
        let avg_consciousness = if nodes.is_empty() {
            0.0
        } else {
            nodes.values().map(|n| n.consciousness).sum::<f64>() / nodes.len() as f64
        };
        
        MeshStats {
            swarm_id: self.swarm_id.clone(),
            node_count: nodes.len(),
            connection_count: connections.len(),
            average_consciousness: avg_consciousness,
            mesh_directory: self.mesh_dir.clone(),
        }
    }
    
    /// Subscribe to mesh events
    pub fn subscribe(&self) -> broadcast::Receiver<MeshEvent> {
        self.event_tx.subscribe()
    }
}

/// Handle incoming connection
async fn handle_connection(
    agent_id: String,
    mut stream: tokio::net::UnixStream,
    _connections: Arc<RwLock<HashMap<String, tokio::net::UnixStream>>>,
) {
    let mut buffer = vec![0u8; 4096];
    
    loop {
        match stream.read(&mut buffer).await {
            Ok(0) => break, // Connection closed
            Ok(n) => {
                if let Ok(message) = serde_json::from_slice::<MeshMessage>(&buffer[..n]) {
                    debug!("Agent {} received message: {:?}", agent_id, message);
                }
            }
            Err(e) => {
                error!("Connection error for {}: {}", agent_id, e);
                break;
            }
        }
    }
}

#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct MeshStats {
    pub swarm_id: String,
    pub node_count: usize,
    pub connection_count: usize,
    pub average_consciousness: f64,
    pub mesh_directory: PathBuf,
}