noob-components/.trellis/scripts/multi_agent/cleanup.py

#!/usr/bin/env python3
"""
Multi-Agent Pipeline: Cleanup Worktree.

Usage:
    python3 cleanup.py <branch-name>      Remove specific worktree
    python3 cleanup.py --list             List all worktrees
    python3 cleanup.py --merged           Remove merged worktrees
    python3 cleanup.py --all              Remove all worktrees (with confirmation)

Options:
    -y, --yes                       Skip confirmation prompts
    --keep-branch                   Don't delete the git branch

This script:
1. Archives task directory to archive/{YYYY-MM}/
2. Removes agent from registry
3. Removes git worktree
4. Optionally deletes git branch
"""

from __future__ import annotations

import argparse
import json
import shutil
import subprocess
import sys
from pathlib import Path

import _bootstrap  # noqa: F401 — adds parent scripts/ dir to sys.path

from common.git import run_git
from common.log import Colors, log_info, log_success, log_warn, log_error
from common.paths import FILE_TASK_JSON, get_repo_root
from common.registry import (
    registry_get_file,
    registry_get_task_dir,
    registry_remove_by_id,
    registry_remove_by_worktree,
    registry_search_agent,
)
from common.task_utils import (
    archive_task_complete,
    is_safe_task_path,
)

# Colors, log_info, log_success, log_warn, log_error
# are now imported from common.log above.


def confirm(prompt: str, skip_confirm: bool) -> bool:
    """Ask for confirmation."""
    if skip_confirm:
        return True

    if not sys.stdin.isatty():
        log_error("Non-interactive mode detected. Use -y to skip confirmation.")
        return False

    response = input(f"{prompt} [y/N] ")
    return response.lower() in ("y", "yes")


def _warn_submodule_prs(task_dir_abs: Path) -> None:
    """Print reminders for any open submodule PRs found in task.json."""
    task_json = task_dir_abs / FILE_TASK_JSON
    if not task_json.is_file():
        return

    try:
        task_data = json.loads(task_json.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, OSError):
        return

    submodule_prs = task_data.get("submodule_prs")
    if not isinstance(submodule_prs, dict) or not submodule_prs:
        return

    for name, url in submodule_prs.items():
        log_warn(f"Submodule PR still open: {name} -> {url}")
    log_info("Remember to close/merge submodule PRs before cleanup")


# =============================================================================
# Commands
# =============================================================================


def cmd_list(repo_root: Path) -> int:
    """List worktrees."""
    print(f"{Colors.BLUE}=== Git Worktrees ==={Colors.NC}")
    print()

    subprocess.run(["git", "worktree", "list"], cwd=repo_root)
    print()

    # Show registry info
    registry_file = registry_get_file(repo_root)
    if registry_file and registry_file.is_file():
        print(f"{Colors.BLUE}=== Registered Agents ==={Colors.NC}")
        print()

        data = json.loads(registry_file.read_text(encoding="utf-8"))
        agents = data.get("agents", [])

        if agents:
            for agent in agents:
                print(
                    f"  {agent.get('id', '?')}: PID={agent.get('pid', '?')} [{agent.get('worktree_path', '?')}]"
                )
        else:
            print("  (none)")
        print()

    return 0


def archive_task(worktree_path: str, repo_root: Path) -> None:
    """Archive task directory."""
    task_dir = registry_get_task_dir(worktree_path, repo_root)

    if not task_dir or not is_safe_task_path(task_dir, repo_root):
        return

    task_dir_abs = repo_root / task_dir
    if not task_dir_abs.is_dir():
        return

    result = archive_task_complete(task_dir_abs, repo_root)
    if "archived_to" in result:
        dest = Path(result["archived_to"])
        log_success(f"Archived task: {dest.name} -> archive/{dest.parent.name}/")


