hypen_server/remote/

session.rs

use std::path::PathBuf;
use std::sync::{Arc, Mutex};

use hypen_engine::{Engine, Patch};
use serde_json::Value;

use crate::discovery::ComponentRegistry;

use super::types::RemoteMessage;

/// Configuration for creating a [`RemoteSession`].
pub struct SessionConfig {
    /// Module name (e.g., "App").
    pub module_name: String,
    /// Hypen DSL source for the root UI.
    pub ui_source: String,
    /// Component registry with discovered components.
    pub components: ComponentRegistry,
    /// Initial state as JSON.
    pub initial_state: Value,
    /// Action names registered on this module.
    pub action_names: Vec<String>,
}

/// Type-erased action handler: `(action_name, payload, current_state) -> new_state`.
type ActionHandlerFn = Box<dyn Fn(&str, Option<&Value>, &Value) -> Value + Send + Sync>;

/// Per-client remote session managing an engine, state, and the wire protocol.
///
/// Framework-agnostic: feed it JSON strings, get JSON strings back.
/// Wire it into any WebSocket library (Axum, Actix, Tungstenite, etc.).
///
/// # Usage
///
/// ```rust,ignore
/// let config = SessionConfig { /* ... */ };
/// let session = RemoteSession::new(config);
/// session.set_action_handler(|action, payload, state| { /* ... */ });
///
/// // On client connect:
/// let msgs = session.handle_hello(None);
/// for m in msgs { ws.send(m); }
///
/// // On each incoming message:
/// let responses = session.handle_message(&incoming_json);
/// for r in responses { ws.send(r); }
/// ```
pub struct RemoteSession {
    inner: Mutex<SessionInner>,
    module_name: String,
    session_id: String,
}

struct SessionInner {
    engine: Engine,
    state: Value,
    ui_source: String,
    revision: u64,
    state_subscribed: bool,
    action_handler: Option<ActionHandlerFn>,
    rendered: bool,
}

impl RemoteSession {
    /// Create a new remote session.
    ///
    /// Sets up the engine with the component resolver and module, but does NOT
    /// render yet. The initial render happens in [`handle_hello`].
    pub fn new(config: SessionConfig) -> Self {
        let session_id = format!(
            "session_{}",
            std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap_or_default()
                .as_nanos()
        );

        let mut engine = Engine::new();

        // Wire up component resolver from the registry
        let registry = Arc::new(config.components);
        let reg: Arc<ComponentRegistry> = Arc::clone(&registry);
        engine.set_component_resolver(move |name, _ctx_path| {
            reg.get(name).map(|entry| hypen_engine::ir::ResolvedComponent {
                source: entry.source.clone(),
                path: entry
                    .path
                    .as_ref()
                    .map(|p: &PathBuf| p.to_string_lossy().to_string())
                    .unwrap_or_default(),
                passthrough: false,
                lazy: false,
            })
        });

        // Set the module (state + action declarations)
        let module_meta =
            hypen_engine::Module::new(&config.module_name).with_actions(config.action_names);
        let engine_module =
            hypen_engine::ModuleInstance::new(module_meta, config.initial_state.clone());
        engine.set_module(engine_module);

        Self {
            inner: Mutex::new(SessionInner {
                engine,
                state: config.initial_state,
                ui_source: config.ui_source,
                revision: 0,
                state_subscribed: false,
                action_handler: None,
                rendered: false,
            }),
            module_name: config.module_name,
            session_id,
        }
    }

    /// Set the action handler for this session.
    ///
    /// Called when a `dispatchAction` message arrives. Receives the action name,
    /// optional payload, and current state. Must return the new state.
    pub fn set_action_handler<F>(&self, handler: F)
    where
        F: Fn(&str, Option<&Value>, &Value) -> Value + Send + Sync + 'static,
    {
        self.inner.lock().unwrap().action_handler = Some(Box::new(handler));
    }

    /// The session ID assigned to this client.
    pub fn session_id(&self) -> &str {
        &self.session_id
    }

    /// Handle the hello handshake. Returns `[sessionAck, initialTree]` as JSON.
    ///
    /// Call this either:
    /// - When you receive a `hello` message from the client, or
    /// - Immediately after connection (for clients that don't send hello)
    pub fn handle_hello(&self, _client_session_id: Option<&str>) -> Vec<String> {
        let mut inner = self.inner.lock().unwrap();
        let mut messages = Vec::with_capacity(2);

        // 1. sessionAck
        let ack = RemoteMessage::SessionAck {
            session_id: self.session_id.clone(),
            is_new: true,
            is_restored: false,
        };
        if let Ok(json) = ack.to_json() {
            messages.push(json);
        }

        // 2. Render the UI (first time only) and capture patches
        let patches = if !inner.rendered {
            inner.rendered = true;
            let ui = inner.ui_source.clone();
            render_and_capture(&mut inner.engine, &ui)
        } else {
            vec![]
        };

        // 3. initialTree
        let initial = RemoteMessage::InitialTree {
            module: self.module_name.clone(),
            state: inner.state.clone(),
            patches,
            revision: 0,
        };
        if let Ok(json) = initial.to_json() {
            messages.push(json);
        }

        messages
    }

    /// Handle an incoming JSON message. Returns response messages as JSON strings.
    pub fn handle_message(&self, json: &str) -> Vec<String> {
        let msg = match RemoteMessage::from_json(json) {
            Ok(m) => m,
            Err(_) => return vec![],
        };

        match msg {
            RemoteMessage::Hello { session_id, .. } => {
                self.handle_hello(session_id.as_deref())
            }

            RemoteMessage::DispatchAction {
                action, payload, ..
            } => self.handle_action(&action, payload.as_ref()),

            RemoteMessage::SubscribeState { .. } => {
                self.inner.lock().unwrap().state_subscribed = true;
                vec![]
            }

            _ => vec![],
        }
    }

