aimux_server/

layout_save.rs

use anyhow::{bail, Context, Result};
use serde::{Deserialize, Serialize};
use aimux_common::config::LayoutNode;
use aimux_protocol::{LayoutListEntry, LayoutSource};

use crate::session::{Layout, LayoutResult, SessionManager};
use aimux_protocol::types::{PaneId, WindowId};

// ---------------------------------------------------------------------------
// Types
// ---------------------------------------------------------------------------

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum SavedLayout {
    Preset {
        layout: String,
        panes: Vec<SavedPane>,
    },
    Custom(LayoutNode),
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SavedPane {
    #[serde(default)]
    pub command: Option<String>,
    #[serde(default)]
    pub mark: Option<String>,
}

// ---------------------------------------------------------------------------
// Validation
// ---------------------------------------------------------------------------

pub fn validate_layout_name(name: &str) -> Result<()> {
    if name.is_empty() || name.len() > 64 {
        bail!("layout name must be 1-64 characters, got {}", name.len());
    }
    if !name
        .chars()
        .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_')
    {
        bail!(
            "layout name must be alphanumeric, hyphens, or underscores: {}",
            name
        );
    }
    Ok(())
}

// ---------------------------------------------------------------------------
// Preset name mapping
// ---------------------------------------------------------------------------

fn preset_to_string(layout: &Layout) -> String {
    match layout {
        Layout::Single => "single",
        Layout::EvenHorizontal => "even-horizontal",
        Layout::EvenVertical => "even-vertical",
        Layout::MainVertical => "main-vertical",
        Layout::Tiled => "tiled",
        Layout::Custom(_) => unreachable!("custom handled separately"),
    }
    .to_string()
}

fn parse_preset_name(name: &str) -> Result<Layout> {
    match name {
        "single" => Ok(Layout::Single),
        "even-horizontal" => Ok(Layout::EvenHorizontal),
        "even-vertical" => Ok(Layout::EvenVertical),
        "main-vertical" => Ok(Layout::MainVertical),
        "tiled" => Ok(Layout::Tiled),
        other => bail!("unknown preset layout: {}", other),
    }
}

const PRESET_NAMES: &[&str] = &[
    "single",
    "even-horizontal",
    "even-vertical",
    "main-vertical",
    "tiled",
];

// ---------------------------------------------------------------------------
// Save
// ---------------------------------------------------------------------------

fn build_saved_custom(
    node: &crate::layout_dsl::LayoutTree,
    panes: &[crate::session::Pane],
    marks: &std::collections::HashMap<String, String>,
    pane_idx: &mut usize,
) -> LayoutNode {
    match node {
        crate::layout_dsl::LayoutTree::Leaf => {
            let mut ln = LayoutNode {
                size: None,
                direction: None,
                children: None,
                command: None,
                mark: None,
            };
            if let Some(pane) = panes.get(*pane_idx) {
                ln.command = Some(pane.shell_command.clone());
                // Find first mark pointing to this pane
                let mark = marks
                    .iter()
                    .find(|(_, v)| v.as_str() == pane.id)
                    .map(|(k, _)| k.clone());
                ln.mark = mark;
            }
            *pane_idx += 1;
            ln
        }
        crate::layout_dsl::LayoutTree::Split {
            direction,
            children,
        } => {
            let dir = match direction {
                aimux_common::config::SplitDirection::Horizontal => {
                    aimux_common::config::SplitDirection::Horizontal
                }
                aimux_common::config::SplitDirection::Vertical => {
                    aimux_common::config::SplitDirection::Vertical
                }
            };
            let child_nodes: Vec<LayoutNode> = children
                .iter()
                .map(|(frac, child)| {
                    let mut ln = build_saved_custom(child, panes, marks, pane_idx);
                    ln.size = Some(format!("{}%", (frac * 100.0).round() as u32));
                    ln
                })
                .collect();
            LayoutNode {
                size: None,
                direction: Some(dir),
                children: Some(child_nodes),
                command: None,
                mark: None,
            }
        }
    }
}

fn ensure_layouts_dir() -> Result<std::path::PathBuf> {
    let dir = aimux_common::paths::layouts_dir();
    #[cfg(unix)]
    {
        use std::os::unix::fs::DirBuilderExt;
        std::fs::DirBuilder::new()
            .recursive(true)
            .mode(0o700)
            .create(&dir)
            .with_context(|| format!("failed to create layouts dir: {}", dir.display()))?;
    }
    #[cfg(not(unix))]
    {
        std::fs::create_dir_all(&dir)
            .with_context(|| format!("failed to create layouts dir: {}", dir.display()))?;
    }
    Ok(dir)
}

impl SessionManager {
    pub fn save_layout(
        &self,
        session_name: &str,
        window_id: WindowId,
        name: &str,
    ) -> Result<()> {
        validate_layout_name(name)?;

        let session = self
            .sessions
            .get(session_name)
            .ok_or_else(|| anyhow::anyhow!("session not found: {}", session_name))?;
        let window = session
            .find_window(window_id)
            .ok_or_else(|| anyhow::anyhow!("window not found: {}", window_id))?;

        let saved = match &window.layout {
            Layout::Custom(tree) => {
                let mut pane_idx = 0;
                SavedLayout::Custom(build_saved_custom(
                    tree,
                    &window.panes,
                    &self.marks,
                    &mut pane_idx,
                ))
            }
            preset => {
                let layout_name = preset_to_string(preset);
                let panes: Vec<SavedPane> = window
                    .panes
                    .iter()
                    .map(|p| SavedPane {
                        command: Some(p.shell_command.clone()),
                        mark: self
                            .marks
                            .iter()
                            .find(|(_, v)| v.as_str() == p.id)
                            .map(|(k, _)| k.clone()),
                    })
                    .collect();
                SavedLayout::Preset {
                    layout: layout_name,
                    panes,
                }
            }
        };

        let dir = ensure_layouts_dir()?;
        let path = dir.join(format!("{}.json", name));
        let json = serde_json::to_string_pretty(&saved)
            .context("failed to serialize layout")?;
        std::fs::write(&path, json)
            .with_context(|| format!("failed to write layout file: {}", path.display()))?;
        Ok(())
    }

    pub fn load_layout(
        &mut self,
        session_name: &str,
        window_id: WindowId,
        name: &str,
    ) -> LayoutResult {
        validate_layout_name(name)?;

        let saved = resolve_layout(name, self.config())?;

        match saved {
            SavedLayout::Preset { layout, panes } => {
                self.load_preset_layout(session_name, window_id, &layout, &panes)
            }
            SavedLayout::Custom(node) => {
                self.load_custom_layout(session_name, window_id, &node)
            }
        }
    }

    fn load_preset_layout(
        &mut self,
        session_name: &str,
        window_id: WindowId,
        layout_name: &str,
        saved_panes: &[SavedPane],
    ) -> LayoutResult {
        let layout = parse_preset_name(layout_name)?;
        let target_count = if saved_panes.is_empty() {
            1
        } else {
            saved_panes.len()
        };

        let session = self
            .sessions
            .get(session_name)
            .ok_or_else(|| anyhow::anyhow!("session not found: {}", session_name))?;
        let default_shell = session.options.default_shell.clone();
        let scrollback_limit = session.options.scrollback_limit;
        let window = session
            .find_window(window_id)
            .ok_or_else(|| anyhow::anyhow!("window not found: {}", window_id))?;

        // Kill existing panes — collect backends instead of closing
        let existing_pane_ids: Vec<PaneId> =
            window.panes.iter().map(|p| p.id.clone()).collect();
        for pane_id in &existing_pane_ids {
            self.pane_index.remove(pane_id);
            self.marks.retain(|_, v| v.as_str() != pane_id.as_str());
        }
        let killed_backends;
        {
            let session = self.sessions.get_mut(session_name).unwrap();
            let window = session.find_window_mut(window_id).unwrap();
            killed_backends = window.panes.drain(..).map(|p| p.pty).collect::<Vec<_>>();
        }

        // Create new panes
        let mut new_pane_ids = Vec::with_capacity(target_count);
        for i in 0..target_count {
            let shell = saved_panes
                .get(i)
                .and_then(|sp| sp.command.as_deref())
                .unwrap_or(&default_shell);
            let pane = self.create_pane(shell, 80, 24, scrollback_limit)?;
            let pane_id = pane.id.clone();
            self.pane_index
                .insert(pane_id.clone(), (session_name.to_string(), window_id));
            let session = self.sessions.get_mut(session_name).unwrap();
            let window = session.find_window_mut(window_id).unwrap();
            window.panes.push(pane);
            new_pane_ids.push(pane_id);
        }

        // Set layout
        {
            let session = self.sessions.get_mut(session_name).unwrap();
            let window = session.find_window_mut(window_id).unwrap();
            window.layout = layout;
            window.active_pane = 0;
        }

        // Apply marks
        for (i, sp) in saved_panes.iter().enumerate() {
            if let Some(ref mark_name) = sp.mark {
                if i < new_pane_ids.len() {
                    let _ = self.set_mark(mark_name, &new_pane_ids[i]);
                }
            }
        }

        Ok((new_pane_ids, killed_backends))
    }

    fn load_custom_layout(
        &mut self,
        session_name: &str,
        window_id: WindowId,
        node: &LayoutNode,
    ) -> LayoutResult {
        // Temporarily insert layout node into config, delegate to apply_layout,
        // then remove it. This reuses the existing tree-based layout application.
        let temp_name = "__aimux_load_layout_temp__".to_string();

        // Insert into config layouts
        let layouts = self.config_mut().layouts.get_or_insert_with(Default::default);
        layouts.insert(temp_name.clone(), node.clone());

        // Kill existing panes first — collect backends instead of closing
        let killed_backends;
        {
            let session = self
                .sessions
                .get(session_name)
                .ok_or_else(|| anyhow::anyhow!("session not found: {}", session_name))?;
            let window = session
                .find_window(window_id)
                .ok_or_else(|| anyhow::anyhow!("window not found: {}", window_id))?;
            let existing_pane_ids: Vec<PaneId> =
                window.panes.iter().map(|p| p.id.clone()).collect();
            for pane_id in &existing_pane_ids {
                self.pane_index.remove(pane_id);
                self.marks.retain(|_, v| v.as_str() != pane_id.as_str());
            }
            let session = self.sessions.get_mut(session_name).unwrap();
            let window = session.find_window_mut(window_id).unwrap();
            killed_backends = window.panes.drain(..).map(|p| p.pty).collect::<Vec<_>>();
            window.layout = Layout::Single;
        }

        // Create one placeholder pane so apply_layout has something to work with
        {
            let session = self.sessions.get(session_name).unwrap();
            let default_shell = session.options.default_shell.clone();
            let scrollback_limit = session.options.scrollback_limit;
            let pane = self.create_pane(&default_shell, 80, 24, scrollback_limit)?;
            let pane_id = pane.id.clone();
            self.pane_index
                .insert(pane_id, (session_name.to_string(), window_id));
            let session = self.sessions.get_mut(session_name).unwrap();
            let window = session.find_window_mut(window_id).unwrap();
            window.panes.push(pane);
        }

        let (pane_ids, mut apply_backends) = self.apply_layout(session_name, window_id, &temp_name)?;

        // Clean up temp config entry
        if let Some(ref mut layouts) = self.config_mut().layouts {
            layouts.remove(&temp_name);
        }

        // Combine backends from the killed panes and any from apply_layout
        let mut all_backends = killed_backends;
        all_backends.append(&mut apply_backends);
        Ok((pane_ids, all_backends))
    }

    pub fn list_layouts(&self) -> Vec<LayoutListEntry> {
        let mut entries = Vec::new();

        // 1. Config-defined layouts
        if let Some(ref layouts) = self.config().layouts {
            for name in layouts.keys() {
                entries.push(LayoutListEntry {
                    name: name.clone(),
                    source: LayoutSource::Config,
                });
            }
        }

        // 2. Saved layouts from disk
        let dir = aimux_common::paths::layouts_dir();
        if let Ok(read_dir) = std::fs::read_dir(&dir) {
            for entry in read_dir.flatten() {
                let path = entry.path();
                if path.extension().and_then(|e| e.to_str()) == Some("json") {
                    if let Some(stem) = path.file_stem().and_then(|s| s.to_str()) {
                        // Skip if already listed as config
                        if !entries.iter().any(|e| e.name == stem) {
                            entries.push(LayoutListEntry {
                                name: stem.to_string(),
                                source: LayoutSource::Saved,
                            });
                        }
                    }
                }
            }
        }

        // 3. Built-in presets
        for &preset in PRESET_NAMES {
            if !entries.iter().any(|e| e.name == preset) {
                entries.push(LayoutListEntry {
                    name: preset.to_string(),
                    source: LayoutSource::Preset,
                });
            }
        }

        entries
    }

    /// Resolve a layout name: config → saved files → presets.
    /// Used by --layout flag on session/window creation.
    pub fn resolve_layout_name(&self, name: &str) -> Result<()> {
        // Just validate that the name can be resolved; the actual application
        // is handled by apply_layout or load_layout.
        let _ = resolve_layout(name, self.config())?;
        Ok(())
    }
}

/// Resolution order: config → saved files → presets.
fn resolve_layout(name: &str, config: &aimux_common::config::AimuxConfig) -> Result<SavedLayout> {
    // 1. Config file's layouts section
    if let Some(ref layouts) = config.layouts {
        if let Some(node) = layouts.get(name) {
            return Ok(SavedLayout::Custom(node.clone()));
        }
    }

    // 2. Saved layouts directory
    let dir = aimux_common::paths::layouts_dir();
    let path = dir.join(format!("{}.json", name));
    if path.exists() {
        let json = std::fs::read_to_string(&path)
            .with_context(|| format!("failed to read layout file: {}", path.display()))?;
        let saved: SavedLayout = serde_json::from_str(&json)
            .with_context(|| format!("failed to parse layout file: {}", path.display()))?;
        return Ok(saved);
    }

    // 3. Built-in presets
    if parse_preset_name(name).is_ok() {
        return Ok(SavedLayout::Preset {
            layout: name.to_string(),
            panes: vec![],
        });
    }

    bail!(
        "layout not found: {} (checked config, saved layouts, and presets)",
        name
    )
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    #[test]
    fn validate_layout_name_valid() {
        validate_layout_name("dev").unwrap();
        validate_layout_name("my-layout").unwrap();
        validate_layout_name("test_1").unwrap();
        validate_layout_name("A").unwrap();
        validate_layout_name(&"a".repeat(64)).unwrap();
    }

    #[test]
    fn validate_layout_name_invalid() {
        // Empty
        assert!(validate_layout_name("").is_err());
        // Too long
        assert!(validate_layout_name(&"a".repeat(65)).is_err());
        // Invalid chars
        assert!(validate_layout_name("hello world").is_err());
        assert!(validate_layout_name("foo/bar").is_err());
        assert!(validate_layout_name("foo\\bar").is_err());
        assert!(validate_layout_name("a..b").is_err());
        assert!(validate_layout_name("foo@bar").is_err());
    }

    #[test]
    fn saved_layout_preset_serde_roundtrip() {
        let saved = SavedLayout::Preset {
            layout: "even-horizontal".to_string(),
            panes: vec![
                SavedPane {
                    command: Some("nvim".to_string()),
                    mark: Some("editor".to_string()),
                },
                SavedPane {
                    command: Some("cargo watch".to_string()),
                    mark: None,
                },
            ],
        };
        let json = serde_json::to_string_pretty(&saved).unwrap();
        let parsed: SavedLayout = serde_json::from_str(&json).unwrap();
        match parsed {
            SavedLayout::Preset { layout, panes } => {
                assert_eq!(layout, "even-horizontal");
                assert_eq!(panes.len(), 2);
                assert_eq!(panes[0].command.as_deref(), Some("nvim"));
                assert_eq!(panes[0].mark.as_deref(), Some("editor"));
                assert_eq!(panes[1].command.as_deref(), Some("cargo watch"));
                assert!(panes[1].mark.is_none());
            }
            _ => panic!("expected Preset variant"),
        }
    }

    #[test]
    fn saved_layout_custom_serde_roundtrip() {
        let node = LayoutNode {
            size: None,
            direction: Some(aimux_common::config::SplitDirection::Horizontal),
            children: Some(vec![
                LayoutNode {
                    size: Some("60%".to_string()),
                    direction: None,
                    children: None,
                    command: Some("nvim".to_string()),
                    mark: Some("editor".to_string()),
                },
                LayoutNode {
                    size: Some("40%".to_string()),
                    direction: None,
                    children: None,
                    command: None,
                    mark: None,
                },
            ]),
            command: None,
            mark: None,
        };
        let saved = SavedLayout::Custom(node);
        let json = serde_json::to_string_pretty(&saved).unwrap();
        let parsed: SavedLayout = serde_json::from_str(&json).unwrap();
        match parsed {
            SavedLayout::Custom(n) => {
                assert!(n.direction.is_some());
                let children = n.children.unwrap();
                assert_eq!(children.len(), 2);
                assert_eq!(children[0].command.as_deref(), Some("nvim"));
            }
            _ => panic!("expected Custom variant"),
        }
    }

    #[test]
    fn preset_names_roundtrip() {
        for &name in PRESET_NAMES {
            let layout = parse_preset_name(name).unwrap();
            let back = preset_to_string(&layout);
            assert_eq!(back, name);
        }
    }

    #[tokio::test]
    async fn save_and_load_preset_layout() {
        use crate::session::testing::new_manager;

        let mut mgr = new_manager();
        let (_, _p0) = mgr.new_session("s1", None).unwrap();
        let _p1 = mgr.split_pane("%0", true, None).unwrap();

        // Save using a temp dir
        let tmp = std::env::temp_dir().join("aimux_test_layouts");
        let _ = std::fs::remove_dir_all(&tmp);
        std::fs::create_dir_all(&tmp).unwrap();

        // Manually serialize and write
        let session = mgr.sessions.get("s1").unwrap();
        let window = &session.windows[0];
        let saved = SavedLayout::Preset {
            layout: "even-horizontal".to_string(),
            panes: window
                .panes
                .iter()
                .map(|p| SavedPane {
                    command: Some(p.shell_command.clone()),
                    mark: None,
                })
                .collect(),
        };
        let json = serde_json::to_string_pretty(&saved).unwrap();
        std::fs::write(tmp.join("test-layout.json"), &json).unwrap();

        // Verify it was serialized correctly
        let read_back: SavedLayout =
            serde_json::from_str(&std::fs::read_to_string(tmp.join("test-layout.json")).unwrap())
                .unwrap();
        match read_back {
            SavedLayout::Preset { layout, panes } => {
                assert_eq!(layout, "even-horizontal");
                assert_eq!(panes.len(), 2);
            }
            _ => panic!("expected Preset"),
        }

        let _ = std::fs::remove_dir_all(&tmp);
    }

    #[test]
    fn list_layouts_includes_presets() {
        use crate::session::testing::new_manager;

        let mgr = new_manager();
        let entries = mgr.list_layouts();
        // Should include all preset names
        for &name in PRESET_NAMES {
            assert!(
                entries.iter().any(|e| e.name == name),
                "missing preset: {}",
                name
            );
        }
    }
}