worktrunk 0.35.1

use std::path::{Path, PathBuf};
use std::process;
use std::time::{SystemTime, UNIX_EPOCH};

use super::worktree::{BranchDeletionMode, RemoveResult, path_mismatch};
use anyhow::{Context, bail};
use color_print::cformat;
use worktrunk::config::UserConfig;
use worktrunk::git::{
    GitError, IntegrationReason, Repository, parse_porcelain_z, parse_untracked_files,
};
use worktrunk::path::format_path_for_display;
use worktrunk::styling::{eprintln, format_with_gutter, progress_message, warning_message};

/// Target for worktree removal.
#[derive(Debug)]
pub enum RemoveTarget<'a> {
    /// Remove worktree by branch name
    Branch(&'a str),
    /// Remove the current worktree (supports detached HEAD)
    Current,
    /// Remove worktree by path (supports detached HEAD)
    Path(&'a std::path::Path),
}

/// CLI-only helpers implemented on [`Repository`] via an extension trait so we can keep orphan
/// implementations inside the binary crate.
pub trait RepositoryCliExt {
    /// Warn about untracked files being auto-staged.
    fn warn_if_auto_staging_untracked(&self) -> anyhow::Result<()>;

    /// Prepare a worktree removal by branch name or current worktree.
    ///
    /// Returns a `RemoveResult` describing what will be removed. The actual
    /// removal is performed by the output handler.
    ///
    /// `current_path` overrides process-CWD discovery for determining which
    /// worktree is "current". Pass `None` for normal CLI usage (discovers from
    /// CWD). Pass `Some` when calling from a context where CWD may have changed
    /// (e.g., background threads in the picker).
    fn prepare_worktree_removal(
        &self,
        target: RemoveTarget,
        deletion_mode: BranchDeletionMode,
        force_worktree: bool,
        config: &UserConfig,
        current_path: Option<PathBuf>,
    ) -> anyhow::Result<RemoveResult>;

    /// Prepare the target worktree for push by auto-stashing non-overlapping changes when safe.
    fn prepare_target_worktree(
        &self,
        target_worktree: Option<&PathBuf>,
        target_branch: &str,
    ) -> anyhow::Result<Option<TargetWorktreeStash>>;

    /// Check if HEAD is a linear extension of the target branch.
    ///
    /// Returns true when:
    /// 1. The merge-base equals target's SHA (target hasn't advanced), AND
    /// 2. There are no merge commits between target and HEAD (history is linear)
    ///
    /// This detects branches that have merged the target into themselves — such
    /// branches need rebasing to linearize history even though merge-base equals target.
    fn is_rebased_onto(&self, target: &str) -> anyhow::Result<bool>;
}

impl RepositoryCliExt for Repository {
    fn warn_if_auto_staging_untracked(&self) -> anyhow::Result<()> {
        // Use -z for NUL-separated output to handle filenames with spaces/newlines
        let status = self
            .run_command(&["status", "--porcelain", "-z"])
            .context("Failed to get status")?;
        warn_about_untracked_files(&status)
    }

    fn prepare_worktree_removal(
        &self,
        target: RemoveTarget,
        deletion_mode: BranchDeletionMode,
        force_worktree: bool,
        config: &UserConfig,
        current_path: Option<PathBuf>,
    ) -> anyhow::Result<RemoveResult> {
        let current_path = current_path.map_or_else(|| self.current_worktree().root(), Ok)?;
        let worktrees = self.list_worktrees()?;
        // Primary worktree path: prefer default branch's worktree, fall back to first
        // worktree, then repo base for bare repos with no worktrees.
        let primary_path = self.home_path()?;

        // Phase 1: Resolve target to branch name and worktree disposition.
        // BranchOnly variants don't early-return — they go through shared validation below.
        enum Resolved {
            Worktree {
                path: PathBuf,
                branch: Option<String>,
                is_current: bool,
            },
            BranchOnly {
                branch: String,
                pruned: bool,
            },
        }

        let resolved = match target {
            RemoveTarget::Branch(branch) => {
                match worktrees
                    .iter()
                    .find(|wt| wt.branch.as_deref() == Some(branch))
                {
                    Some(wt) => {
                        if !wt.path.exists() {
                            // Directory missing - prune and continue
                            self.prune_worktrees()?;
                            Resolved::BranchOnly {
                                branch: branch.to_string(),
                                pruned: true,
                            }
                        } else if wt.locked.is_some() {
                            return Err(GitError::WorktreeLocked {
                                branch: branch.into(),
                                path: wt.path.clone(),
                                reason: wt.locked.clone(),
                            }
                            .into());
                        } else {
                            let is_current = current_path == wt.path;
                            Resolved::Worktree {
                                path: wt.path.clone(),
                                branch: Some(branch.to_string()),
                                is_current,
                            }
                        }
                    }
                    None => {
                        // No worktree found - check if the branch exists locally
                        let branch_handle = self.branch(branch);
                        if !branch_handle.exists_locally()? {
                            let remotes = branch_handle.remotes()?;
                            if !remotes.is_empty() {
                                return Err(GitError::RemoteOnlyBranch {
                                    branch: branch.into(),
                                    remote: remotes[0].clone(),
                                }
                                .into());
                            }
                            return Err(GitError::BranchNotFound {
                                branch: branch.into(),
                                show_create_hint: false,
                                last_fetch_ago: None,
                            }
                            .into());
                        }
                        Resolved::BranchOnly {
                            branch: branch.to_string(),
                            pruned: false,
                        }
                    }
                }
            }
            RemoveTarget::Current | RemoveTarget::Path(_) => {
                let lookup_path = match target {
                    RemoveTarget::Path(p) => p,
                    _ => current_path.as_path(),
                };
                let wt = worktrees
                    .iter()
                    .find(|wt| wt.path == lookup_path)
                    .ok_or_else(|| {
                        anyhow::anyhow!("Worktree not found at {}", lookup_path.display())
                    })?;
                if wt.locked.is_some() {
                    let name = wt
                        .branch
                        .clone()
                        .unwrap_or_else(|| wt.dir_name().to_string());
                    return Err(GitError::WorktreeLocked {
                        branch: name,
                        path: wt.path.clone(),
                        reason: wt.locked.clone(),
                    }
                    .into());
                }
                let is_current = wt.path == current_path;
                Resolved::Worktree {
                    path: wt.path.clone(),
                    branch: wt.branch.clone(),
                    is_current,
                }
            }
        };

        // Phase 2: Main-worktree guard (before default-branch check, since
        // -D can't override the main worktree restriction).
        if let Resolved::Worktree { ref path, .. } = resolved
            && !self.worktree_at(path).is_linked()?
        {
            return Err(GitError::CannotRemoveMainWorktree.into());
        }

        // Phase 3: Branch-level validation (applies to ALL paths).
        let branch_name = match &resolved {
            Resolved::Worktree { branch, .. } => branch.as_deref(),
            Resolved::BranchOnly { branch, .. } => Some(branch.as_str()),
        };
        if let Some(branch) = branch_name {
            check_not_default_branch(self, branch, &deletion_mode)?;
        }

        // Phase 4: Return BranchOnly early (after validation), or continue to
        // worktree-level checks. Compute integration here (same as Phase 5 does
        // for RemovedWorktree) so the output handler doesn't re-derive it.
        let (worktree_path, branch_name, is_current) = match resolved {
            Resolved::BranchOnly { branch, pruned } => {
                let default_branch = self.default_branch();
                let target = default_branch.as_deref().or(Some("HEAD"));
                let (integration_reason, target_branch) =
                    compute_integration_reason(self, Some(&branch), target, deletion_mode);
                return Ok(RemoveResult::BranchOnly {
                    branch_name: branch,
                    deletion_mode,
                    pruned,
                    target_branch,
                    integration_reason,
                });
            }
            Resolved::Worktree {
                path,
                branch,
                is_current,
            } => (path, branch, is_current),
        };

        // Phase 5: Remaining worktree-level validation.
        let target_wt = self.worktree_at(&worktree_path);

        if !force_worktree {
            target_wt.ensure_clean("remove worktree", branch_name.as_deref(), true)?;
        }

        // main_path: where post-remove hooks run from and background removal
        // executes. Prefer the primary worktree for stability (the removed worktree
        // is gone, and cwd may itself be a removal candidate during prune).
        // Fall back to cwd when the primary worktree IS the one being removed
        // (bare repo only — normal repos guard this in Phase 2 above).
        // changed_directory: whether the user needs to cd away from cwd.
        let changed_directory = is_current;
        let main_path = if worktree_path == primary_path {
            current_path
        } else {
            primary_path
        };

        // Resolve target branch for integration reason display
        let default_branch = self.default_branch();
        let target_branch = match (&default_branch, &branch_name) {
            (Some(db), Some(bn)) if db == bn => None,
            _ => default_branch,
        };

        // Pre-compute integration reason to avoid race conditions when removing
        // multiple worktrees in background mode. Use the effective target for
        // display (e.g., origin/main when upstream is ahead).
        let (integration_reason, effective_target) = compute_integration_reason(
            self,
            branch_name.as_deref(),
            target_branch.as_deref(),
            deletion_mode,
        );
        let target_branch = effective_target.or(target_branch);

        // Compute expected_path for path mismatch detection
        // Only set if actual path differs from expected (path mismatch)
        let expected_path = branch_name
            .as_ref()
            .and_then(|branch| path_mismatch(self, branch, &worktree_path, config));

        // Capture commit SHA before removal for post-remove hook template variables.
        // This ensures {{ commit }} references the removed worktree's state.
        let removed_commit = target_wt
            .run_command(&["rev-parse", "HEAD"])
            .ok()
            .map(|s| s.trim().to_string());

        Ok(RemoveResult::RemovedWorktree {
            main_path,
            worktree_path,
            changed_directory,
            branch_name,
            deletion_mode,
            target_branch,
            integration_reason,
            force_worktree,
            expected_path,
            removed_commit,
        })
    }

    fn prepare_target_worktree(
        &self,
        target_worktree: Option<&PathBuf>,
        target_branch: &str,
    ) -> anyhow::Result<Option<TargetWorktreeStash>> {
        let Some(wt_path) = target_worktree else {
            return Ok(None);
        };

        // Skip if target worktree directory is missing (prunable worktree)
        if !wt_path.exists() {
            return Ok(None);
        }

        let wt = self.worktree_at(wt_path);
        if !wt.is_dirty()? {
            return Ok(None);
        }

        let push_files = self.changed_files(target_branch, "HEAD")?;
        // Use -z for NUL-separated output: handles filenames with spaces and renames correctly
        // Format: "XY path\0" for normal files, "XY new_path\0old_path\0" for renames/copies
        let wt_status_output = wt.run_command(&["status", "--porcelain", "-z"])?;

        let wt_files: Vec<String> = parse_porcelain_z(&wt_status_output);

        let overlapping: Vec<String> = push_files
            .iter()
            .filter(|f| wt_files.contains(f))
            .cloned()
            .collect();

        if !overlapping.is_empty() {
            return Err(GitError::ConflictingChanges {
                target_branch: target_branch.to_string(),
                files: overlapping,
                worktree_path: wt_path.clone(),
            }
            .into());
        }

        let nanos = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_nanos();
        let stash_name = format!(
            "worktrunk autostash::{}::{}::{}",
            target_branch,
            process::id(),
            nanos
        );

        eprintln!(
            "{}",
            progress_message(cformat!(
                "Stashing changes in <bold>{}</>...",
                format_path_for_display(wt_path)
            ))
        );

        // Stash all changes including untracked files.
        // Note: git stash push returns exit code 0 whether or not anything was stashed.
        wt.run_command(&["stash", "push", "--include-untracked", "-m", &stash_name])?;

        // Verify stash was created by checking the stash list for our entry.
        let list_output = wt.run_command(&["stash", "list", "--format=%gd%x00%gs%x00"])?;
        let mut parts = list_output.split('\0');
        while let Some(id) = parts.next() {
            if id.is_empty() {
                continue;
            }
            if let Some(message) = parts.next()
                && (message == stash_name || message.ends_with(&stash_name))
            {
                return Ok(Some(TargetWorktreeStash::new(wt_path, id.to_string())));
            }
        }

        // Stash entry not found. Verify the worktree is now clean — if it's still
        // dirty, stashing may have failed silently or our lookup missed the entry.
        if wt.is_dirty()? {
            bail!(cformat!(
                "Failed to stash changes in {}; worktree still has uncommitted changes. \
                 Expected stash entry: <bold>{}</>. Check <bold>git stash list</>.",
                format_path_for_display(wt_path),
                stash_name
            ));
        }

        // Worktree is clean and no stash entry — nothing needed to be stashed
        Ok(None)
    }

    fn is_rebased_onto(&self, target: &str) -> anyhow::Result<bool> {
        // Orphan branches have no common ancestor, so they can't be "rebased onto" target
        let Some(merge_base) = self.merge_base("HEAD", target)? else {
            return Ok(false);
        };
        let target_sha = self.run_command(&["rev-parse", target])?.trim().to_string();

        if merge_base != target_sha {
            return Ok(false); // Target has advanced past merge-base
        }

        // Check for merge commits — if present, history is not linear
        let merge_commits = self
            .run_command(&["rev-list", "--merges", &format!("{}..HEAD", target)])?
            .trim()
            .to_string();

        Ok(merge_commits.is_empty())
    }
}

/// Check if the current worktree is the primary worktree (should not be removed).
///
/// Returns true for the main worktree in normal repos and the default branch
/// worktree in bare repos. Used by `wt merge` to skip removal silently, and
/// by `prepare_worktree_removal` Phase 2 (which errors instead of skipping).
pub(crate) fn is_primary_worktree(repo: &Repository) -> anyhow::Result<bool> {
    let current_root = repo.current_worktree().root()?;
    let primary = repo.primary_worktree()?;
    Ok(primary.as_deref() == Some(current_root.as_path()))
}

/// Compute integration reason and effective target for branch deletion.
///
/// Returns `(None, None)` if:
/// - `deletion_mode` is `ForceDelete` (skip integration check)
/// - `branch_name` is `None` (detached HEAD)
/// - `target_branch` is `None` (no target to check against)
///
/// When `Some`, the effective target may differ from the local default branch
/// (e.g., `origin/main` when upstream is ahead).
///
/// Note: Integration is computed even for `Keep` mode so we can inform the user
/// if the flag had an effect (branch was integrated) or not (branch was unmerged).
pub(crate) fn compute_integration_reason(
    repo: &Repository,
    branch_name: Option<&str>,
    target_branch: Option<&str>,
    deletion_mode: BranchDeletionMode,
) -> (Option<IntegrationReason>, Option<String>) {
    // Skip for force delete (we'll delete regardless of integration status)
    // But compute for keep mode so we can inform user if the flag had no effect
    if deletion_mode.is_force() {
        return (None, None);
    }
    let (branch, target) = match branch_name.zip(target_branch) {
        Some(pair) => pair,
        None => return (None, None),
    };
    // On error, return None (informational only)
    match repo.integration_reason(branch, target) {
        Ok((effective_target, reason)) => (reason, Some(effective_target)),
        Err(_) => (None, None),
    }
}

/// Reject removing the default branch unless force-delete is set.
///
/// The default branch is the integration target — checking it against itself is
/// tautological (same logic as `wt list`'s `is_main` guard in `check_integration_state`).
pub(crate) fn check_not_default_branch(
    repo: &Repository,
    branch: &str,
    deletion_mode: &BranchDeletionMode,
) -> anyhow::Result<()> {
    if !deletion_mode.is_force() && repo.default_branch().as_deref() == Some(branch) {
        return Err(GitError::CannotRemoveDefaultBranch {
            branch: branch.to_string(),
        }
        .into());
    }
    Ok(())
}

/// Warn about untracked files that will be auto-staged.
pub(crate) fn warn_about_untracked_files(status_output: &str) -> anyhow::Result<()> {
    let files = parse_untracked_files(status_output);
    if files.is_empty() {
        return Ok(());
    }

    let count = files.len();
    let path_word = if count == 1 { "path" } else { "paths" };
    eprintln!(
        "{}",
        warning_message(format!("Auto-staging {count} untracked {path_word}:"))
    );

    let joined_files = files.join("\n");
    eprintln!("{}", format_with_gutter(&joined_files, None));

    Ok(())
}

/// Stash guard that auto-restores on drop.
///
/// Created by `prepare_target_worktree()` when the target worktree has changes
/// that don't conflict with the push. Automatically restores the stash when
/// dropped, ensuring cleanup happens in both success and error paths.
#[must_use = "stash guard restores immediately if dropped; hold it until push completes"]
pub(crate) struct TargetWorktreeStash {
    /// Inner data wrapped in Option so we can take() in Drop.
    /// None means already restored (or disarmed).
    inner: Option<StashData>,
}

struct StashData {
    path: PathBuf,
    stash_ref: String,
}

impl StashData {
    /// Restore the stash, printing progress and warning on failure.
    fn restore(self) {
        eprintln!(
            "{}",
            progress_message(cformat!(
                "Restoring stashed changes in <bold>{}</>...",
                format_path_for_display(&self.path)
            ))
        );

        // Don't use --quiet so git shows conflicts if any
        let success = Repository::current()
            .ok()
            .and_then(|repo| {
                repo.worktree_at(&self.path)
                    .run_command(&["stash", "pop", &self.stash_ref])
                    .ok()
            })
            .is_some();

        if !success {
            eprintln!(
                "{}",
                warning_message(cformat!(
                    "Failed to restore stash <bold>{stash_ref}</>; run <bold>git stash pop {stash_ref}</> in <bold>{path}</>",
                    stash_ref = self.stash_ref,
                    path = format_path_for_display(&self.path),
                ))
            );
        }
    }
}

impl Drop for TargetWorktreeStash {
    fn drop(&mut self) {
        if let Some(data) = self.inner.take() {
            data.restore();
        }
    }
}

impl TargetWorktreeStash {
    pub(crate) fn new(path: &Path, stash_ref: String) -> Self {
        Self {
            inner: Some(StashData {
                path: path.to_path_buf(),
                stash_ref,
            }),
        }
    }

    /// Explicitly restore the stash now, preventing Drop from restoring again.
    ///
    /// Use this when you need the restore to happen at a specific point
    /// (e.g., before a success message). Drop handles errors/early returns.
    pub(crate) fn restore_now(&mut self) {
        if let Some(data) = self.inner.take() {
            data.restore();
        }
    }
}

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

    #[test]
    fn test_parse_porcelain_z_modified_staged() {
        // "M  file.txt\0" - staged modification
        let output = "M  file.txt\0";
        assert_eq!(parse_porcelain_z(output), vec!["file.txt"]);
    }

    #[test]
    fn test_parse_porcelain_z_modified_unstaged() {
        // " M file.txt\0" - unstaged modification (this was the bug case)
        let output = " M file.txt\0";
        assert_eq!(parse_porcelain_z(output), vec!["file.txt"]);
    }

    #[test]
    fn test_parse_porcelain_z_modified_both() {
        // "MM file.txt\0" - both staged and unstaged
        let output = "MM file.txt\0";
        assert_eq!(parse_porcelain_z(output), vec!["file.txt"]);
    }

    #[test]
    fn test_parse_porcelain_z_untracked() {
        // "?? new.txt\0" - untracked file
        let output = "?? new.txt\0";
        assert_eq!(parse_porcelain_z(output), vec!["new.txt"]);
    }

    #[test]
    fn test_parse_porcelain_z_rename() {
        // "R  new.txt\0old.txt\0" - rename includes both paths
        let output = "R  new.txt\0old.txt\0";
        let result = parse_porcelain_z(output);
        assert_eq!(result, vec!["new.txt", "old.txt"]);
    }

    #[test]
    fn test_parse_porcelain_z_copy() {
        // "C  copy.txt\0original.txt\0" - copy includes both paths
        let output = "C  copy.txt\0original.txt\0";
        let result = parse_porcelain_z(output);
        assert_eq!(result, vec!["copy.txt", "original.txt"]);
    }

    #[test]
    fn test_parse_porcelain_z_multiple_files() {
        // Multiple files with different statuses
        let output = " M file1.txt\0M  file2.txt\0?? untracked.txt\0R  new.txt\0old.txt\0";
        let result = parse_porcelain_z(output);
        assert_eq!(
            result,
            vec![
                "file1.txt",
                "file2.txt",
                "untracked.txt",
                "new.txt",
                "old.txt"
            ]
        );
    }

    #[test]
    fn test_parse_porcelain_z_filename_with_spaces() {
        // "M  file with spaces.txt\0"
        let output = "M  file with spaces.txt\0";
        assert_eq!(parse_porcelain_z(output), vec!["file with spaces.txt"]);
    }

    #[test]
    fn test_parse_porcelain_z_empty() {
        assert_eq!(parse_porcelain_z(""), Vec::<String>::new());
    }

    #[test]
    fn test_parse_porcelain_z_short_entry_skipped() {
        // Entry too short to have path (malformed, shouldn't happen in practice)
        let output = "M\0";
        assert_eq!(parse_porcelain_z(output), Vec::<String>::new());
    }

    #[test]
    fn test_parse_porcelain_z_rename_missing_old_path() {
        // Rename without old path (malformed, but should handle gracefully)
        let output = "R  new.txt\0";
        let result = parse_porcelain_z(output);
        // Should include new.txt, old path is simply not added
        assert_eq!(result, vec!["new.txt"]);
    }

    #[test]
    fn test_parse_untracked_files_single() {
        assert_eq!(parse_untracked_files("?? new.txt\0"), vec!["new.txt"]);
    }

    #[test]
    fn test_parse_untracked_files_multiple() {
        assert_eq!(
            parse_untracked_files("?? file1.txt\0?? file2.txt\0?? file3.txt\0"),
            vec!["file1.txt", "file2.txt", "file3.txt"]
        );
    }

    #[test]
    fn test_parse_untracked_files_ignores_modified() {
        // Only untracked files should be collected
        assert_eq!(
            parse_untracked_files(" M modified.txt\0?? untracked.txt\0"),
            vec!["untracked.txt"]
        );
    }

    #[test]
    fn test_parse_untracked_files_ignores_staged() {
        assert_eq!(
            parse_untracked_files("M  staged.txt\0?? untracked.txt\0"),
            vec!["untracked.txt"]
        );
    }

    #[test]
    fn test_parse_untracked_files_empty() {
        assert!(parse_untracked_files("").is_empty());
    }

    #[test]
    fn test_parse_untracked_files_skips_rename_old_path() {
        // Rename entries have old path as second NUL-separated field
        // Should only have untracked file, not the rename paths
        assert_eq!(
            parse_untracked_files("R  new.txt\0old.txt\0?? untracked.txt\0"),
            vec!["untracked.txt"]
        );
    }

    #[test]
    fn test_parse_untracked_files_with_spaces() {
        assert_eq!(
            parse_untracked_files("?? file with spaces.txt\0"),
            vec!["file with spaces.txt"]
        );
    }

    #[test]
    fn test_parse_untracked_files_no_untracked() {
        // All files are tracked (modified, staged, etc.)
        assert!(parse_untracked_files(" M file1.txt\0M  file2.txt\0").is_empty());
    }

    #[test]
    fn test_stash_guard_restore_now_clears_inner() {
        // Create a guard - note: this doesn't actually create a stash since we're not
        // in a real git repo with that stash ref. We're just testing the state machine.
        let mut guard = TargetWorktreeStash::new(std::path::Path::new("/tmp"), "stash@{0}".into());

        // Inner should be populated
        assert!(guard.inner.is_some());

        // restore_now() should clear inner (the restore itself will fail since no real repo,
        // but that's expected - we're testing the state transition)
        guard.restore_now();

        // Inner should now be None
        assert!(guard.inner.is_none());

        // Calling restore_now() again is a no-op
        guard.restore_now();
        assert!(guard.inner.is_none());
    }

    #[test]
    fn test_stash_guard_drop_clears_inner() {
        // Test that Drop also consumes the inner
        let guard = TargetWorktreeStash::new(std::path::Path::new("/tmp"), "stash@{0}".into());

        // Just drop it - the restore will fail (no real repo) but Drop shouldn't panic
        drop(guard);
        // If we get here, Drop worked without panicking
    }
}