def cleanup_registry_only(search: str, repo_root: Path, skip_confirm: bool) -> int:
    """Cleanup from registry only (no worktree)."""
    agent_info = registry_search_agent(search, repo_root)

    if not agent_info:
        log_error(f"No agent found in registry matching: {search}")
        return 1

    agent_id = agent_info.get("id", "?")
    task_dir = agent_info.get("task_dir", "?")

    print()
    print(f"{Colors.BLUE}=== Cleanup Agent (no worktree) ==={Colors.NC}")
    print(f"  Agent ID:  {agent_id}")
    print(f"  Task Dir:  {task_dir}")
    print()

    if not confirm("Archive task and remove from registry?", skip_confirm):
        log_info("Aborted")
        return 0

    # Check for submodule PRs and archive task directory
    if task_dir and is_safe_task_path(task_dir, repo_root):
        task_dir_abs = repo_root / task_dir
        if task_dir_abs.is_dir():
            _warn_submodule_prs(task_dir_abs)
            result = archive_task_complete(task_dir_abs, repo_root)
            if "archived_to" in result:
                dest = Path(result["archived_to"])
                log_success(
                    f"Archived task: {dest.name} -> archive/{dest.parent.name}/"
                )
    else:
        log_warn("Invalid task_dir in registry, skipping archive")

    # Remove from registry
    registry_remove_by_id(agent_id, repo_root)
    log_success(f"Removed from registry: {agent_id}")

    log_success("Cleanup complete")
    return 0


def cleanup_worktree(
    branch: str, repo_root: Path, skip_confirm: bool, keep_branch: bool
) -> int:
    """Cleanup single worktree."""
    # Find worktree path for branch
    _, worktree_list, _ = run_git(
        ["worktree", "list", "--porcelain"], cwd=repo_root
    )

    worktree_path = None
    current_worktree = None

    for line in worktree_list.splitlines():
        if line.startswith("worktree "):
            current_worktree = line[9:]  # Remove "worktree " prefix
        elif line.startswith("branch refs/heads/"):
            current_branch = line[18:]  # Remove "branch refs/heads/" prefix
            if current_branch == branch:
                worktree_path = current_worktree
                break

    if not worktree_path:
        # No worktree found, try to cleanup from registry only
        log_warn(f"No worktree found for: {branch}")
        log_info("Trying to cleanup from registry...")
        return cleanup_registry_only(branch, repo_root, skip_confirm)

    print()
    print(f"{Colors.BLUE}=== Cleanup Worktree ==={Colors.NC}")
    print(f"  Branch:   {branch}")
    print(f"  Worktree: {worktree_path}")
    print()

    if not confirm("Remove this worktree?", skip_confirm):
        log_info("Aborted")
        return 0

    # 1. Archive task (and check for submodule PRs)
    task_dir = registry_get_task_dir(worktree_path, repo_root)
    if task_dir and is_safe_task_path(task_dir, repo_root):
        task_dir_abs_for_warn = repo_root / task_dir
        if task_dir_abs_for_warn.is_dir():
            _warn_submodule_prs(task_dir_abs_for_warn)
    archive_task(worktree_path, repo_root)

    # 2. Remove from registry
    registry_remove_by_worktree(worktree_path, repo_root)
    log_info("Removed from registry")

    # 3. Remove worktree
    log_info("Removing worktree...")
    ret, _, _ = run_git(
        ["worktree", "remove", worktree_path, "--force"], cwd=repo_root
    )
    if ret != 0:
        # Try removing directory manually
        try:
            shutil.rmtree(worktree_path)
        except Exception as e:
            log_error(f"Failed to remove worktree: {e}")

    log_success("Worktree removed")

    # 4. Delete branch (optional)
    if not keep_branch:
        log_info("Deleting branch...")
        ret, _, _ = run_git(["branch", "-D", branch], cwd=repo_root)
        if ret != 0:
            log_warn("Could not delete branch (may be checked out elsewhere)")

    log_success(f"Cleanup complete for: {branch}")
    return 0


