storkit/server/src/worktree.rs

use crate::slog;
use crate::config::ProjectConfig;
use std::path::{Path, PathBuf};
use std::process::Command;

/// Write a `.mcp.json` file in the given directory pointing to the MCP server
/// at the given port.
pub fn write_mcp_json(dir: &Path, port: u16) -> Result<(), String> {
    let content = format!(
        "{{\n  \"mcpServers\": {{\n    \"story-kit\": {{\n      \"type\": \"http\",\n      \"url\": \"http://localhost:{port}/mcp\"\n    }}\n  }}\n}}\n"
    );
    std::fs::write(dir.join(".mcp.json"), content).map_err(|e| format!("Write .mcp.json: {e}"))
}

#[derive(Debug, Clone)]
pub struct WorktreeInfo {
    pub path: PathBuf,
    pub branch: String,
    pub base_branch: String,
}

#[derive(Debug, Clone)]
pub struct WorktreeListEntry {
    pub story_id: String,
    pub path: PathBuf,
}

/// Worktree path inside the project: `{project_root}/.story_kit/worktrees/{story_id}`.
pub fn worktree_path(project_root: &Path, story_id: &str) -> PathBuf {
    project_root
        .join(".story_kit")
        .join("worktrees")
        .join(story_id)
}

fn branch_name(story_id: &str) -> String {
    format!("feature/story-{story_id}")
}

/// Detect the current branch of the project root (the base branch worktrees fork from).
fn detect_base_branch(project_root: &Path) -> String {
    Command::new("git")
        .args(["rev-parse", "--abbrev-ref", "HEAD"])
        .current_dir(project_root)
        .output()
        .ok()
        .and_then(|o| {
            if o.status.success() {
                Some(String::from_utf8_lossy(&o.stdout).trim().to_string())
            } else {
                None
            }
        })
        .unwrap_or_else(|| "master".to_string())
}

/// Create a git worktree for the given story.
///
/// - Creates the worktree at `{project_root}/.story_kit/worktrees/{story_id}`
///   on branch `feature/story-{story_id}`.
/// - Writes `.mcp.json` in the worktree pointing to the MCP server at `port`.
/// - Runs setup commands from the config for each component.
/// - If the worktree/branch already exists, reuses rather than errors.
pub async fn create_worktree(
    project_root: &Path,
    story_id: &str,
    config: &ProjectConfig,
    port: u16,
) -> Result<WorktreeInfo, String> {
    let wt_path = worktree_path(project_root, story_id);
    let branch = branch_name(story_id);
    let base_branch = detect_base_branch(project_root);
    let root = project_root.to_path_buf();

    // Already exists — reuse (ensure sparse checkout is configured)
    if wt_path.exists() {
        let wt_clone = wt_path.clone();
        tokio::task::spawn_blocking(move || configure_sparse_checkout(&wt_clone))
            .await
            .map_err(|e| format!("spawn_blocking: {e}"))??;
        write_mcp_json(&wt_path, port)?;
        run_setup_commands(&wt_path, config).await?;
        return Ok(WorktreeInfo {
            path: wt_path,
            branch,
            base_branch,
        });
    }

    let wt = wt_path.clone();
    let br = branch.clone();

    tokio::task::spawn_blocking(move || create_worktree_sync(&root, &wt, &br))
        .await
        .map_err(|e| format!("spawn_blocking: {e}"))??;

    write_mcp_json(&wt_path, port)?;
    run_setup_commands(&wt_path, config).await?;

    Ok(WorktreeInfo {
        path: wt_path,
        branch,
        base_branch,
    })
}

fn create_worktree_sync(
    project_root: &Path,
    wt_path: &Path,
    branch: &str,
) -> Result<(), String> {
    // Ensure the parent directory exists
    if let Some(parent) = wt_path.parent() {
        std::fs::create_dir_all(parent)
            .map_err(|e| format!("Create worktree dir: {e}"))?;
    }

    // Prune stale worktree references (e.g. directories deleted externally)
    let _ = Command::new("git")
        .args(["worktree", "prune"])
        .current_dir(project_root)
        .output();

    // Try to create branch. If it already exists that's fine.
    let _ = Command::new("git")
        .args(["branch", branch])
        .current_dir(project_root)
        .output();

    // Create worktree
    let output = Command::new("git")
        .args([
            "worktree",
            "add",
            &wt_path.to_string_lossy(),
            branch,
        ])
        .current_dir(project_root)
        .output()
        .map_err(|e| format!("git worktree add: {e}"))?;

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        // If it says already checked out, that's fine
        if stderr.contains("already checked out") || stderr.contains("already exists") {
            return Ok(());
        }
        return Err(format!("git worktree add failed: {stderr}"));
    }

    // Enable sparse checkout to exclude pipeline files from the worktree.
    // This prevents .story_kit/work/ changes from ending up in feature branches,
    // which cause rename/delete merge conflicts when merging back to master.
    configure_sparse_checkout(wt_path)?;

    Ok(())
}

