archy/core/archipelago/src/state.rs

use crate::data_model::{DataModel, WebSocketMessage};
use std::sync::Arc;
use tokio::sync::{broadcast, RwLock};
use tracing::debug;

/// Manages the application state and broadcasts updates to WebSocket clients
pub struct StateManager {
    data: Arc<RwLock<DataModel>>,
    revision: Arc<RwLock<u32>>,
    broadcast_tx: broadcast::Sender<WebSocketMessage>,
}

impl StateManager {
    pub fn new() -> Self {
        let (broadcast_tx, _) = broadcast::channel(100);
        Self {
            data: Arc::new(RwLock::new(DataModel::new())),
            revision: Arc::new(RwLock::new(0)),
            broadcast_tx,
        }
    }

    /// Get the current data model and revision
    pub async fn get_snapshot(&self) -> (DataModel, u32) {
        let data = self.data.read().await.clone();
        let rev = *self.revision.read().await;
        (data, rev)
    }

    /// Subscribe to state updates
    pub fn subscribe(&self) -> broadcast::Receiver<WebSocketMessage> {
        self.broadcast_tx.subscribe()
    }

    /// Update the data model and broadcast to all connected clients
    pub async fn update_data(&self, new_data: DataModel) {
        let mut data = self.data.write().await;
        let mut rev = self.revision.write().await;
        
        *data = new_data.clone();
        *rev += 1;
        
        debug!("Data model updated to revision {}", *rev);
        
        // Broadcast full data dump to all connected clients
        // In the future, we can optimize this by computing and sending JSON patches
        let message = WebSocketMessage {
            rev: *rev,
            data: Some(new_data),
            patch: None,
        };
        
        // Ignore errors if no receivers are connected
        let _ = self.broadcast_tx.send(message);
    }

    /// Get a WebSocket message with the current state
    pub async fn get_initial_message(&self) -> WebSocketMessage {
        let (data, rev) = self.get_snapshot().await;
        WebSocketMessage {
            rev,
            data: Some(data),
            patch: None,
        }
    }
}

impl Default for StateManager {
    fn default() -> Self {
        Self::new()
    }
}