def cmd_merged(repo_root: Path, skip_confirm: bool, keep_branch: bool) -> int:
    """Cleanup merged worktrees."""
    # Get main branch
    _, head_out, _ = run_git(
        ["symbolic-ref", "refs/remotes/origin/HEAD"], cwd=repo_root
    )
    main_branch = head_out.strip().replace("refs/remotes/origin/", "") or "main"

    print(f"{Colors.BLUE}=== Finding Merged Worktrees ==={Colors.NC}")
    print()

    # Get merged branches
    _, merged_out, _ = run_git(
        ["branch", "--merged", main_branch], cwd=repo_root
    )
    merged_branches = []
    for line in merged_out.splitlines():
        branch = line.strip().lstrip("* ")
        if branch and branch != main_branch:
            merged_branches.append(branch)

    if not merged_branches:
        log_info("No merged branches found")
        return 0

    # Get worktree list
    _, worktree_list, _ = run_git(["worktree", "list"], cwd=repo_root)

    worktree_branches = []
    for branch in merged_branches:
        if f"[{branch}]" in worktree_list:
            worktree_branches.append(branch)
            print(f"  - {branch}")

    if not worktree_branches:
        log_info("No merged worktrees found")
        return 0

    print()
    if not confirm("Remove these merged worktrees?", skip_confirm):
        log_info("Aborted")
        return 0

    for branch in worktree_branches:
        cleanup_worktree(branch, repo_root, True, keep_branch)

    return 0


def cmd_all(repo_root: Path, skip_confirm: bool, keep_branch: bool) -> int:
    """Cleanup all worktrees."""
    print(f"{Colors.BLUE}=== All Worktrees ==={Colors.NC}")
    print()

    # Get worktree list
    _, worktree_list, _ = run_git(
        ["worktree", "list", "--porcelain"], cwd=repo_root
    )

    worktrees = []
    main_worktree = str(repo_root.resolve())

    for line in worktree_list.splitlines():
        if line.startswith("worktree "):
            wt = line[9:]
            if wt != main_worktree:
                worktrees.append(wt)

    if not worktrees:
        log_info("No worktrees to remove")
        return 0

    for wt in worktrees:
        print(f"  - {wt}")

    print()
    print(f"{Colors.RED}WARNING: This will remove ALL worktrees!{Colors.NC}")

    if not confirm("Are you sure?", skip_confirm):
        log_info("Aborted")
        return 0

    # Get branch for each worktree
    for wt in worktrees:
        # Find branch name from worktree list
        _, wt_list, _ = run_git(["worktree", "list"], cwd=repo_root)
        for line in wt_list.splitlines():
            if wt in line:
                # Extract branch from [branch] format
                import re

                match = re.search(r"\[([^\]]+)\]", line)
                if match:
                    branch = match.group(1)
                    cleanup_worktree(branch, repo_root, True, keep_branch)
                break

    return 0


# =============================================================================
# Main
# =============================================================================


def main() -> int:
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Multi-Agent Pipeline: Cleanup Worktree"
    )
    parser.add_argument("branch", nargs="?", help="Branch name to cleanup")
    parser.add_argument("-y", "--yes", action="store_true", help="Skip confirmation")
    parser.add_argument(
        "--keep-branch", action="store_true", help="Don't delete git branch"
    )
    parser.add_argument("--list", action="store_true", help="List all worktrees")
    parser.add_argument("--merged", action="store_true", help="Remove merged worktrees")
    parser.add_argument("--all", action="store_true", help="Remove all worktrees")

    args = parser.parse_args()
    repo_root = get_repo_root()

    if args.list:
        return cmd_list(repo_root)
    elif args.merged:
        return cmd_merged(repo_root, args.yes, args.keep_branch)
    elif args.all:
        return cmd_all(repo_root, args.yes, args.keep_branch)
    elif args.branch:
        return cleanup_worktree(args.branch, repo_root, args.yes, args.keep_branch)
    else:
        print("""Usage:
  python3 cleanup.py <branch-name>      Remove specific worktree
  python3 cleanup.py --list             List all worktrees
  python3 cleanup.py --merged           Remove merged worktrees
  python3 cleanup.py --all              Remove all worktrees

Options:
  -y, --yes             Skip confirmation
  --keep-branch         Don't delete git branch
""")
        return 1


if __name__ == "__main__":
    sys.exit(main())