/// Placeholder for worktree isolation of `.story_kit/work/`.
///
/// Previous approaches (sparse checkout, skip-worktree) all leaked state
/// from worktrees back to the main checkout's config/index. For now this
/// is a no-op — merge conflicts from pipeline file moves are handled at
/// merge time by the mergemaster (squash merge ignores work/ diffs).
fn configure_sparse_checkout(_wt_path: &Path) -> Result<(), String> {
    Ok(())
}

/// Remove a git worktree and its branch.
pub async fn remove_worktree(
    project_root: &Path,
    info: &WorktreeInfo,
    config: &ProjectConfig,
) -> Result<(), String> {
    run_teardown_commands(&info.path, config).await?;

    let root = project_root.to_path_buf();
    let wt_path = info.path.clone();
    let branch = info.branch.clone();

    tokio::task::spawn_blocking(move || remove_worktree_sync(&root, &wt_path, &branch))
        .await
        .map_err(|e| format!("spawn_blocking: {e}"))?
}

/// Remove a git worktree by story ID, deriving the path and branch deterministically.
pub async fn remove_worktree_by_story_id(
    project_root: &Path,
    story_id: &str,
    config: &ProjectConfig,
) -> Result<(), String> {
    let path = worktree_path(project_root, story_id);
    if !path.exists() {
        return Err(format!("Worktree not found for story: {story_id}"));
    }
    let branch = branch_name(story_id);
    let base_branch = detect_base_branch(project_root);
    let info = WorktreeInfo {
        path,
        branch,
        base_branch,
    };
    remove_worktree(project_root, &info, config).await
}

/// List all worktrees under `{project_root}/.story_kit/worktrees/`.
pub fn list_worktrees(project_root: &Path) -> Result<Vec<WorktreeListEntry>, String> {
    let worktrees_dir = project_root.join(".story_kit").join("worktrees");
    if !worktrees_dir.exists() {
        return Ok(Vec::new());
    }
    let mut entries = Vec::new();
    for entry in
        std::fs::read_dir(&worktrees_dir).map_err(|e| format!("list worktrees: {e}"))?
    {
        let entry = entry.map_err(|e| format!("list worktrees entry: {e}"))?;
        let path = entry.path();
        if path.is_dir() {
            let story_id = path
                .file_name()
                .map(|n| n.to_string_lossy().to_string())
                .unwrap_or_default();
            entries.push(WorktreeListEntry { story_id, path });
        }
    }
    entries.sort_by(|a, b| a.story_id.cmp(&b.story_id));
    Ok(entries)
}

fn remove_worktree_sync(
    project_root: &Path,
    wt_path: &Path,
    branch: &str,
) -> Result<(), String> {
    // Remove worktree
    let output = Command::new("git")
        .args([
            "worktree",
            "remove",
            "--force",
            &wt_path.to_string_lossy(),
        ])
        .current_dir(project_root)
        .output()
        .map_err(|e| format!("git worktree remove: {e}"))?;

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        slog!("[worktree] remove warning: {stderr}");
    }

    // Delete branch (best effort)
    let _ = Command::new("git")
        .args(["branch", "-d", branch])
        .current_dir(project_root)
        .output();

    Ok(())
}

async fn run_setup_commands(wt_path: &Path, config: &ProjectConfig) -> Result<(), String> {
    for component in &config.component {
        let cmd_dir = wt_path.join(&component.path);
        for cmd in &component.setup {
            run_shell_command(cmd, &cmd_dir).await?;
        }
    }
    Ok(())
}