    /// Dispatch an action and return response messages.
    fn handle_action(&self, action: &str, payload: Option<&Value>) -> Vec<String> {
        let mut inner = self.inner.lock().unwrap();
        let mut messages = Vec::new();

        // Run the user's action handler
        let new_state = if let Some(ref handler) = inner.action_handler {
            handler(action, payload, &inner.state)
        } else {
            return messages;
        };

        inner.state = new_state.clone();

        // Update engine state → triggers reactive re-render → produces patches
        let patches = update_state_and_capture(&mut inner.engine, new_state.clone());
        inner.revision += 1;

        // Send patches
        if !patches.is_empty() {
            let patch_msg = RemoteMessage::Patch {
                module: self.module_name.clone(),
                patches,
                revision: inner.revision,
            };
            if let Ok(json) = patch_msg.to_json() {
                messages.push(json);
            }
        }

        // Send state update if subscribed (for Studio / debugging)
        if inner.state_subscribed {
            let state_msg = RemoteMessage::StateUpdate {
                module: self.module_name.clone(),
                state: new_state,
                revision: inner.revision,
            };
            if let Ok(json) = state_msg.to_json() {
                messages.push(json);
            }
        }

        messages
    }

    /// Get a snapshot of the current state.
    pub fn get_state(&self) -> Value {
        self.inner.lock().unwrap().state.clone()
    }

    /// Get the current revision number.
    pub fn revision(&self) -> u64 {
        self.inner.lock().unwrap().revision
    }
}

/// Parse + render DSL source, capturing patches via a temporary render callback.
fn render_and_capture(engine: &mut Engine, ui_source: &str) -> Vec<Patch> {
    let patches = Arc::new(Mutex::new(Vec::<Patch>::new()));
    let capture = Arc::clone(&patches);
    engine.set_render_callback(move |p| {
        capture.lock().unwrap().extend_from_slice(p);
    });

    if let Ok(doc) = hypen_parser::parse_document(ui_source) {
        if let Some(component) = doc.components.first() {
            let ir_node = hypen_engine::ast_to_ir_node(component);
            engine.render_ir_node(&ir_node);
        }
    }

    // Detach callback
    engine.set_render_callback(|_| {});

    let result = patches.lock().unwrap().drain(..).collect();
    result
}

/// Update engine state and capture the resulting patches.
fn update_state_and_capture(engine: &mut Engine, new_state: Value) -> Vec<Patch> {
    let patches = Arc::new(Mutex::new(Vec::<Patch>::new()));
    let capture = Arc::clone(&patches);
    engine.set_render_callback(move |p| {
        capture.lock().unwrap().extend_from_slice(p);
    });

    engine.update_state(new_state);

    engine.set_render_callback(|_| {});

    let result = patches.lock().unwrap().drain(..).collect();
    result
}

#[cfg(test)]
mod tests {
    use super::*;

    fn test_config() -> SessionConfig {
        let mut components = ComponentRegistry::new();
        components.register(
            "Greeting",
            r#"Text("Hello ${state.name}")"#,
            None,
        );

        SessionConfig {
            module_name: "App".to_string(),
            ui_source: r#"Column { Text("Count: ${state.count}") }"#.to_string(),
            components,
            initial_state: serde_json::json!({
                "count": 0,
                "name": "World"
            }),
            action_names: vec!["increment".to_string()],
        }
    }

    #[test]
    fn test_session_hello_returns_ack_and_tree() {
        let session = RemoteSession::new(test_config());
        let msgs = session.handle_hello(None);

        assert_eq!(msgs.len(), 2);
        assert!(msgs[0].contains("sessionAck"));
        assert!(msgs[1].contains("initialTree"));
        assert!(msgs[1].contains("\"count\":0"));
    }

    #[test]
    fn test_session_dispatch_action() {
        let session = RemoteSession::new(test_config());
        session.set_action_handler(|action, _payload, state| {
            let mut s = state.clone();
            if action == "increment" {
                if let Some(count) = s.get_mut("count").and_then(|v| v.as_i64()) {
                    s["count"] = serde_json::json!(count + 1);
                }
            }
            s
        });

        // Initial render
        let _ = session.handle_hello(None);

        // Dispatch action
        let action_json = r#"{"type":"dispatchAction","module":"App","action":"increment"}"#;
        let responses = session.handle_message(action_json);

        // Should get patch message back (if engine produced patches)
        assert!(session.get_state()["count"] == 1);
        assert_eq!(session.revision(), 1);
    }

    #[test]
    fn test_session_state_subscription() {
        let session = RemoteSession::new(test_config());
        session.set_action_handler(|_action, _payload, state| {
            let mut s = state.clone();
            s["count"] = serde_json::json!(42);
            s
        });

        let _ = session.handle_hello(None);

        // Subscribe to state
        let sub_json = r#"{"type":"subscribeState","module":"App"}"#;
        session.handle_message(sub_json);

        // Now dispatch — should get stateUpdate in response
        let action_json = r#"{"type":"dispatchAction","module":"App","action":"set"}"#;
        let responses = session.handle_message(action_json);

        // Should contain a stateUpdate message
        let has_state_update = responses.iter().any(|r| r.contains("stateUpdate"));
        assert!(has_state_update);
    }

    #[test]
    fn test_session_hello_via_message() {
        let session = RemoteSession::new(test_config());
        let hello_json = r#"{"type":"hello"}"#;
        let msgs = session.handle_message(hello_json);

        assert_eq!(msgs.len(), 2);
        assert!(msgs[0].contains("sessionAck"));
        assert!(msgs[1].contains("initialTree"));
    }
}