gbranch

#!/usr/bin/env -S uv run --quiet --script
# /// script
# requires-python = ">=3.11"
# dependencies = []
# ///

"""
gbranch - Manage branches with naming convention {USER}/{datestamp}-{name}

Usage:
    gbranch                    # List active branches
    gbranch <name>             # Create or checkout branch matching name
    gbranch --cleanup          # Delete merged branches (dry-run by default)
    gbranch --cleanup --no-dry-run  # Actually delete merged branches
    gbranch --completions      # Generate shell completions
    gbranch --complete=zsh     # Install zsh completions to ~/.zshrc
    gbranch --complete=bash    # Install bash completions to ~/.bashrc
"""

USER="rjpower"

import argparse
import re
import subprocess
import sys
from datetime import datetime
from pathlib import Path


def run_git(*args: str) -> str:
    """Run git command and return output."""
    result = subprocess.run(
        ["git", *args],
        capture_output=True,
        text=True,
        check=True,
    )
    return result.stdout.strip()


def get_all_branches() -> list[str]:
    """Get all local and remote branches matching {USER}/ pattern."""
    try:
        output = run_git("branch", "-a")
    except subprocess.CalledProcessError:
        return []

    branches = []
    for line in output.splitlines():
        line = line.strip()
        # Remove markers and remotes/ prefix
        line = re.sub(r"^\*?\s+", "", line)
        line = re.sub(r"^remotes/origin/", "", line)

        if line.startswith(f"{USER}/") and not line.endswith("/HEAD"):
            branches.append(line)

    # Deduplicate while preserving order
    seen = set()
    unique_branches = []
    for branch in branches:
        if branch not in seen:
            seen.add(branch)
            unique_branches.append(branch)

    return unique_branches


def get_current_branch() -> str | None:
    """Get the current branch name."""
    try:
        return run_git("branch", "--show-current")
    except subprocess.CalledProcessError:
        return None


def is_branch_merged(branch: str, base: str = "main") -> bool:
    """Check if branch is fully merged into base (handles regular, squash, and rebase merges).

    Uses two strategies:
    1. git cherry for fast detection of regular merges
    2. git patch-id for detecting squash merges by comparing content changes
    """
    try:
        # Strategy 1: Check git cherry (fast, handles non-squash merges)
        cherry_output = run_git("cherry", base, branch)
        if not cherry_output:
            return True
        lines = cherry_output.splitlines()
        if all(line.startswith("-") for line in lines if line.strip()):
            return True

        # Strategy 2: Use patch-id for squash detection
        merge_base = run_git("merge-base", base, branch)

        # Get patch-id for entire branch diff
        branch_diff = run_git("diff", f"{merge_base}..{branch}")
        if not branch_diff:
            return True  # Empty diff = already merged

        # Calculate patch-id for branch
        result = subprocess.run(
            ["git", "patch-id", "--stable"],
            input=branch_diff,
            capture_output=True,
            text=True,
            check=True,
        )
        branch_patch_id_line = result.stdout.strip()
        if not branch_patch_id_line:
            return False
        branch_patch_id = branch_patch_id_line.split()[0]

        # Check each commit in base for matching patch-id
        rev_list = run_git("rev-list", f"{branch}..{base}")
        for commit in rev_list.splitlines():
            commit = commit.strip()
            if not commit:
                continue

            commit_diff = run_git("diff", f"{commit}~..{commit}")
            if not commit_diff:
                continue

            result = subprocess.run(
                ["git", "patch-id", "--stable"],
                input=commit_diff,
                capture_output=True,
                text=True,
                check=True,
            )
            commit_patch_id_line = result.stdout.strip()
            if not commit_patch_id_line:
                continue
            commit_patch_id = commit_patch_id_line.split()[0]

            if branch_patch_id == commit_patch_id:
                return True

        return False
    except (subprocess.CalledProcessError, IndexError):
        return False


def list_branches() -> None:
    """List active branches with current branch highlighted and merged status."""
    branches = get_all_branches()
    current = get_current_branch()

    if not branches:
        print(f"No {USER}/ branches found")
        return

    # Sort by date (most recent first)
    pattern = re.compile(USER + r"/(\d{8})-(.+)")

    def sort_key(branch: str) -> tuple[str, str]:
        match = pattern.match(branch)
        if match:
            return (match.group(1), match.group(2))
        return ("", branch)

    branches.sort(key=sort_key, reverse=True)

    # Filter to local branches only for merged check
    local_branches = [b for b in branches if not b.startswith("remotes/")]

    for branch in branches:
        marker = "* " if branch == current else "  "
        # Show merged status for local branches only
        merged_marker = ""
        if branch in local_branches and branch != current:
            merged_marker = "[x] " if is_branch_merged(branch) else "    "
        print(f"{marker}{merged_marker}{branch}")