async fn run_teardown_commands(wt_path: &Path, config: &ProjectConfig) -> Result<(), String> {
    for component in &config.component {
        let cmd_dir = wt_path.join(&component.path);
        for cmd in &component.teardown {
            // Best effort — don't fail teardown
            if let Err(e) = run_shell_command(cmd, &cmd_dir).await {
                slog!("[worktree] teardown warning for {}: {e}", component.name);
            }
        }
    }
    Ok(())
}

async fn run_shell_command(cmd: &str, cwd: &Path) -> Result<(), String> {
    let cmd = cmd.to_string();
    let cwd = cwd.to_path_buf();

    tokio::task::spawn_blocking(move || {
        slog!("[worktree] Running: {cmd} in {}", cwd.display());
        let output = Command::new("sh")
            .args(["-c", &cmd])
            .current_dir(&cwd)
            .output()
            .map_err(|e| format!("Run '{cmd}': {e}"))?;

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            return Err(format!("Command '{cmd}' failed: {stderr}"));
        }
        Ok(())
    })
    .await
    .map_err(|e| format!("spawn_blocking: {e}"))?
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::ComponentConfig;
    use std::fs;
    use tempfile::TempDir;

    /// Initialise a bare-minimum git repo so worktree operations work.
    fn init_git_repo(dir: &Path) {
        Command::new("git")
            .args(["init"])
            .current_dir(dir)
            .output()
            .expect("git init");
        Command::new("git")
            .args(["commit", "--allow-empty", "-m", "init"])
            .current_dir(dir)
            .output()
            .expect("git commit");
    }

    #[test]
    fn write_mcp_json_uses_given_port() {
        let tmp = TempDir::new().unwrap();
        write_mcp_json(tmp.path(), 4242).unwrap();
        let content = std::fs::read_to_string(tmp.path().join(".mcp.json")).unwrap();
        assert!(content.contains("http://localhost:4242/mcp"));
    }

    #[test]
    fn write_mcp_json_default_port() {
        let tmp = TempDir::new().unwrap();
        write_mcp_json(tmp.path(), 3001).unwrap();
        let content = std::fs::read_to_string(tmp.path().join(".mcp.json")).unwrap();
        assert!(content.contains("http://localhost:3001/mcp"));
    }

    #[test]
    fn worktree_path_is_inside_project() {
        let project_root = Path::new("/home/user/my-project");
        let path = worktree_path(project_root, "42_my_story");
        assert_eq!(
            path,
            Path::new("/home/user/my-project/.story_kit/worktrees/42_my_story")
        );
    }

    #[test]
    fn list_worktrees_empty_when_no_dir() {
        let tmp = TempDir::new().unwrap();
        let entries = list_worktrees(tmp.path()).unwrap();
        assert!(entries.is_empty());
    }

    #[test]
    fn list_worktrees_returns_subdirs() {
        let tmp = TempDir::new().unwrap();
        let worktrees_dir = tmp.path().join(".story_kit").join("worktrees");
        fs::create_dir_all(worktrees_dir.join("42_story_a")).unwrap();
        fs::create_dir_all(worktrees_dir.join("43_story_b")).unwrap();
        // A file (not dir) — should be ignored
        fs::write(worktrees_dir.join("readme.txt"), "").unwrap();

        let entries = list_worktrees(tmp.path()).unwrap();
        assert_eq!(entries.len(), 2);
        assert_eq!(entries[0].story_id, "42_story_a");
        assert_eq!(entries[1].story_id, "43_story_b");
    }

    #[test]
    fn create_worktree_after_stale_reference() {
        let tmp = TempDir::new().unwrap();
        let project_root = tmp.path().join("my-project");
        fs::create_dir_all(&project_root).unwrap();
        init_git_repo(&project_root);

        let wt_path = tmp.path().join("my-worktree");
        let branch = "feature/test-stale";

        // First creation should succeed
        create_worktree_sync(&project_root, &wt_path, branch).unwrap();
        assert!(wt_path.exists());

        // Simulate external deletion (e.g., rm -rf by another agent)
        fs::remove_dir_all(&wt_path).unwrap();
        assert!(!wt_path.exists());

        // Second creation should succeed despite stale git reference.
        // Without `git worktree prune`, this fails with "already checked out"
        // or "already exists".
        let result = create_worktree_sync(&project_root, &wt_path, branch);
        assert!(
            result.is_ok(),
            "Expected worktree creation to succeed after stale reference, got: {:?}",
            result.err()
        );
        assert!(wt_path.exists());
    }

    #[test]
    fn worktree_has_all_files_including_work() {
        let tmp = TempDir::new().unwrap();
        let project_root = tmp.path().join("my-project");
        fs::create_dir_all(&project_root).unwrap();
        init_git_repo(&project_root);

        // Create a tracked file under .story_kit/work/ on the initial branch
        let work_dir = project_root.join(".story_kit").join("work");
        fs::create_dir_all(&work_dir).unwrap();
        fs::write(work_dir.join("test_story.md"), "# Test").unwrap();
        Command::new("git")
            .args(["add", "."])
            .current_dir(&project_root)
            .output()
            .unwrap();
        Command::new("git")
            .args(["commit", "-m", "add work file"])
            .current_dir(&project_root)
            .output()
            .unwrap();

        let wt_path = tmp.path().join("my-worktree");
        let branch = "feature/test-sparse";
        create_worktree_sync(&project_root, &wt_path, branch).unwrap();

        // Worktree should have all files including .story_kit/work/
        assert!(wt_path.join(".story_kit").join("work").exists());
        assert!(wt_path.join(".git").exists());

        // Main checkout must NOT be affected by worktree creation.
        assert!(
            work_dir.exists(),
            ".story_kit/work/ must still exist in the main checkout"
        );
    }

    #[test]
    fn branch_name_format() {
        assert_eq!(branch_name("42_my_story"), "feature/story-42_my_story");
        assert_eq!(branch_name("1_test"), "feature/story-1_test");
    }

    #[test]
    fn detect_base_branch_returns_branch_in_git_repo() {
        let tmp = TempDir::new().unwrap();
        let project_root = tmp.path().join("my-project");
        fs::create_dir_all(&project_root).unwrap();
        init_git_repo(&project_root);

        let branch = detect_base_branch(&project_root);
        assert!(!branch.is_empty());
    }

    #[test]
    fn detect_base_branch_falls_back_to_master_for_non_git_dir() {
        let tmp = TempDir::new().unwrap();
        let branch = detect_base_branch(tmp.path());
        assert_eq!(branch, "master");
    }

    #[test]
    fn configure_sparse_checkout_is_noop() {
        let tmp = TempDir::new().unwrap();
        assert!(configure_sparse_checkout(tmp.path()).is_ok());
    }

    #[tokio::test]
    async fn run_shell_command_succeeds_for_echo() {
        let tmp = TempDir::new().unwrap();
        let result = run_shell_command("echo hello", tmp.path()).await;
        assert!(result.is_ok(), "Expected success: {:?}", result.err());
    }

    #[tokio::test]
    async fn run_shell_command_fails_for_nonzero_exit() {
        let tmp = TempDir::new().unwrap();
        let result = run_shell_command("exit 1", tmp.path()).await;
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("failed"));
    }

    #[tokio::test]
    async fn run_setup_commands_no_components_succeeds() {
        let tmp = TempDir::new().unwrap();
        let config = ProjectConfig {
            component: vec![],
            agent: vec![],
        };
        assert!(run_setup_commands(tmp.path(), &config).await.is_ok());
    }

    #[tokio::test]
    async fn run_setup_commands_runs_each_command_successfully() {
        let tmp = TempDir::new().unwrap();
        let config = ProjectConfig {
            component: vec![ComponentConfig {
                name: "test".to_string(),
                path: ".".to_string(),
                setup: vec!["echo setup_ok".to_string()],
                teardown: vec![],
            }],
            agent: vec![],
        };
        assert!(run_setup_commands(tmp.path(), &config).await.is_ok());
    }

    #[tokio::test]
    async fn run_setup_commands_propagates_failure() {
        let tmp = TempDir::new().unwrap();
        let config = ProjectConfig {
            component: vec![ComponentConfig {
                name: "test".to_string(),
                path: ".".to_string(),
                setup: vec!["exit 1".to_string()],
                teardown: vec![],
            }],
            agent: vec![],
        };
        let result = run_setup_commands(tmp.path(), &config).await;
        assert!(result.is_err(), "Expected failure from failing setup command");
        assert!(result.unwrap_err().contains("failed"));
    }

    #[tokio::test]
    async fn run_teardown_commands_ignores_failures() {
        let tmp = TempDir::new().unwrap();
        let config = ProjectConfig {
            component: vec![ComponentConfig {
                name: "test".to_string(),
                path: ".".to_string(),
                setup: vec![],
                teardown: vec!["exit 1".to_string()],
            }],
            agent: vec![],
        };
        // Teardown failures are best-effort — should not propagate
        assert!(run_teardown_commands(tmp.path(), &config).await.is_ok());
    }

    #[tokio::test]
    async fn create_worktree_fresh_creates_dir_and_mcp_json() {
        let tmp = TempDir::new().unwrap();
        let project_root = tmp.path().join("my-project");
        fs::create_dir_all(&project_root).unwrap();
        init_git_repo(&project_root);

        let config = ProjectConfig {
            component: vec![],
            agent: vec![],
        };
        let info = create_worktree(&project_root, "42_fresh_test", &config, 3001)
            .await
            .unwrap();

        assert!(info.path.exists());
        assert!(info.path.join(".mcp.json").exists());
        let mcp = fs::read_to_string(info.path.join(".mcp.json")).unwrap();
        assert!(mcp.contains("3001"));
        assert_eq!(info.branch, "feature/story-42_fresh_test");
    }

    #[tokio::test]
    async fn create_worktree_reuses_existing_path_and_updates_port() {
        let tmp = TempDir::new().unwrap();
        let project_root = tmp.path().join("my-project");
        fs::create_dir_all(&project_root).unwrap();
        init_git_repo(&project_root);

        let config = ProjectConfig {
            component: vec![],
            agent: vec![],
        };
        // First creation
        let _info1 = create_worktree(&project_root, "43_reuse_test", &config, 3001)
            .await
            .unwrap();
        // Second call — worktree already exists, reuse path, update port
        let info2 = create_worktree(&project_root, "43_reuse_test", &config, 3002)
            .await
            .unwrap();

        let mcp = fs::read_to_string(info2.path.join(".mcp.json")).unwrap();
        assert!(mcp.contains("3002"), "MCP json should be updated to new port");
    }

    #[test]
    fn remove_worktree_sync_cleans_up_directory() {
        let tmp = TempDir::new().unwrap();
        let project_root = tmp.path().join("my-project");
        fs::create_dir_all(&project_root).unwrap();
        init_git_repo(&project_root);

        let wt_path = project_root
            .join(".story_kit")
            .join("worktrees")
            .join("test_rm");
        create_worktree_sync(&project_root, &wt_path, "feature/test-rm").unwrap();
        assert!(wt_path.exists());

        remove_worktree_sync(&project_root, &wt_path, "feature/test-rm").unwrap();
        assert!(!wt_path.exists());
    }

    #[tokio::test]
    async fn remove_worktree_by_story_id_returns_err_when_not_found() {
        let tmp = TempDir::new().unwrap();
        let config = ProjectConfig {
            component: vec![],
            agent: vec![],
        };

        let result = remove_worktree_by_story_id(tmp.path(), "99_nonexistent", &config).await;
        assert!(result.is_err());
        assert!(
            result
                .unwrap_err()
                .contains("Worktree not found for story: 99_nonexistent")
        );
    }

    #[tokio::test]
    async fn remove_worktree_by_story_id_removes_existing_worktree() {
        let tmp = TempDir::new().unwrap();
        let project_root = tmp.path().join("my-project");
        fs::create_dir_all(&project_root).unwrap();
        init_git_repo(&project_root);

        let config = ProjectConfig {
            component: vec![],
            agent: vec![],
        };
        create_worktree(&project_root, "88_remove_by_id", &config, 3001)
            .await
            .unwrap();

        let result =
            remove_worktree_by_story_id(&project_root, "88_remove_by_id", &config).await;
        assert!(result.is_ok(), "Expected removal to succeed: {:?}", result.err());
    }

    #[tokio::test]
    async fn remove_worktree_async_removes_directory() {
        let tmp = TempDir::new().unwrap();
        let project_root = tmp.path().join("my-project");
        fs::create_dir_all(&project_root).unwrap();
        init_git_repo(&project_root);

        let config = ProjectConfig {
            component: vec![],
            agent: vec![],
        };
        let info = create_worktree(&project_root, "77_remove_async", &config, 3001)
            .await
            .unwrap();

        let path = info.path.clone();
        assert!(path.exists());
        remove_worktree(&project_root, &info, &config).await.unwrap();
        assert!(!path.exists());
    }
}