def find_matching_branches(name: str) -> list[str]:
    """Find branches matching the given name (partial match supported)."""
    branches = get_all_branches()
    pattern = re.compile(USER + r"/\d{8}-(.+)")

    matches = []
    for branch in branches:
        match = pattern.match(branch)
        if match:
            suffix = match.group(1)
            # Match if the suffix contains the search term or ends with it
            if name in suffix or suffix.endswith(name):
                matches.append(branch)

    return matches


def checkout_or_create(name: str) -> None:
    """Checkout existing branch or create new one."""
    matches = find_matching_branches(name)

    if len(matches) == 1:
        branch = matches[0]
        print(f"Checking out: {branch}")
        subprocess.run(["git", "checkout", branch], check=True)
    elif len(matches) > 1:
        print(f"Multiple branches match '{name}':")
        for branch in matches:
            print(f"  {branch}")
        print(f"\nPlease be more specific or use full branch name")
        sys.exit(1)
    else:
        # Create new branch
        datestamp = datetime.now().strftime("%Y%m%d")
        full_branch = f"{USER}/{datestamp}-{name}"
        print(f"Creating new branch: {full_branch}")
        subprocess.run(["git", "checkout", "-b", full_branch, "main"], check=True)


def cleanup_merged_branches(dry_run: bool = False, force: bool = False) -> None:
    """Delete local branches that have been fully merged into main."""
    subprocess.run(["git", "remote", "prune", "origin"])
    all_branches = get_all_branches()
    current = get_current_branch()

    # Only consider local branches (not remotes/)
    local_branches = [b for b in all_branches if not b.startswith("remotes/")]

    # Find merged branches (excluding current)
    merged_branches = []
    for branch in local_branches:
        if branch == current:
            continue
        if is_branch_merged(branch):
            merged_branches.append(branch)

    if not merged_branches:
        print("No merged branches to clean up")
        return

    mode = "[DRY RUN] Would delete" if dry_run else "Deleting"
    print(f"{mode} {len(merged_branches)} merged branch(es):")
    for branch in merged_branches:
        print(f"  {branch}")

    if not dry_run:
        print()
        for branch in merged_branches:
            try:
                if not force:
                    run_git("branch", "-d", branch)
                else:
                    run_git("branch", "-D", branch)
                print(f"✓ Deleted: {branch}")
            except subprocess.CalledProcessError as e:
                print(f"✗ Failed to delete {branch}: {e}")


def generate_completions() -> None:
    """Output branch suffixes for shell completion."""
    branches = get_all_branches()
    pattern = re.compile(USER + r"/\d{8}-(.+)")

    suffixes = []
    for branch in branches:
        match = pattern.match(branch)
        if match:
            suffixes.append(match.group(1))

    for suffix in sorted(set(suffixes)):
        print(suffix)


def install_completion(shell: str) -> None:
    """Install shell completion by adding to rc file."""
    home = Path.home()

    if shell == "zsh":
        rc_file = home / ".zshrc"
        completion_code = """
# gbranch completion
_gbranch() {
    local -a branches
    branches=(${(f)"$(gbranch --completions 2>/dev/null)"})
    _describe 'branches' branches
}
compdef _gbranch gbranch
"""
    elif shell == "bash":
        rc_file = home / ".bashrc"
        completion_code = """
# gbranch completion
_gbranch_completions() {
    local cur="${COMP_WORDS[COMP_CWORD]}"
    if [ $COMP_CWORD -eq 1 ]; then
        COMPREPLY=($(compgen -W "$(gbranch --completions 2>/dev/null)" -- "$cur"))
    fi
}
complete -F _gbranch_completions gbranch
"""
    else:
        print(f"Unknown shell: {shell}")
        sys.exit(1)

    if not rc_file.exists():
        print(f"Error: {rc_file} not found")
        sys.exit(1)

    content = rc_file.read_text()
    marker = "# gbranch completion"

    if marker in content:
        print(f"Completion already installed in {rc_file}")
        return

    with rc_file.open("a") as f:
        f.write(completion_code)

    print(f"Completion installed to {rc_file}")
    print(f"Run: source {rc_file}")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Manage branches with {USER}/{datestamp}-{name} convention",
        add_help=False,
    )
    parser.add_argument("name", nargs="?", help="Branch name or partial match")
    parser.add_argument("--cleanup", action="store_true", help="Delete merged branches")
    parser.add_argument("--dry-run", default=False, action="store_true", help="Actually delete.")
    parser.add_argument("--force", default=False, action="store_true", help="-D")
    parser.add_argument("--completions", action="store_true", help="Generate completions")
    parser.add_argument("--complete", choices=["bash", "zsh"], help="Install completions for shell")
    parser.add_argument("-h", "--help", action="store_true", help="Show help")

    args = parser.parse_args()

    if args.help:
        print(__doc__)
        sys.exit(0)

    if args.completions:
        generate_completions()
        return

    if args.complete:
        install_completion(args.complete)
        return

    if args.cleanup:
        cleanup_merged_branches(dry_run=args.dry_run, force=args.force)
        return

    if args.name:
        checkout_or_create(args.name)
    else:
        list_branches()


if __name__ == "__main__":
    main()